stack-mos-dt-0/apps/web/src/spa/chat/use-chat-connection.ts

import { useEffect, useReducer, useRef } from 'react';
import { destroySocket, getSocket } from '@/lib/socket';
import type {
  AgentEndPayload,
  AgentStartPayload,
  AgentTextPayload,
  AgentThinkingPayload,
  CommandManifest,
  CommandManifestPayload,
  ErrorPayload,
  MessageAckPayload,
  SessionInfoPayload,
  SessionUsagePayload,
  SlashCommandApprovalResultPayload,
  SlashCommandResultPayload,
  SystemReloadPayload,
  ToolEndPayload,
  ToolStartPayload,
} from '@/lib/chat-contract';

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

export interface ChatTranscriptMessage {
  id: string;
  role: 'user' | 'assistant';
  text: string;
  thinking?: string;
}

export interface PendingApproval {
  command: string;
  args?: string;
}

export interface ChatConnectionState {
  conversationId: string | null;
  /** True once a message has been sent while no conversation is active yet, so the
   * first scoped server event naming a conversation may establish it (the gateway
   * does not guarantee message:ack is the first event for a new conversation). */
  pendingSend: boolean;
  ack: MessageAckPayload | null;
  streaming: boolean;
  text: string;
  thinking: string;
  tools: ToolCallState[];
  usage: SessionUsagePayload | null;
  sessionInfo: SessionInfoPayload | null;
  manifest: CommandManifest | null;
  commandResults: SlashCommandResultPayload[];
  approval: SlashCommandApprovalResultPayload | null;
  pendingApproval: PendingApproval | null;
  /** True while an approval request has been sent and no response has arrived yet. */
  approvalRequestPending: boolean;
  systemReload: SystemReloadPayload | null;
  error: string | null;
  messages: ChatTranscriptMessage[];
}

export interface ChatConnectionActions {
  sendMessage: (input: { content: string; provider?: string; modelId?: string }) => void;
  abort: () => void;
  setThinking: (level: string) => void;
  executeCommand: (input: { command: string; args?: string }) => void;
  approveCommand: (input: { command: string; args?: string }) => void;
  runApprovedCommand: () => void;
}

export interface ChatConnectionValue {
  state: ChatConnectionState;
  actions: ChatConnectionActions;
}

const initialState: ChatConnectionState = {
  conversationId: null,
  pendingSend: false,
  ack: null,
  streaming: false,
  text: '',
  thinking: '',
  tools: [],
  usage: null,
  sessionInfo: null,
  manifest: null,
  commandResults: [],
  approval: null,
  pendingApproval: null,
  approvalRequestPending: false,
  systemReload: null,
  error: null,
  messages: [],
};

type Action =
  | { type: 'server/message:ack'; payload: MessageAckPayload }
  | { type: 'server/agent:start'; payload: AgentStartPayload }
  | { type: 'server/agent:text'; payload: AgentTextPayload }
  | { type: 'server/agent:thinking'; payload: AgentThinkingPayload }
  | { type: 'server/agent:tool:start'; payload: ToolStartPayload }
  | { type: 'server/agent:tool:end'; payload: ToolEndPayload }
  | { type: 'server/agent:end'; payload: AgentEndPayload }
  | { type: 'server/session:info'; payload: SessionInfoPayload }
  | { type: 'server/commands:manifest'; payload: CommandManifestPayload }
  | { type: 'server/command:result'; payload: SlashCommandResultPayload }
  | { type: 'server/command:approval'; payload: SlashCommandApprovalResultPayload }
  | { type: 'server/system:reload'; payload: SystemReloadPayload }
  | { type: 'server/error'; payload: ErrorPayload }
  | { type: 'local/send'; content: string }
  | { type: 'local/approve-request'; command: string; args?: string }
  | { type: 'local/consume-approval' };

/**
 * Resolves whether a scoped server event (one carrying a conversationId) belongs to
 * the active conversation. The gateway does not guarantee message:ack is the first
 * event for a new conversation (session:info, and error on auth/session-creation
 * failure, can both arrive first) — so while a message is pending and no conversation
 * is active yet, the first scoped event names the conversation instead of being
 * dropped. Once a conversation is active, only its own events pass.
 */
function resolveScopedConversation(
  state: ChatConnectionState,
  conversationId: string,
): { active: true; state: ChatConnectionState } | { active: false; state: null } {
  if (state.conversationId === conversationId) {
    return { active: true, state };
  }
  if (state.conversationId === null && state.pendingSend) {
    return { active: true, state: { ...state, conversationId, pendingSend: false } };
  }
  return { active: false, state: null };
}

