fix(ai): stop streams getting cut off and keep tool steps across messages

plugins/ai/src/frontend/AiPanel.tsx

@@ -11,7 +11,6 @@ import {
   Textarea,
 } from "@termix/plugin-sdk/ui";
 import {
-  getAiConversation,
   getAiProviders,
   getAiStatus,
   setAiOptIn,
@@ -22,6 +21,12 @@ import { AiMessage } from "./AiMessage";
 import { AiProviderSettings } from "./AiProviderSettings";
 import { AiToolCall } from "./AiToolCall";
 import { ProposalCard } from "./ProposalCard";
+import {
+  buildTimeline,
+  finishedRunEntries,
+  userEntry,
+  type HistoryEntry,
+} from "./transcript";
 import { useAiStream, type ToolActivity } from "./use-ai-stream";
 import {
   activeMentionQuery,
@@ -30,22 +35,6 @@ import {
 } from "./useMentions";
 import { mentionLabel } from "./labels";
 
-interface HistoryEntry {
-  role: "user" | "assistant";
-  content: string;
-}
-
-/** One entry in the rendered conversation, in the order it happened. */
-type TimelineItem =
-  | {
-      kind: "message";
-      key: string;
-      role: "user" | "assistant";
-      content: string;
-    }
-  | { kind: "tool"; key: string; tool: ToolActivity }
-  | { kind: "proposal"; key: string; proposal: AiProposal };
-
 export function AiPanel({ activeTab }: { activeTab?: string | null }) {
   const { t } = useTranslation();
   const { state, send, stop } = useAiStream();
@@ -59,6 +48,7 @@ export function AiPanel({ activeTab }: { activeTab?: string | null }) {
   const [resolvedProposals, setResolvedProposals] = useState<
     Record<number, { status: "applied" | "rejected"; resultSummary?: string }>
   >({});
+  const runCountRef = useRef(0);
   const scrollRef = useRef<HTMLDivElement>(null);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
 
@@ -157,33 +147,15 @@ export function AiPanel({ activeTab }: { activeTab?: string | null }) {
     });
   }, [state.assistantText, state.tools.length, history.length]);
 
-  const reloadHistory = useCallback(
-    async (conversationId: number | null, reply: string) => {
-      // Reload from the server so the transcript is the stored one rather
-      // than a client-side reconstruction.
-      if (conversationId) {
-        try {
-          const loaded = await getAiConversation(conversationId);
-          setHistory(
-            loaded.messages
-              .filter((entry) => entry.role !== "tool")
-              .map((entry) => ({
-                role: entry.role as "user" | "assistant",
-                content: entry.content,
-              })),
-          );
-          return;
-        } catch {
-          // Fall through to keeping the streamed copy.
-        }
-      }
-
-      // No stored copy to fall back on, so the streamed reply is appended
-      // locally. Without this it would vanish when the next turn resets the
-      // stream state.
-      if (reply) {
-        setHistory((prev) => [...prev, { role: "assistant", content: reply }]);
-      }
+  // The run's steps stay in the transcript, so the next message does not wipe
+  // what the assistant already looked at.
+  const keepRun = useCallback(
+    (_conversationId: number | null, reply: string, tools: ToolActivity[]) => {
+      const runId = runCountRef.current++;
+      setHistory((prev) => [
+        ...prev,
+        ...finishedRunEntries(runId, tools, reply),
+      ]);
     },
     [],
   );
@@ -193,14 +165,14 @@ export function AiPanel({ activeTab }: { activeTab?: string | null }) {
     if (!message || !providerId || state.streaming) return;
 
     setInput("");
-    setHistory((prev) => [...prev, { role: "user", content: message }]);
+    setHistory((prev) => [...prev, userEntry(message)]);
 
     await send({
       message,
       providerId,
       conversationId: state.conversationId,
       activeTab,
-      onComplete: reloadHistory,
+      onComplete: keepRun,
     });
   }
 
@@ -233,14 +205,14 @@ export function AiPanel({ activeTab }: { activeTab?: string | null }) {
           : `I approved "${label}". It completed.`
         : `I rejected "${label}". Do not retry it unless I ask.`;
 
-    setHistory((prev) => [...prev, { role: "user", content: followUp }]);
+    setHistory((prev) => [...prev, userEntry(followUp)]);
 
     void send({
       message: followUp,
       providerId,
       conversationId: state.conversationId,
       activeTab,
-      onComplete: reloadHistory,
+      onComplete: keepRun,
     });
   }
 
@@ -254,70 +226,7 @@ export function AiPanel({ activeTab }: { activeTab?: string | null }) {
     };
   });
 
-  /**
-   * The reply exists in two places once a turn finishes: the streamed text in
-   * `state`, and the stored copy that `onComplete` reloads into `history`.
-   * Whichever arrives is shown exactly once, and the stored copy wins because
-   * it sits in the right place in the timeline.
-   */
-  const last = history[history.length - 1];
-  const historyEndsWithThisReply =
-    state.assistantText.length > 0 &&
-    last?.role === "assistant" &&
-    last.content === state.assistantText;
-
-  const streamingReply = historyEndsWithThisReply ? "" : state.assistantText;
-
-  /**
-   * One ordered list rather than three stacked blocks.
-   *
-   * Tool calls happen before the reply that describes them, but once a turn
-   * ends its reply moves into `history`, which used to render above a separate
-   * tools block. The tool call appeared to jump below the answer it produced.
-   * Splicing the current turn's activity in ahead of its own reply keeps the
-   * order stable across that handover.
-   */
-  const timeline: TimelineItem[] = [];
-
-  const trailingReply = historyEndsWithThisReply ? history.length - 1 : -1;
-
-  history.forEach((entry, index) => {
-    if (index === trailingReply) return;
-    timeline.push({
-      kind: "message",
-      key: `history-${index}`,
-      role: entry.role,
-      content: entry.content,
-    });
-  });
-
-  for (const tool of state.tools) {
-    timeline.push({ kind: "tool", key: tool.id, tool });
-  }
-
-  if (trailingReply >= 0) {
-    timeline.push({
-      kind: "message",
-      key: `history-${trailingReply}`,
-      role: "assistant",
-      content: history[trailingReply].content,
-    });
-  } else if (streamingReply) {
-    timeline.push({
-      kind: "message",
-      key: "streaming",
-      role: "assistant",
-      content: streamingReply,
-    });
-  }
-
-  for (const proposal of proposals) {
-    timeline.push({
-      kind: "proposal",
-      key: `proposal-${proposal.id}`,
-      proposal,
-    });
-  }
+  const timeline = buildTimeline(history, state, proposals);
 
   if (loading) {
     return (
@@ -405,7 +314,7 @@ export function AiPanel({ activeTab }: { activeTab?: string | null }) {
               <AiToolCall
                 key={item.key}
                 tool={item.tool}
-                streaming={state.streaming}
+                streaming={item.live && state.streaming}
               />
             );
           }

plugins/ai/src/frontend/terminal/TerminalAiPanel.tsx

@@ -11,6 +11,12 @@ import { Button } from "@termix/plugin-sdk/ui";
 import { getAiProviders, type AiProposal, type AiProvider } from "../ai-api";
 import { AiMessage } from "../AiMessage";
 import { AiToolCall } from "../AiToolCall";
+import {
+  buildTimeline,
+  finishedRunEntries,
+  userEntry,
+  type HistoryEntry,
+} from "../transcript";
 import { useAiStream } from "../use-ai-stream";
 import { TerminalProposalCard } from "./TerminalProposalCard";
 import {
@@ -20,25 +26,6 @@ import {
   type ToolbarPosition,
 } from "./panel-geometry";
 
-interface HistoryEntry {
-  role: "user" | "assistant";
-  content: string;
-}
-
-type TimelineItem =
-  | {
-      kind: "message";
-      key: string;
-      role: "user" | "assistant";
-      content: string;
-    }
-  | {
-      kind: "tool";
-      key: string;
-      tool: import("../use-ai-stream").ToolActivity;
-    }
-  | { kind: "proposal"; key: string; proposal: AiProposal };
-
 interface TerminalAiPanelProps {
   hostLabel: string;
   hostId: number;
@@ -75,6 +62,7 @@ export function TerminalAiPanel({
     Record<number, { status: "applied" | "rejected"; resultSummary?: string }>
   >({});
   const conversationIdRef = useRef<number | null>(null);
+  const runCountRef = useRef(0);
   const scrollRef = useRef<HTMLDivElement>(null);
 
   // Prefilled rather than sent: the user sees exactly what the terminal is
@@ -202,7 +190,7 @@ export function TerminalAiPanel({
     const message = input.trim();
     if (!message || !providerId || state.streaming) return;
 
-    setHistory((current) => [...current, { role: "user", content: message }]);
+    setHistory((current) => [...current, userEntry(message)]);
     setInput("");
 
     void send({
@@ -210,14 +198,13 @@ export function TerminalAiPanel({
       providerId,
       conversationId: conversationIdRef.current,
       activeTab: activeTab ?? `terminal:${hostLabel}`,
-      onComplete: (conversationId, reply) => {
+      onComplete: (conversationId, reply, tools) => {
         conversationIdRef.current = conversationId;
-        if (reply) {
-          setHistory((current) => [
-            ...current,
-            { role: "assistant", content: reply },
-          ]);
-        }
+        const runId = runCountRef.current++;
+        setHistory((current) => [
+          ...current,
+          ...finishedRunEntries(runId, tools, reply),
+        ]);
       },
     });
   }, [activeTab, hostLabel, input, providerId, send, state.streaming]);
@@ -232,58 +219,7 @@ export function TerminalAiPanel({
     };
   });
 
-  // The reply exists in two places once a turn finishes: the streamed text in
-  // `state`, and the copy `onComplete` appended to `history`. Show it once,
-  // and keep tool calls spliced in ahead of the reply they produced instead
-  // of in a separate trailing block, so the timeline reads in the order
-  // things actually happened rather than text-then-tools-then-proposals.
-  const last = history[history.length - 1];
-  const historyEndsWithThisReply =
-    state.assistantText.length > 0 &&
-    last?.role === "assistant" &&
-    last.content === state.assistantText;
-  const streamingReply = historyEndsWithThisReply ? "" : state.assistantText;
-  const trailingReply = historyEndsWithThisReply ? history.length - 1 : -1;
-
-  const timeline: TimelineItem[] = [];
-
-  history.forEach((entry, index) => {
-    if (index === trailingReply) return;
-    timeline.push({
-      kind: "message",
-      key: `history-${index}`,
-      role: entry.role,
-      content: entry.content,
-    });
-  });
-
-  for (const tool of state.tools) {
-    timeline.push({ kind: "tool", key: tool.id, tool });
-  }
-
-  if (trailingReply >= 0) {
-    timeline.push({
-      kind: "message",
-      key: `history-${trailingReply}`,
-      role: "assistant",
-      content: history[trailingReply].content,
-    });
-  } else if (streamingReply) {
-    timeline.push({
-      kind: "message",
-      key: "streaming",
-      role: "assistant",
-      content: streamingReply,
-    });
-  }
-
-  for (const proposal of proposals) {
-    timeline.push({
-      kind: "proposal",
-      key: `proposal-${proposal.id}`,
-      proposal,
-    });
-  }
+  const timeline = buildTimeline(history, state, proposals);
 
   return (
     <div
@@ -371,7 +307,7 @@ export function TerminalAiPanel({
               <AiToolCall
                 key={item.key}
                 tool={item.tool}
-                streaming={state.streaming}
+                streaming={item.live && state.streaming}
               />
             );
           }

plugins/ai/src/frontend/transcript.ts

@@ -0,0 +1,82 @@
+import type { AiProposal } from "./ai-api";
+import type { StreamState, ToolActivity } from "./use-ai-stream";
+
+/** A finished part of the conversation, kept after the stream resets. */
+export type HistoryEntry =
+  | { kind: "message"; role: "user" | "assistant"; content: string }
+  | { kind: "tool"; tool: ToolActivity };
+
+/** One entry in the rendered conversation, in the order it happened. */
+export type TimelineItem =
+  | {
+      kind: "message";
+      key: string;
+      role: "user" | "assistant";
+      content: string;
+    }
+  /** live is false for a past run's step, which can never still be running. */
+  | { kind: "tool"; key: string; tool: ToolActivity; live: boolean }
+  | { kind: "proposal"; key: string; proposal: AiProposal };
+
+export function userEntry(content: string): HistoryEntry {
+  return { kind: "message", role: "user", content };
+}
+
+/**
+ * What a finished run leaves in the transcript: its tool steps, then its
+ * reply. Tool ids are only unique within one run, so they get a run prefix.
+ */
+export function finishedRunEntries(
+  runId: number,
+  tools: ToolActivity[],
+  reply: string,
+): HistoryEntry[] {
+  const entries: HistoryEntry[] = tools.map((tool) => ({
+    kind: "tool",
+    tool: { ...tool, id: `run-${runId}-${tool.id}` },
+  }));
+  if (reply)
+    entries.push({ kind: "message", role: "assistant", content: reply });
+  return entries;
+}
+
+/** Past runs, then the run in progress, then its proposals. */
+export function buildTimeline(
+  history: HistoryEntry[],
+  state: Pick<StreamState, "tools" | "assistantText">,
+  proposals: AiProposal[],
+): TimelineItem[] {
+  const timeline: TimelineItem[] = history.map((entry, index) =>
+    entry.kind === "tool"
+      ? { kind: "tool", key: entry.tool.id, tool: entry.tool, live: false }
+      : {
+          kind: "message",
+          key: `history-${index}`,
+          role: entry.role,
+          content: entry.content,
+        },
+  );
+
+  for (const tool of state.tools) {
+    timeline.push({ kind: "tool", key: tool.id, tool, live: true });
+  }
+
+  if (state.assistantText) {
+    timeline.push({
+      kind: "message",
+      key: "streaming",
+      role: "assistant",
+      content: state.assistantText,
+    });
+  }
+
+  for (const proposal of proposals) {
+    timeline.push({
+      kind: "proposal",
+      key: `proposal-${proposal.id}`,
+      proposal,
+    });
+  }
+
+  return timeline;
+}

plugins/ai/src/frontend/use-ai-stream.ts

@@ -43,9 +43,9 @@ export function useAiStream() {
     setState(INITIAL);
   }, []);
 
+  // The run's own cleanup clears abortRef, after it has handed its steps over.
   const stop = useCallback(() => {
     abortRef.current?.abort();
-    abortRef.current = null;
     setState((prev) => ({ ...prev, streaming: false }));
   }, []);
 
@@ -56,7 +56,15 @@ export function useAiStream() {
       model?: string;
       conversationId?: number | null;
       activeTab?: string | null;
-      onComplete?: (conversationId: number | null, reply: string) => void;
+      /**
+       * Called however the run ends, with its steps, which are then cleared
+       * from the live state. Keeping them is up to the caller.
+       */
+      onComplete?: (
+        conversationId: number | null,
+        reply: string,
+        tools: ToolActivity[],
+      ) => void;
     }) => {
       abortRef.current?.abort();
       const controller = new AbortController();
@@ -74,6 +82,20 @@ export function useAiStream() {
       let conversationId = input.conversationId ?? null;
       let replyText = "";
       let toolSequence = 0;
+      const runTools: ToolActivity[] = [];
+
+      const finish = (error?: string | null) => {
+        // A newer send took over; this run's state is already gone.
+        if (abortRef.current !== controller) return;
+        input.onComplete?.(conversationId, replyText, runTools);
+        setState((prev) => ({
+          ...prev,
+          streaming: false,
+          assistantText: "",
+          tools: [],
+          ...(error ? { error } : {}),
+        }));
+      };
 
       try {
         const response = await aiApp().fetch("chat/stream", {
@@ -152,6 +174,11 @@ export function useAiStream() {
               // prev.tools.length appended a duplicate entry that no result
               // could ever match, leaving it stuck on "running".
               const activityId = `tool-${toolSequence++}`;
+              runTools.push({
+                id: activityId,
+                name: event.name,
+                arguments: event.arguments ?? {},
+              });
               setState((prev) => ({
                 ...prev,
                 tools: [
@@ -164,6 +191,15 @@ export function useAiStream() {
                 ],
               }));
             } else if (event.type === "tool_result") {
+              for (let i = runTools.length - 1; i >= 0; i -= 1) {
+                if (
+                  runTools[i].name === event.name &&
+                  !("result" in runTools[i])
+                ) {
+                  runTools[i] = { ...runTools[i], result: event.result };
+                  break;
+                }
+              }
               setState((prev) => {
                 const tools = [...prev.tools];
                 // Attach to the most recent call of this tool awaiting a result.
@@ -186,20 +222,15 @@ export function useAiStream() {
           }
         }
 
-        setState((prev) => ({ ...prev, streaming: false }));
-        input.onComplete?.(conversationId, replyText);
+        finish();
       } catch (error) {
-        if (controller.signal.aborted) {
-          setState((prev) => ({ ...prev, streaming: false }));
-          return;
-        }
-        setState((prev) => ({
-          ...prev,
-          streaming: false,
-          error: getErrorMessage(error, "The assistant stopped unexpectedly"),
-        }));
+        finish(
+          controller.signal.aborted
+            ? null
+            : getErrorMessage(error, "The assistant stopped unexpectedly"),
+        );
       } finally {
-        abortRef.current = null;
+        if (abortRef.current === controller) abortRef.current = null;
       }
     },
     [],