stack/packages/cli/src/tui/hooks/use-socket.ts

import { useState, useEffect, useRef, useCallback } from 'react';
import { io, type Socket } from 'socket.io-client';
import type {
  ServerToClientEvents,
  ClientToServerEvents,
  MessageAckPayload,
  AgentTextPayload,
  AgentThinkingPayload,
  ToolStartPayload,
  ToolEndPayload,
  ErrorPayload,
} from '@mosaic/types';

export interface ToolCall {
  toolCallId: string;
  toolName: string;
  status: 'running' | 'success' | 'error';
}

export interface Message {
  role: 'user' | 'assistant' | 'thinking' | 'tool';
  content: string;
  timestamp: Date;
  toolCalls?: ToolCall[];
}

export interface TokenUsage {
  input: number;
  output: number;
  total: number;
}

export interface UseSocketOptions {
  gatewayUrl: string;
  sessionCookie?: string;
  initialConversationId?: string;
}

export interface UseSocketReturn {
  connected: boolean;
  connecting: boolean;
  messages: Message[];
  conversationId: string | undefined;
  isStreaming: boolean;
  currentStreamText: string;
  currentThinkingText: string;
  activeToolCalls: ToolCall[];
  tokenUsage: TokenUsage;
  modelName: string | null;
  sendMessage: (content: string) => void;
  connectionError: string | null;
}

type TypedSocket = Socket<ServerToClientEvents, ClientToServerEvents>;

export function useSocket(opts: UseSocketOptions): UseSocketReturn {
  const { gatewayUrl, sessionCookie, initialConversationId } = opts;

  const [connected, setConnected] = useState(false);
  const [connecting, setConnecting] = useState(true);
  const [messages, setMessages] = useState<Message[]>([]);
  const [conversationId, setConversationId] = useState(initialConversationId);
  const [isStreaming, setIsStreaming] = useState(false);
  const [currentStreamText, setCurrentStreamText] = useState('');
  const [currentThinkingText, setCurrentThinkingText] = useState('');
  const [activeToolCalls, setActiveToolCalls] = useState<ToolCall[]>([]);
  // TODO: wire up once gateway emits token-usage and model-info events
  const tokenUsage: TokenUsage = { input: 0, output: 0, total: 0 };
  const modelName: string | null = null;
  const [connectionError, setConnectionError] = useState<string | null>(null);

  const socketRef = useRef<TypedSocket | null>(null);

  useEffect(() => {
    const socket = io(`${gatewayUrl}/chat`, {
      transports: ['websocket'],
      extraHeaders: sessionCookie ? { Cookie: sessionCookie } : undefined,
      reconnection: true,
      reconnectionDelay: 2000,
      reconnectionAttempts: Infinity,
    }) as TypedSocket;

    socketRef.current = socket;

    socket.on('connect', () => {
      setConnected(true);
      setConnecting(false);
      setConnectionError(null);
    });

    socket.on('disconnect', () => {
      setConnected(false);
      setIsStreaming(false);
      setCurrentStreamText('');
      setCurrentThinkingText('');
      setActiveToolCalls([]);
    });

    socket.io.on('error', (err: Error) => {
      setConnecting(false);
      setConnectionError(err.message);
    });

    socket.on('message:ack', (data: MessageAckPayload) => {
      setConversationId(data.conversationId);
    });

    socket.on('agent:start', () => {
      setIsStreaming(true);
      setCurrentStreamText('');
      setCurrentThinkingText('');
      setActiveToolCalls([]);
    });

    socket.on('agent:text', (data: AgentTextPayload) => {
      setCurrentStreamText((prev) => prev + data.text);
    });

    socket.on('agent:thinking', (data: AgentThinkingPayload) => {
      setCurrentThinkingText((prev) => prev + data.text);
    });

    socket.on('agent:tool:start', (data: ToolStartPayload) => {
      setActiveToolCalls((prev) => [
        ...prev,
        { toolCallId: data.toolCallId, toolName: data.toolName, status: 'running' },
      ]);
    });

    socket.on('agent:tool:end', (data: ToolEndPayload) => {
      setActiveToolCalls((prev) =>
        prev.map((tc) =>
          tc.toolCallId === data.toolCallId
            ? { ...tc, status: data.isError ? 'error' : 'success' }
            : tc,
        ),
      );
    });

    socket.on('agent:end', () => {
      setCurrentStreamText((prev) => {
        if (prev) {
          setMessages((msgs) => [
            ...msgs,
            { role: 'assistant', content: prev, timestamp: new Date() },
          ]);
        }
        return '';
      });
      setCurrentThinkingText('');
      setActiveToolCalls([]);
      setIsStreaming(false);
    });

    socket.on('error', (data: ErrorPayload) => {
      setMessages((msgs) => [
        ...msgs,
        { role: 'assistant', content: `Error: ${data.error}`, timestamp: new Date() },
      ]);
      setIsStreaming(false);
    });

    return () => {
      socket.disconnect();
    };
  }, [gatewayUrl, sessionCookie]);

  const sendMessage = useCallback(
    (content: string) => {
      if (!content.trim() || isStreaming) return;
      if (!socketRef.current?.connected) return;

      setMessages((msgs) => [...msgs, { role: 'user', content, timestamp: new Date() }]);

      socketRef.current.emit('message', {
        conversationId,
        content,
      });
    },
    [conversationId, isStreaming],
  );

  return {
    connected,
    connecting,
    messages,
    conversationId,
    isStreaming,
    currentStreamText,
    currentThinkingText,
    activeToolCalls,
    tokenUsage,
    modelName,
    sendMessage,
    connectionError,
  };
}