function reduce(state: ChatConnectionState, action: Action): ChatConnectionState {
  switch (action.type) {
    case 'server/message:ack': {
      const { payload } = action;
      if (state.conversationId === null) {
        return {
          ...state,
          conversationId: payload.conversationId,
          pendingSend: false,
          ack: payload,
        };
      }
      if (payload.conversationId !== state.conversationId) return state;
      return { ...state, ack: payload };
    }

    case 'server/agent:start': {
      const resolved = resolveScopedConversation(state, action.payload.conversationId);
      if (!resolved.active) return state;
      return { ...resolved.state, streaming: true, text: '', thinking: '', tools: [], error: null };
    }

    case 'server/agent:text': {
      const resolved = resolveScopedConversation(state, action.payload.conversationId);
      if (!resolved.active) return state;
      return { ...resolved.state, text: resolved.state.text + action.payload.text };
    }

    case 'server/agent:thinking': {
      const resolved = resolveScopedConversation(state, action.payload.conversationId);
      if (!resolved.active) return state;
      return { ...resolved.state, thinking: resolved.state.thinking + action.payload.text };
    }

    case 'server/agent:tool:start': {
      const { payload } = action;
      const resolved = resolveScopedConversation(state, payload.conversationId);
      if (!resolved.active) return state;
      const tool: ToolCallState = {
        toolCallId: payload.toolCallId,
        toolName: payload.toolName,
        status: 'running',
      };
      return { ...resolved.state, tools: [...resolved.state.tools, tool] };
    }

    case 'server/agent:tool:end': {
      const { payload } = action;
      const resolved = resolveScopedConversation(state, payload.conversationId);
      if (!resolved.active) return state;
      return {
        ...resolved.state,
        tools: resolved.state.tools.map((tool) =>
          tool.toolCallId === payload.toolCallId
            ? { ...tool, status: payload.isError ? 'error' : 'success' }
            : tool,
        ),
      };
    }

    case 'server/agent:end': {
      const { payload } = action;
      const resolved = resolveScopedConversation(state, payload.conversationId);
      if (!resolved.active) return state;
      const next = resolved.state;
      const hasContent = next.text.length > 0 || next.thinking.length > 0;
      const messages = hasContent
        ? [
            ...next.messages,
            {
              id: `assistant-${payload.conversationId}-${next.messages.length}`,
              role: 'assistant' as const,
              text: next.text,
              thinking: next.thinking || undefined,
            },
          ]
        : next.messages;
      return {
        ...next,
        streaming: false,
        text: '',
        thinking: '',
        usage: payload.usage ?? next.usage,
        messages,
      };
    }

    case 'server/session:info': {
      const resolved = resolveScopedConversation(state, action.payload.conversationId);
      if (!resolved.active) return state;
      return { ...resolved.state, sessionInfo: action.payload };
    }

    case 'server/commands:manifest': {
      return { ...state, manifest: action.payload.manifest };
    }

    case 'server/command:result': {
      const { payload } = action;
      const resolved = resolveScopedConversation(state, payload.conversationId);
      if (!resolved.active) return state;
      return { ...resolved.state, commandResults: [...resolved.state.commandResults, payload] };
    }

    case 'server/command:approval': {
      const { payload } = action;
      const resolved = resolveScopedConversation(state, payload.conversationId);
      if (!resolved.active) return state;
      const next = resolved.state;
      if (!next.pendingApproval || next.pendingApproval.command !== payload.command) {
        // Stale or mismatched response for a request that is no longer (or never was)
        // the one outstanding approval — do not let it replace active approval state.
        return next;
      }
      return { ...next, approval: payload, approvalRequestPending: false };
    }

    case 'server/system:reload': {
      return {
        ...state,
        systemReload: action.payload,
        manifest: {
          commands: action.payload.commands,
          skills: action.payload.skills,
          version: state.manifest?.version ?? 0,
        },
      };
    }

    case 'server/error': {
      const { payload } = action;
      const resolved = resolveScopedConversation(state, payload.conversationId);
      if (!resolved.active) return state;
      return { ...resolved.state, error: payload.error, streaming: false };
    }

    case 'local/send': {
      const message: ChatTranscriptMessage = {
        id: `user-${Date.now()}-${state.messages.length}`,
        role: 'user',
        text: action.content,
      };
      return {
        ...state,
        messages: [...state.messages, message],
        error: null,
        pendingSend: state.conversationId === null ? true : state.pendingSend,
      };
    }

    case 'local/approve-request': {
      // Only one approval request may be outstanding at a time; a second request
      // before the first resolves is ignored so it cannot overwrite the original
      // command+args pair with a mismatched one.
      if (state.approvalRequestPending) return state;
      return {
        ...state,
        pendingApproval: { command: action.command, args: action.args },
        approvalRequestPending: true,
        approval: null,
      };
    }

    case 'local/consume-approval': {
      return { ...state, approval: null, pendingApproval: null };
    }

    default:
      return state;
  }
}

/**
 * Owns the `/chat` socket for the lifetime of the consuming page: subscribes
 * once, narrows the socket to the typed contract, and tears the socket down
 * on cleanup so the next visit gets a fresh authenticated connection.
 */
