initial

apps/web/src/hooks/useSessionStream.ts

@@ -0,0 +1,139 @@
+import { useEffect, useRef, useState } from 'react';
+import type { Message, WsFrame } from '@/api/types';
+
+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: '',
+        role: frame.role,
+        content: '',
+        tool_calls: null,
+        tool_results: null,
+        status: 'streaming',
+        last_seq: 0,
+        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: '',
+        role: 'tool',
+        content: '',
+        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,
+        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 } : m
+      );
+      return { ...state, messages: next };
+    }
+    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;
+}