Fix chat tool calling and prompt breaking (#23457)

* Implement tool call history keeping

* Refactor to match single message implementation

* Simplify data representation

* Cleanup chat page rendering

* Include system message to not break cache

* Formatting

* Update tests and update .gitignore

web/src/pages/Chat.tsx

@@ -13,6 +13,7 @@ import { ChatComposer } from "@/components/chat/ChatComposer";
 import ChatSettings from "@/components/chat/ChatSettings";
 import type {
   ChatMessage,
+  ChatStats,
   GenAIModelsResponse,
   ShowStatsMode,
 } from "@/types/chat";
@@ -22,12 +23,28 @@ import {
   getFindSimilarObjectsFromToolCalls,
   prependAttachment,
   streamChatCompletion,
+  toolCallsForMessage,
+  toolResponsesById,
 } from "@/utils/chatUtil";
 
+type StreamingTurn = {
+  content: string;
+  reasoning: string;
+  chain: ChatMessage[];
+  stats?: ChatStats;
+};
+
+const hasText = (content: unknown): content is string =>
+  typeof content === "string" && content.trim().length > 0;
+
+const toWire = (messages: ChatMessage[]): ChatMessage[] =>
+  messages.map(({ reasoning: _r, stats: _s, ...rest }) => rest);
+
 export default function ChatPage() {
   const { t } = useTranslation(["views/chat"]);
   const [input, setInput] = useState("");
   const [messages, setMessages] = useState<ChatMessage[]>([]);
+  const [streaming, setStreaming] = useState<StreamingTurn | null>(null);
   const [isLoading, setIsLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
   const [attachedEventId, setAttachedEventId] = useState<string | null>(null);
@@ -72,28 +89,19 @@ export default function ChatPage() {
     if (isNearBottom) {
       el.scrollTo({ top: el.scrollHeight, behavior: "smooth" });
     }
-  }, [messages, autoScroll]);
+  }, [messages, streaming, autoScroll]);
 
   const submitConversation = useCallback(
     async (messagesToSend: ChatMessage[]) => {
       if (isLoading) return;
       const last = messagesToSend[messagesToSend.length - 1];
-      if (!last || last.role !== "user" || !last.content.trim()) return;
+      if (!last || last.role !== "user" || !hasText(last.content)) return;
 
       setError(null);
-      const assistantPlaceholder: ChatMessage = {
-        role: "assistant",
-        content: "",
-        toolCalls: undefined,
-      };
-      setMessages([...messagesToSend, assistantPlaceholder]);
+      setMessages(messagesToSend);
+      setStreaming({ content: "", reasoning: "", chain: [] });
       setIsLoading(true);
 
-      const apiMessages = messagesToSend.map((m) => ({
-        role: m.role,
-        content: m.content,
-      }));
-
       const baseURL = axios.defaults.baseURL ?? "";
       const url = `${baseURL}chat/completion`;
       const headers: Record<string, string> = {
@@ -104,16 +112,50 @@ export default function ChatPage() {
       const controller = new AbortController();
       abortRef.current = controller;
 
+      let chain: ChatMessage[] = [];
+      let stats: ChatStats | undefined;
+      let reasoning = "";
+      let hadError = false;
+
       await streamChatCompletion(
         url,
         headers,
-        apiMessages,
+        toWire(messagesToSend),
         {
-          updateMessages: (updater) => setMessages(updater),
-          onError: (message) => setError(message),
+          onContentDelta: (delta) =>
+            setStreaming((s) => (s ? { ...s, content: s.content + delta } : s)),
+          onReasoningDelta: (delta) => {
+            reasoning += delta;
+            setStreaming((s) =>
+              s ? { ...s, reasoning: s.reasoning + delta } : s,
+            );
+          },
+          onChain: (fullChain) => {
+            chain = fullChain;
+            setStreaming((s) => (s ? { ...s, chain: fullChain } : s));
+          },
+          onStats: (s) => {
+            stats = s;
+            setStreaming((cur) => (cur ? { ...cur, stats: s } : cur));
+          },
+          onError: (message) => {
+            hadError = true;
+            setError(message);
+          },
           onDone: () => {
             abortRef.current = null;
             setIsLoading(false);
+            setStreaming(null);
+            const lastMsg = chain[chain.length - 1];
+            if (!hadError && lastMsg?.role === "assistant") {
+              setMessages(
+                chain.map((m, i) =>
+                  i === chain.length - 1
+                    ? { ...m, reasoning: reasoning || undefined, stats }
+                    : m,
+                ),
+              );
+            }
           },
           defaultErrorMessage: t("error"),
         },
@@ -125,12 +167,14 @@ export default function ChatPage() {
   );
 
   const recentEventIds = useMemo(() => {
+    const responses = toolResponsesById(messages);
     for (let i = messages.length - 1; i >= 0; i--) {
       const msg = messages[i];
-      if (msg.role !== "assistant" || !msg.toolCalls) continue;
-      const similar = getFindSimilarObjectsFromToolCalls(msg.toolCalls);
+      if (msg.role !== "assistant" || !msg.tool_calls?.length) continue;
+      const calls = toolCallsForMessage(msg, responses);
+      const similar = getFindSimilarObjectsFromToolCalls(calls);
       if (similar) return similar.results.map((e) => e.id);
-      const events = getEventIdsFromSearchObjectsToolCalls(msg.toolCalls);
+      const events = getEventIdsFromSearchObjectsToolCalls(calls);
       if (events.length > 0) return events.map((e) => e.id);
     }
     return [];
@@ -154,12 +198,14 @@ export default function ChatPage() {
     abortRef.current?.abort();
     abortRef.current = null;
     setIsLoading(false);
+    setStreaming(null);
   }, []);
 
   const startNewChat = useCallback(() => {
     abortRef.current?.abort();
     abortRef.current = null;
     setIsLoading(false);
+    setStreaming(null);
     setMessages([]);
     setInput("");
     setAttachedEventId(null);
@@ -181,7 +227,83 @@ export default function ChatPage() {
     setAttachedEventId(null);
   }, []);
 
-  const hasStarted = messages.length > 0;
+  const hasStarted = messages.length > 0 || streaming != null;
+
+  // While streaming, the backend's in-flight chain is the source of truth;
+  // otherwise the committed conversation is.
+  const renderList =
+    streaming && streaming.chain.length ? streaming.chain : messages;
+  const responses = toolResponsesById(renderList);
+  const renderTail = renderList[renderList.length - 1];
+  const finalShown =
+    renderTail?.role === "assistant" && hasText(renderTail.content);
+
+  const renderMessage = (msg: ChatMessage, i: number) => {
+    if (msg.role === "system" || msg.role === "tool") return null;
+
+    if (msg.role === "user") {
+      if (!hasText(msg.content)) return null;
+      return (
+        <div key={i} className="flex flex-col gap-2">
+          <MessageBubble
+            role="user"
+            content={msg.content}
+            messageIndex={i}
+            onEditSubmit={handleEditSubmit}
+            isComplete
+            showStats={showStats}
+          />
+        </div>
+      );
+    }
+
+    const calls = toolCallsForMessage(msg, responses);
+    const contentText = hasText(msg.content) ? msg.content : "";
+    const similar = getFindSimilarObjectsFromToolCalls(calls);
+    const events = similar ? [] : getEventIdsFromSearchObjectsToolCalls(calls);
+
+    return (
+      <div key={i} className="flex flex-col gap-2">
+        {calls.length > 0 && <ToolCallsGroup toolCalls={calls} />}
+        {hasText(msg.reasoning) && (
+          <ReasoningBubble
+            reasoning={msg.reasoning}
+            answerStarted={!!contentText}
+          />
+        )}
+        {contentText && (
+          <MessageBubble
+            role="assistant"
+            content={contentText}
+            messageIndex={i}
+            isComplete
+            stats={msg.stats}
+            showStats={showStats}
+          />
+        )}
+        {similar ? (
+          <ChatEventThumbnailsRow
+            events={similar.results}
+            anchor={similar.anchor}
+            onAttach={setAttachedEventId}
+          />
+        ) : (
+          <ChatEventThumbnailsRow
+            events={events}
+            onAttach={setAttachedEventId}
+          />
+        )}
+      </div>
+    );
+  };
+
+  const processingDots = (
+    <div className="flex items-center gap-2 self-start rounded-2xl bg-muted px-5 py-4">
+      <span className="size-2.5 animate-bounce rounded-full bg-muted-foreground/60 [animation-delay:-0.32s]" />
+      <span className="size-2.5 animate-bounce rounded-full bg-muted-foreground/60 [animation-delay:-0.16s]" />
+      <span className="size-2.5 animate-bounce rounded-full bg-muted-foreground/60" />
+    </div>
+  );
 
   return (
     <div className="flex size-full flex-col">
@@ -212,102 +334,31 @@ export default function ChatPage() {
           <div className="flex w-full flex-col xl:w-[50%] 3xl:w-[35%]">
             {hasStarted ? (
               <div className="flex w-full flex-1 flex-col gap-3 pb-3">
-                {messages.map((msg, i) => {
-                  const isLastAssistant =
-                    i === messages.length - 1 && msg.role === "assistant";
-                  const isComplete =
-                    msg.role === "user" || !isLoading || !isLastAssistant;
-                  const hasToolCalls =
-                    msg.toolCalls && msg.toolCalls.length > 0;
-                  const hasContent = !!msg.content?.trim();
-                  const hasReasoning = !!msg.reasoning?.trim();
-                  const showProcessing =
-                    isLastAssistant &&
-                    isLoading &&
-                    !hasContent &&
-                    !hasReasoning;
-
-                  // Hide empty placeholder only when there are no tool calls
-                  // and no reasoning streaming yet
-                  if (
-                    isLastAssistant &&
-                    isLoading &&
-                    !hasContent &&
-                    !hasToolCalls &&
-                    !hasReasoning
-                  )
-                    return (
-                      <div
-                        key={i}
-                        className="flex items-center gap-2 self-start rounded-2xl bg-muted px-5 py-4"
-                      >
-                        <span className="size-2.5 animate-bounce rounded-full bg-muted-foreground/60 [animation-delay:-0.32s]" />
-                        <span className="size-2.5 animate-bounce rounded-full bg-muted-foreground/60 [animation-delay:-0.16s]" />
-                        <span className="size-2.5 animate-bounce rounded-full bg-muted-foreground/60" />
-                      </div>
-                    );
-
-                  return (
-                    <div key={i} className="flex flex-col gap-2">
-                      {msg.role === "assistant" && hasToolCalls && (
-                        <ToolCallsGroup toolCalls={msg.toolCalls!} />
-                      )}
-                      {msg.role === "assistant" && hasReasoning && (
+                {renderList.map((msg, i) => renderMessage(msg, i))}
+                {streaming &&
+                  !finalShown &&
+                  (streaming.content || streaming.reasoning ? (
+                    <div className="flex flex-col gap-2">
+                      {hasText(streaming.reasoning) && (
                         <ReasoningBubble
-                          reasoning={msg.reasoning!}
-                          answerStarted={hasContent}
+                          reasoning={streaming.reasoning}
+                          answerStarted={!!streaming.content}
                         />
                       )}
-                      {showProcessing ? (
-                        <div className="flex items-center gap-2 self-start rounded-2xl bg-muted px-5 py-4">
-                          <span className="size-2 animate-bounce rounded-full bg-muted-foreground/60 [animation-delay:-0.3s]" />
-                          <span className="size-2 animate-bounce rounded-full bg-muted-foreground/60 [animation-delay:-0.15s]" />
-                          <span className="size-2 animate-bounce rounded-full bg-muted-foreground/60" />
-                        </div>
-                      ) : msg.role === "assistant" &&
-                        !hasContent &&
-                        hasReasoning &&
-                        !isComplete ? null : (
+                      {streaming.content && (
                         <MessageBubble
-                          role={msg.role}
-                          content={msg.content}
-                          messageIndex={i}
-                          onEditSubmit={
-                            msg.role === "user" ? handleEditSubmit : undefined
-                          }
-                          isComplete={isComplete}
-                          stats={msg.stats}
+                          role="assistant"
+                          content={streaming.content}
+                          messageIndex={-1}
+                          isComplete={false}
+                          stats={streaming.stats}
                           showStats={showStats}
                         />
                       )}
-                      {msg.role === "assistant" &&
-                        isComplete &&
-                        (() => {
-                          const similar = getFindSimilarObjectsFromToolCalls(
-                            msg.toolCalls,
-                          );
-                          if (similar) {
-                            return (
-                              <ChatEventThumbnailsRow
-                                events={similar.results}
-                                anchor={similar.anchor}
-                                onAttach={setAttachedEventId}
-                              />
-                            );
-                          }
-                          const events = getEventIdsFromSearchObjectsToolCalls(
-                            msg.toolCalls,
-                          );
-                          return (
-                            <ChatEventThumbnailsRow
-                              events={events}
-                              onAttach={setAttachedEventId}
-                            />
-                          );
-                        })()}
                     </div>
-                  );
-                })}
+                  ) : (
+                    processingDots
+                  ))}
                 {error && (
                   <p
                     className="flex items-center gap-1.5 self-start text-sm text-destructive"

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: "system" | "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,20 @@
 import type { ChatMessage, ChatStats, ToolCall } from "@/types/chat";
 
 export type StreamChatCallbacks = {
-  /** Update the messages array (e.g. pass to setState). */
-  updateMessages: (updater: (prev: ChatMessage[]) => ChatMessage[]) => void;
+  /** Streamed delta of the assistant's final answer text. */
+  onContentDelta: (delta: string) => void;
+  /** Streamed delta of the assistant's reasoning trace. */
+  onReasoningDelta: (delta: string) => void;
+  /** The full conversation chain so far (system message, history, this turn's
+   * tool-call turns, tool results, and — on the final emission — the final
+   * assistant message). */
+  onChain: (chain: 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 +29,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 +45,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,
+    onChain,
+    onStats,
     onError,
     onDone,
-    onStats,
     defaultErrorMessage = "Something went wrong. Please try again.",
   } = callbacks;
 
@@ -91,65 +97,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) {
-        updateMessages((prev) => {
-          const next = [...prev];
-          const lastMsg = next[next.length - 1];
-          if (lastMsg?.role === "assistant")
-            next[next.length - 1] = {
-              ...lastMsg,
-              toolCalls: data.tool_calls,
-            };
-          return next;
-        });
+      if (data.type === "messages") {
+        onChain(data.messages ?? []);
         return "continue";
       }
       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;
-        });
+        onContentDelta(data.delta);
         return "continue";
       }
       if (data.type === "reasoning" && data.delta !== undefined) {
-        updateMessages((prev) => {
-          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;
-        });
+        onReasoningDelta(data.delta);
         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 +133,8 @@ export async function streamChatCompletion(
         const trimmed = line.trim();
         if (!trimmed) continue;
         try {
-          const data = JSON.parse(trimmed) as StreamChunk & { type: string };
-          const result = applyChunk(data as StreamChunk);
-          if (result === "break") {
+          const data = JSON.parse(trimmed) as StreamChunk;
+          if (applyChunk(data) === "break") {
             hadStreamError = true;
             break;
           }
@@ -181,50 +148,63 @@ export async function streamChatCompletion(
     // Flush remaining buffer
     if (!hadStreamError && buffer.trim()) {
       try {
-        const data = JSON.parse(buffer.trim()) as StreamChunk & {
-          type: string;
-          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;
-          });
-        }
+        const data = JSON.parse(buffer.trim()) as StreamChunk;
+        applyChunk(data);
       } 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.
  */