export function useChatConnection(): ChatConnectionValue {
  const [state, dispatch] = useReducer(reduce, initialState);
  // Reducer state updates are batched/async; a ref lets a double-click on "Run
  // approved command" be rejected synchronously, before React ever re-renders.
  const executedApprovalIds = useRef<Set<string>>(new Set());

  useEffect(() => {
    const socket = getSocket();

    const onMessageAck = (payload: MessageAckPayload): void =>
      dispatch({ type: 'server/message:ack', payload });
    const onAgentStart = (payload: AgentStartPayload): void =>
      dispatch({ type: 'server/agent:start', payload });
    const onAgentText = (payload: AgentTextPayload): void =>
      dispatch({ type: 'server/agent:text', payload });
    const onAgentThinking = (payload: AgentThinkingPayload): void =>
      dispatch({ type: 'server/agent:thinking', payload });
    const onToolStart = (payload: ToolStartPayload): void =>
      dispatch({ type: 'server/agent:tool:start', payload });
    const onToolEnd = (payload: ToolEndPayload): void =>
      dispatch({ type: 'server/agent:tool:end', payload });
    const onAgentEnd = (payload: AgentEndPayload): void =>
      dispatch({ type: 'server/agent:end', payload });
    const onSessionInfo = (payload: SessionInfoPayload): void =>
      dispatch({ type: 'server/session:info', payload });
    const onCommandsManifest = (payload: CommandManifestPayload): void =>
      dispatch({ type: 'server/commands:manifest', payload });
    const onCommandResult = (payload: SlashCommandResultPayload): void =>
      dispatch({ type: 'server/command:result', payload });
    const onCommandApproval = (payload: SlashCommandApprovalResultPayload): void =>
      dispatch({ type: 'server/command:approval', payload });
    const onSystemReload = (payload: SystemReloadPayload): void =>
      dispatch({ type: 'server/system:reload', payload });
    const onError = (payload: ErrorPayload): void => dispatch({ type: 'server/error', payload });

    socket.on('message:ack', onMessageAck);
    socket.on('agent:start', onAgentStart);
    socket.on('agent:text', onAgentText);
    socket.on('agent:thinking', onAgentThinking);
    socket.on('agent:tool:start', onToolStart);
    socket.on('agent:tool:end', onToolEnd);
    socket.on('agent:end', onAgentEnd);
    socket.on('session:info', onSessionInfo);
    socket.on('commands:manifest', onCommandsManifest);
    socket.on('command:result', onCommandResult);
    socket.on('command:approval', onCommandApproval);
    socket.on('system:reload', onSystemReload);
    socket.on('error', onError);

    if (!socket.connected) {
      socket.connect();
    }

    return () => {
      socket.off('message:ack', onMessageAck);
      socket.off('agent:start', onAgentStart);
      socket.off('agent:text', onAgentText);
      socket.off('agent:thinking', onAgentThinking);
      socket.off('agent:tool:start', onToolStart);
      socket.off('agent:tool:end', onToolEnd);
      socket.off('agent:end', onAgentEnd);
      socket.off('session:info', onSessionInfo);
      socket.off('commands:manifest', onCommandsManifest);
      socket.off('command:result', onCommandResult);
      socket.off('command:approval', onCommandApproval);
      socket.off('system:reload', onSystemReload);
      socket.off('error', onError);
      destroySocket();
    };
  }, []);

  const actions: ChatConnectionActions = {
    sendMessage: ({ content, provider, modelId }) => {
      if (state.streaming) return;
      const socket = getSocket();
      if (!socket.connected) socket.connect();
      dispatch({ type: 'local/send', content });
      socket.emit('message', {
        conversationId: state.conversationId ?? undefined,
        content,
        provider,
        modelId,
      });
    },

    abort: () => {
      if (state.conversationId === null) return;
      const socket = getSocket();
      socket.emit('abort', { conversationId: state.conversationId });
    },

    setThinking: (level) => {
      if (state.conversationId === null) return;
      const socket = getSocket();
      socket.emit('set:thinking', { conversationId: state.conversationId, level });
    },

    executeCommand: ({ command, args }) => {
      if (state.conversationId === null) return;
      const socket = getSocket();
      socket.emit('command:execute', { conversationId: state.conversationId, command, args });
    },

    approveCommand: ({ command, args }) => {
      if (state.conversationId === null) return;
      if (state.approvalRequestPending) return;
      dispatch({ type: 'local/approve-request', command, args });
      const socket = getSocket();
      socket.emit('command:approve', { conversationId: state.conversationId, command, args });
    },

    runApprovedCommand: () => {
      const { conversationId, approval, pendingApproval } = state;
      if (conversationId === null || !approval?.success || !approval.approvalId) return;
      if (!pendingApproval || pendingApproval.command !== approval.command) return;
      if (executedApprovalIds.current.has(approval.approvalId)) return;
      executedApprovalIds.current.add(approval.approvalId);
      const socket = getSocket();
      socket.emit('command:execute', {
        conversationId,
        command: approval.command,
        args: pendingApproval.args,
        approvalId: approval.approvalId,
      });
      dispatch({ type: 'local/consume-approval' });
    },
  };

  return { state, actions };
}