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

import { type MutableRefObject, useState, useEffect, useRef, useCallback } from 'react';
import { io, type Socket } from 'socket.io-client';
import type {
  ServerToClientEvents,
  ClientToServerEvents,
  MessageAckPayload,
  AgentEndPayload,
  AgentTextPayload,
  AgentThinkingPayload,
  ToolStartPayload,
  ToolEndPayload,
  SessionInfoPayload,
  ErrorPayload,
  CommandManifestPayload,
  SlashCommandResultPayload,
  SystemReloadPayload,
  RoutingDecisionInfo,
} from '@mosaic/types';
import { commandRegistry } from '../commands/index.js';

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

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

export interface TokenUsage {
  input: number;
  output: number;
  total: number;
  cacheRead: number;
  cacheWrite: number;
  cost: number;
  contextPercent: number;
  contextWindow: number;
}

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

type TypedSocket = Socket<ServerToClientEvents, ClientToServerEvents>;

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;
  providerName: string | null;
  thinkingLevel: string;
  availableThinkingLevels: string[];
  /** Last routing decision received from the gateway (M4-008) */
  routingDecision: RoutingDecisionInfo | null;
  sendMessage: (content: string) => void;
  addSystemMessage: (content: string) => void;
  setThinkingLevel: (level: string) => void;
  switchConversation: (id: string) => void;
  clearMessages: () => void;
  connectionError: string | null;
  socketRef: MutableRefObject<TypedSocket | null>;
}

const EMPTY_USAGE: TokenUsage = {
  input: 0,
  output: 0,
  total: 0,
  cacheRead: 0,
  cacheWrite: 0,
  cost: 0,
  contextPercent: 0,
  contextWindow: 0,
};

export function useSocket(opts: UseSocketOptions): UseSocketReturn {
  const {
    gatewayUrl,
    sessionCookie,
    initialConversationId,
    initialModel,
    initialProvider,
    agentId,
  } = 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[]>([]);
  const [tokenUsage, setTokenUsage] = useState<TokenUsage>(EMPTY_USAGE);
  const [modelName, setModelName] = useState<string | null>(null);
  const [providerName, setProviderName] = useState<string | null>(null);
  const [thinkingLevel, setThinkingLevelState] = useState<string>('off');
  const [availableThinkingLevels, setAvailableThinkingLevels] = useState<string[]>([]);
  const [routingDecision, setRoutingDecision] = useState<RoutingDecisionInfo | null>(null);
  const [connectionError, setConnectionError] = useState<string | null>(null);

  const socketRef = useRef<TypedSocket | null>(null);
  const conversationIdRef = useRef(conversationId);
  conversationIdRef.current = conversationId;

  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('session:info', (data: SessionInfoPayload) => {
      setProviderName(data.provider);
      setModelName(data.modelId);
      setThinkingLevelState(data.thinkingLevel);
      setAvailableThinkingLevels(data.availableThinkingLevels);
      // Update routing decision if provided (M4-008)
      if (data.routingDecision) {
        setRoutingDecision(data.routingDecision);
      }
    });

    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', (data: AgentEndPayload) => {
      setCurrentStreamText((prev) => {
        if (prev) {
          setMessages((msgs) => [
            ...msgs,
            { role: 'assistant', content: prev, timestamp: new Date() },
          ]);
        }
        return '';
      });
      setCurrentThinkingText('');
      setActiveToolCalls([]);
      setIsStreaming(false);

      // Update usage from the payload
      if (data.usage) {
        setProviderName(data.usage.provider);
        setModelName(data.usage.modelId);
        setThinkingLevelState(data.usage.thinkingLevel);
        setTokenUsage({
          input: data.usage.tokens.input,
          output: data.usage.tokens.output,
          total: data.usage.tokens.total,
          cacheRead: data.usage.tokens.cacheRead,
          cacheWrite: data.usage.tokens.cacheWrite,
          cost: data.usage.cost,
          contextPercent: data.usage.context.percent ?? 0,
          contextWindow: data.usage.context.window,
        });
      }
    });

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

    socket.on('commands:manifest', (data: CommandManifestPayload) => {
      commandRegistry.updateManifest(data.manifest);
    });

    socket.on('command:result', (data: SlashCommandResultPayload) => {
      const prefix = data.success ? '' : 'Error: ';
      const text = data.message ?? (data.success ? 'Done.' : 'Command failed.');
      setMessages((msgs) => [
        ...msgs,
        { role: 'system', content: `${prefix}${text}`, timestamp: new Date() },
      ]);
    });

    socket.on('system:reload', (data: SystemReloadPayload) => {
      commandRegistry.updateManifest({
        commands: data.commands,
        skills: data.skills,
        version: Date.now(),
      });
      setMessages((msgs) => [
        ...msgs,
        { role: 'system', content: data.message, timestamp: new Date() },
      ]);
    });

    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,
        ...(initialProvider ? { provider: initialProvider } : {}),
        ...(initialModel ? { modelId: initialModel } : {}),
        ...(agentId ? { agentId } : {}),
      });
    },
    [conversationId, isStreaming],
  );

  const addSystemMessage = useCallback((content: string) => {
    setMessages((msgs) => [...msgs, { role: 'system', content, timestamp: new Date() }]);
  }, []);

  const setThinkingLevel = useCallback((level: string) => {
    const cid = conversationIdRef.current;
    if (!socketRef.current?.connected || !cid) return;
    socketRef.current.emit('set:thinking', {
      conversationId: cid,
      level,
    });
  }, []);

  const clearMessages = useCallback(() => {
    setMessages([]);
    setCurrentStreamText('');
    setCurrentThinkingText('');
    setActiveToolCalls([]);
    setIsStreaming(false);
  }, []);

  const switchConversation = useCallback(
    (id: string) => {
      clearMessages();
      setConversationId(id);
    },
    [clearMessages],
  );

  return {
    connected,
    connecting,
    messages,
    conversationId,
    isStreaming,
    currentStreamText,
    currentThinkingText,
    activeToolCalls,
    tokenUsage,
    modelName,
    providerName,
    thinkingLevel,
    availableThinkingLevels,
    routingDecision,
    sendMessage,
    addSystemMessage,
    setThinkingLevel,
    switchConversation,
    clearMessages,
    connectionError,
    socketRef,
  };
}