claude-web-ui/frontend/src/hooks/useClaudeSession.js

import { useState, useRef, useCallback, useEffect } from 'react';

const WS_URL = import.meta.env.VITE_WS_URL || 'ws://100.105.142.13:3001';
const API_URL = import.meta.env.VITE_API_URL || 'http://100.105.142.13:3001';

export function useClaudeSession() {
  const [connected, setConnected] = useState(false);
  const [sessionActive, setSessionActive] = useState(false);
  const [messages, setMessages] = useState([]);
  const [currentProject, setCurrentProject] = useState(null);
  const [isProcessing, setIsProcessing] = useState(false);
  const [error, setError] = useState(null);

  const wsRef = useRef(null);
  const currentAssistantMessage = useRef(null);

  const connect = useCallback(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN) return;

    const ws = new WebSocket(WS_URL);
    wsRef.current = ws;

    ws.onopen = () => {
      console.log('WebSocket connected');
      setConnected(true);
      setError(null);
    };

    ws.onclose = () => {
      console.log('WebSocket disconnected');
      setConnected(false);
      setSessionActive(false);
      setIsProcessing(false);
    };

    ws.onerror = (err) => {
      console.error('WebSocket error:', err);
      setError('Connection error');
    };

    ws.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data);
        handleMessage(data);
      } catch (e) {
        console.error('Failed to parse message:', e);
      }
    };
  }, []);

  const handleMessage = useCallback((data) => {
    console.log('Received:', data.type, data);

    switch (data.type) {
      case 'session_started':
        setSessionActive(true);
        setCurrentProject(data.project);
        setMessages(prev => [...prev, {
          type: 'system',
          content: `Session started in ${data.project}`,
          timestamp: data.timestamp
        }]);
        break;

      case 'session_ended':
        setSessionActive(false);
        setIsProcessing(false);
        setMessages(prev => [...prev, {
          type: 'system',
          content: `Session ended (code: ${data.code})`,
          timestamp: data.timestamp
        }]);
        break;

      case 'claude_event':
        handleClaudeEvent(data.event);
        break;

      case 'raw_output':
        console.log('Raw:', data.content);
        break;

      case 'stderr':
        console.log('Stderr:', data.content);
        break;

      case 'error':
        setError(data.message);
        setIsProcessing(false);
        break;
    }
  }, []);

  const handleClaudeEvent = useCallback((event) => {
    // Debug: log all event types to understand the structure
    console.log('Claude Event:', event.type, event);

    // Handle different Claude event types
    if (event.type === 'assistant') {
      setIsProcessing(true);

      // Check if this is a content block
      if (event.message?.content) {
        const newMessages = [];

        for (const block of event.message.content) {
          if (block.type === 'text' && block.text) {
            // Only add text if we don't have a streaming message with similar content
            // (to avoid duplicates from stream_event + final assistant event)
            newMessages.push({
              type: 'assistant',
              content: block.text,
              timestamp: Date.now(),
              final: true  // Mark as final message
            });
          } else if (block.type === 'tool_use') {
            // Tool use is embedded in assistant message content
            newMessages.push({
              type: 'tool_use',
              tool: block.name,
              input: block.input,
              toolUseId: block.id,
              timestamp: Date.now()
            });
          }
        }

        if (newMessages.length > 0) {
          setMessages(prev => {
            // Check if last message is a streaming message - if so, replace it with final
            const last = prev[prev.length - 1];
            if (last?.type === 'assistant' && last.streaming) {
              // Replace streaming message with final content, keep other new messages (tool_use)
              const textMessages = newMessages.filter(m => m.type === 'assistant');
              const otherMessages = newMessages.filter(m => m.type !== 'assistant');

              if (textMessages.length > 0) {
                return [
                  ...prev.slice(0, -1),
                  { ...textMessages[0], streaming: false },
                  ...otherMessages
                ];
              }
            }
            return [...prev, ...newMessages];
          });
        }
      }
    } else if (event.type === 'user' && event.tool_use_result) {
      // Tool results come as 'user' events with tool_use_result
      const result = event.tool_use_result;
      setMessages(prev => [...prev, {
        type: 'tool_result',
        content: result.content,
        toolUseId: result.tool_use_id,
        isError: result.is_error || false,
        timestamp: Date.now()
      }]);
    } else if (event.type === 'content_block_delta') {
      // Streaming delta (direct)
      if (event.delta?.text) {
        setMessages(prev => {
          const last = prev[prev.length - 1];
          if (last?.type === 'assistant' && last.streaming) {
            return [
              ...prev.slice(0, -1),
              { ...last, content: last.content + event.delta.text }
            ];
          }
          return [...prev, {
            type: 'assistant',
            content: event.delta.text,
            streaming: true,
            timestamp: Date.now()
          }];
        });
      }
    } else if (event.type === 'stream_event' && event.event?.type === 'content_block_delta') {
      // Streaming delta (wrapped in stream_event)
      const deltaText = event.event?.delta?.text;
      if (deltaText) {
        setMessages(prev => {
          const last = prev[prev.length - 1];
          if (last?.type === 'assistant' && last.streaming) {
            return [
              ...prev.slice(0, -1),
              { ...last, content: last.content + deltaText }
            ];
          }
          return [...prev, {
            type: 'assistant',
            content: deltaText,
            streaming: true,
            timestamp: Date.now()
          }];
        });
      }
    } else if (event.type === 'result') {
      // Final result - just stop processing
      setIsProcessing(false);
    } else if (event.type === 'system' && event.subtype === 'result') {
      setIsProcessing(false);
    }
  }, []);

  const loadHistory = useCallback(async (project) => {
    try {
      const encodedProject = encodeURIComponent(project);
      const response = await fetch(`${API_URL}/api/history/${encodedProject}`);
      if (response.ok) {
        const data = await response.json();
        if (data.messages && data.messages.length > 0) {
          console.log(`Loaded ${data.messages.length} messages from history`);
          setMessages(data.messages);
          return data.sessionId;
        }
      }
    } catch (err) {
      console.error('Failed to load history:', err);
    }
    return null;
  }, []);

  const startSession = useCallback(async (project = '/projects', resume = true) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
      setError('Not connected');
      return;
    }

    // Load history before starting session if resuming
    if (resume) {
      await loadHistory(project);
    }

    wsRef.current.send(JSON.stringify({
      type: 'start_session',
      project,
      resume
    }));
  }, [loadHistory]);

  const sendMessage = useCallback((message) => {
    if (!wsRef.current || wsRef.current.readyState !== WebSocket.OPEN) {
      setError('Not connected');
      return;
    }

    if (!sessionActive) {
      setError('No active session');
      return;
    }

    // Add user message to display
    setMessages(prev => [...prev, {
      type: 'user',
      content: message,
      timestamp: Date.now()
    }]);

    setIsProcessing(true);

    wsRef.current.send(JSON.stringify({
      type: 'user_message',
      message
    }));
  }, [sessionActive]);

  const stopSession = useCallback(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify({ type: 'stop_session' }));
    }
  }, []);

  const clearMessages = useCallback(() => {
    setMessages([]);
  }, []);

  // Auto-connect on mount
  useEffect(() => {
    connect();
    return () => {
      wsRef.current?.close();
    };
  }, [connect]);

  return {
    connected,
    sessionActive,
    messages,
    currentProject,
    isProcessing,
    error,
    connect,
    startSession,
    sendMessage,
    stopSession,
    clearMessages,
    setError,
    setMessages
  };
}