boocode/apps/web/src/hooks/useSessionStream.ts

import { useEffect, useRef, useState } from 'react';
import type { Message, WsFrame } from '@/api/types';
import { sessionEvents } from './sessionEvents';

// session_renamed frame removed from WsFrame — it was declared but never
// published on the per-session WS channel (server publishes via broker.publishUser
// since v1.4). chat_renamed remains; auto_name.ts publishes it on session WS.

interface State {
  messages: Message[];
  connected: boolean;
  error: string | null;
}

function applyFrame(state: State, frame: WsFrame): State {
  switch (frame.type) {
    case 'snapshot': {
      return { ...state, messages: frame.messages };
    }
    case 'message_started': {
      const exists = state.messages.some((m) => m.id === frame.message_id);
      if (exists) return state;
      const newMsg: Message = {
        id: frame.message_id,
        session_id: '',
        chat_id: frame.chat_id ?? '',
        role: frame.role,
        content: '',
        kind: 'message',
        tool_calls: null,
        tool_results: null,
        status: 'streaming',
        last_seq: 0,
        tokens_used: null,
        ctx_used: null,
        ctx_max: null,
        started_at: null,
        finished_at: null,
        created_at: new Date().toISOString(),
      };
      return { ...state, messages: [...state.messages, newMsg] };
    }
    case 'delta': {
      const next = state.messages.map((m) =>
        m.id === frame.message_id ? { ...m, content: m.content + frame.content } : m
      );
      return { ...state, messages: next };
    }
    case 'tool_call': {
      const next = state.messages.map((m) =>
        m.id === frame.message_id
          ? { ...m, tool_calls: [...(m.tool_calls ?? []), frame.tool_call] }
          : m
      );
      return { ...state, messages: next };
    }
    case 'tool_result': {
      const exists = state.messages.some((m) => m.id === frame.tool_message_id);
      if (exists) {
        const next = state.messages.map((m) =>
          m.id === frame.tool_message_id
            ? {
                ...m,
                role: 'tool' as const,
                tool_results: {
                  tool_call_id: frame.tool_call_id,
                  output: frame.output,
                  truncated: frame.truncated,
                  ...(frame.error ? { error: frame.error } : {}),
                },
                status: 'complete' as const,
              }
            : m
        );
        return { ...state, messages: next };
      }
      const newMsg: Message = {
        id: frame.tool_message_id,
        session_id: '',
        chat_id: frame.chat_id ?? '',
        role: 'tool',
        content: '',
        kind: 'message',
        tool_calls: null,
        tool_results: {
          tool_call_id: frame.tool_call_id,
          output: frame.output,
          truncated: frame.truncated,
          ...(frame.error ? { error: frame.error } : {}),
        },
        status: 'complete',
        last_seq: 0,
        tokens_used: null,
        ctx_used: null,
        ctx_max: null,
        started_at: null,
        finished_at: null,
        created_at: new Date().toISOString(),
      };
      return { ...state, messages: [...state.messages, newMsg] };
    }
    case 'message_complete': {
      const next = state.messages.map((m) =>
        m.id === frame.message_id
          ? {
              ...m,
              status: 'complete' as const,
              ...(frame.tokens_used !== undefined ? { tokens_used: frame.tokens_used } : {}),
              ...(frame.ctx_used !== undefined ? { ctx_used: frame.ctx_used } : {}),
              ...(frame.ctx_max !== undefined ? { ctx_max: frame.ctx_max } : {}),
              ...(frame.started_at !== undefined ? { started_at: frame.started_at } : {}),
              ...(frame.finished_at !== undefined ? { finished_at: frame.finished_at } : {}),
            }
          : m
      );
      return { ...state, messages: next };
    }
    case 'messages_deleted': {
      const removeSet = new Set(frame.message_ids);
      return {
        ...state,
        messages: state.messages.filter((m) => !removeSet.has(m.id)),
      };
    }
    case 'chat_renamed': {
      sessionEvents.emit({
        type: 'chat_updated',
        chat_id: frame.chat_id,
        session_id: '',
        name: frame.name,
        updated_at: new Date().toISOString(),
      });
      return state;
    }
    case 'error': {
      const next = frame.message_id
        ? state.messages.map((m) =>
            m.id === frame.message_id ? { ...m, status: 'failed' as const } : m
          )
        : state.messages;
      return { ...state, messages: next, error: frame.error };
    }
  }
}

export function useSessionStream(sessionId: string | undefined) {
  const [state, setState] = useState<State>({ messages: [], connected: false, error: null });
  const wsRef = useRef<WebSocket | null>(null);

  useEffect(() => {
    if (!sessionId) return;

    setState({ messages: [], connected: false, error: null });

    const proto = window.location.protocol === 'https:' ? 'wss' : 'ws';
    const url = `${proto}://${window.location.host}/api/ws/sessions/${sessionId}`;
    const ws = new WebSocket(url);
    wsRef.current = ws;

    ws.onopen = () => {
      setState((s) => ({ ...s, connected: true, error: null }));
    };
    ws.onmessage = (ev) => {
      try {
        const frame = JSON.parse(typeof ev.data === 'string' ? ev.data : '') as WsFrame;
        setState((s) => applyFrame(s, frame));
      } catch (err) {
        console.warn('bad ws frame', err);
      }
    };
    ws.onerror = () => {
      setState((s) => ({ ...s, error: 'websocket error' }));
    };
    ws.onclose = () => {
      setState((s) => ({ ...s, connected: false }));
    };

    return () => {
      wsRef.current = null;
      ws.close();
    };
  }, [sessionId]);

  return state;
}