Simplify data representation

web/src/types/chat.ts

@@ -1,17 +1,30 @@
+export type ToolCallFunction = {
+  name: string;
+  arguments: string;
+};
+
+export type WireToolCall = {
+  id: string;
+  type?: string;
+  function: ToolCallFunction;
+};
+
+export type ChatMessage = {
+  role: "user" | "assistant" | "tool";
+  content: unknown;
+  tool_call_id?: string;
+  name?: string;
+  tool_calls?: WireToolCall[];
+  reasoning?: string;
+  stats?: ChatStats;
+};
+
 export type ToolCall = {
   name: string;
   arguments?: Record<string, unknown>;
   response?: string;
 };
 
-export type ChatMessage = {
-  role: "user" | "assistant";
-  content: string;
-  reasoning?: string;
-  toolCalls?: ToolCall[];
-  stats?: ChatStats;
-};
-
 export type StartingRequest = {
   label: string;
   prompt: string;

web/src/utils/chatUtil.ts

@@ -1,16 +1,19 @@
 import type { ChatMessage, ChatStats, ToolCall } from "@/types/chat";
 
 export type StreamChatCallbacks = {
-  /** Update the messages array (e.g. pass to setState). */
+  /** Streamed delta of the assistant's final answer text. */
-  updateMessages: (updater: (prev: ChatMessage[]) => ChatMessage[]) => void;
+  onContentDelta: (delta: string) => void;
+  /** Streamed delta of the assistant's reasoning trace. */
+  onReasoningDelta: (delta: string) => void;
+  /** The exact wire messages appended for this turn so far (tool-call turns,
+   * tool results, and — on the final emission — the final assistant message). */
+  onTurnMessages: (messages: ChatMessage[]) => void;
+  /** Token/timing stats for the turn. */
+  onStats: (stats: ChatStats) => void;
   /** Called when the stream sends an error or fetch fails. */
   onError: (message: string) => void;
   /** Called when the stream finishes (success or error). */
   onDone: () => void;
-  /** Called when the stream emits token/timing stats. The stats are also
-   * attached to the last assistant message in updateMessages, so consumers
-   * can usually rely on the message itself rather than wiring this up. */
-  onStats?: (stats: ChatStats) => void;
   /** Message used when fetch throws and no server error is available. */
   defaultErrorMessage?: string;
 };
@@ -25,7 +28,7 @@ type StatsChunk = {
 
 type StreamChunk =
   | { type: "error"; error: string }
-  | { type: "tool_calls"; tool_calls: ToolCall[] }
+  | { type: "messages"; messages: ChatMessage[] }
   | { type: "content"; delta: string }
   | { type: "reasoning"; delta: string }
   | StatsChunk;
@@ -41,16 +44,18 @@ export type StreamChatOptions = {
 export async function streamChatCompletion(
   url: string,
   headers: Record<string, string>,
-  apiMessages: { role: string; content: string }[],
+  apiMessages: ChatMessage[],
   callbacks: StreamChatCallbacks,
   signal?: AbortSignal,
   options: StreamChatOptions = {},
 ): Promise<void> {
   const {
-    updateMessages,
+    onContentDelta,
+    onReasoningDelta,
+    onTurnMessages,
+    onStats,
     onError,
     onDone,
-    onStats,
     defaultErrorMessage = "Something went wrong. Please try again.",
   } = callbacks;
 
@@ -91,65 +96,27 @@ export async function streamChatCompletion(
     const applyChunk = (data: StreamChunk) => {
       if (data.type === "error") {
         onError(data.error);
-        updateMessages((prev) =>
-          prev.filter((m) => !(m.role === "assistant" && m.content === "")),
-        );
         return "break";
       }
-      if (data.type === "tool_calls" && data.tool_calls?.length) {
+      if (data.type === "messages") {
-        updateMessages((prev) => {
+        onTurnMessages(data.messages ?? []);
-          const next = [...prev];
-          const lastMsg = next[next.length - 1];
-          if (lastMsg?.role === "assistant")
-            next[next.length - 1] = {
-              ...lastMsg,
-              toolCalls: data.tool_calls,
-            };
-          return next;
-        });
         return "continue";
       }
       if (data.type === "content" && data.delta !== undefined) {
-        updateMessages((prev) => {
+        onContentDelta(data.delta);
-          const next = [...prev];
-          const lastMsg = next[next.length - 1];
-          if (lastMsg?.role === "assistant")
-            next[next.length - 1] = {
-              ...lastMsg,
-              content: lastMsg.content + data.delta,
-            };
-          return next;
-        });
         return "continue";
       }
       if (data.type === "reasoning" && data.delta !== undefined) {
-        updateMessages((prev) => {
+        onReasoningDelta(data.delta);
-          const next = [...prev];
-          const lastMsg = next[next.length - 1];
-          if (lastMsg?.role === "assistant")
-            next[next.length - 1] = {
-              ...lastMsg,
-              reasoning: (lastMsg.reasoning ?? "") + data.delta,
-            };
-          return next;
-        });
         return "continue";
       }
       if (data.type === "stats") {
-        const stats: ChatStats = {
+        onStats({
           promptTokens: data.prompt_tokens,
           completionTokens: data.completion_tokens,
           completionDurationMs: data.completion_duration_ms,
           tokensPerSecond: data.tokens_per_second,
-        };
-        updateMessages((prev) => {
-          const next = [...prev];
-          const lastMsg = next[next.length - 1];
-          if (lastMsg?.role === "assistant")
-            next[next.length - 1] = { ...lastMsg, stats };
-          return next;
         });
-        onStats?.(stats);
         return "continue";
       }
       return "continue";
@@ -165,9 +132,8 @@ export async function streamChatCompletion(
         const trimmed = line.trim();
         if (!trimmed) continue;
         try {
-          const data = JSON.parse(trimmed) as StreamChunk & { type: string };
+          const data = JSON.parse(trimmed) as StreamChunk;
-          const result = applyChunk(data as StreamChunk);
+          if (applyChunk(data) === "break") {
-          if (result === "break") {
             hadStreamError = true;
             break;
           }
@@ -181,50 +147,63 @@ export async function streamChatCompletion(
     // Flush remaining buffer
     if (!hadStreamError && buffer.trim()) {
       try {
-        const data = JSON.parse(buffer.trim()) as StreamChunk & {
+        const data = JSON.parse(buffer.trim()) as StreamChunk;
-          type: string;
+        applyChunk(data);
-          delta?: string;
-        };
-        if (data.type === "content" && data.delta !== undefined) {
-          updateMessages((prev) => {
-            const next = [...prev];
-            const lastMsg = next[next.length - 1];
-            if (lastMsg?.role === "assistant")
-              next[next.length - 1] = {
-                ...lastMsg,
-                content: lastMsg.content + data.delta!,
-              };
-            return next;
-          });
-        }
       } catch {
         // ignore final malformed chunk
       }
     }
-
-    if (!hadStreamError) {
-      updateMessages((prev) => {
-        const next = [...prev];
-        const lastMsg = next[next.length - 1];
-        if (lastMsg?.role === "assistant" && lastMsg.content === "")
-          next[next.length - 1] = { ...lastMsg, content: " " };
-        return next;
-      });
-    }
   } catch (err) {
     if (err instanceof DOMException && err.name === "AbortError") {
       // User stopped generation — not an error
     } else {
       onError(defaultErrorMessage);
-      updateMessages((prev) =>
-        prev.filter((m) => !(m.role === "assistant" && m.content === "")),
-      );
     }
   } finally {
     onDone();
   }
 }
 
+/** Map each tool result message to its tool_call_id for response lookup. */
+export function toolResponsesById(
+  messages: ChatMessage[],
+): Map<string, string> {
+  const map = new Map<string, string>();
+  for (const m of messages) {
+    if (m.role === "tool" && typeof m.tool_call_id === "string") {
+      map.set(
+        m.tool_call_id,
+        typeof m.content === "string" ? m.content : JSON.stringify(m.content),
+      );
+    }
+  }
+  return map;
+}
+
+/** Derive the display tool calls for one assistant message. */
+export function toolCallsForMessage(
+  message: ChatMessage,
+  responses: Map<string, string>,
+): ToolCall[] {
+  if (!message.tool_calls?.length) return [];
+  return message.tool_calls.map((tc) => {
+    let args: Record<string, unknown> | undefined;
+    const raw = tc.function?.arguments;
+    if (typeof raw === "string") {
+      try {
+        args = JSON.parse(raw) as Record<string, unknown>;
+      } catch {
+        args = undefined;
+      }
+    }
+    return {
+      name: tc.function?.name ?? "",
+      arguments: args,
+      response: responses.get(tc.id),
+    };
+  });
+}
+
 /**
  * Parse search_objects tool call response(s) into event ids for thumbnails.
  */