From ae6ca84f26ce53d677c388ac88762a454c6275a5 Mon Sep 17 00:00:00 2001 From: Khaled Ahmed <61059522+Harvi8@users.noreply.github.com> Date: Sat, 12 Sep 2026 06:09:59 +0300 Subject: [PATCH] feat(chat): stream live tool output while tools run Seed a live tool-result snapshot on tool_execution_start, refresh it on tool_execution_update (coalesced per tool call at display rate), and release it on tool end or the committed toolResult. Rows render a running indicator with streamed output instead of a dead no-result row, and honor the Keep tool calls collapsed setting via toolCallsDefaultCollapsed. --- AGENTS.md | 23 +++++++ components/AppShell.tsx | 1 + components/ChatWindow.tsx | 21 ++++++- components/MessageView.test.mjs | 71 +++++++++++++++++++++ components/MessageView.tsx | 33 ++++++++-- hooks/useAgentSession.rpc.test.mjs | 70 +++++++++++++++++++++ hooks/useAgentSession.ts | 88 ++++++++++++++++++++++++++- lib/message-update-coalescer.test.mjs | 71 +++++++++++++++++++++ lib/message-update-coalescer.ts | 48 +++++++++++++-- lib/types.ts | 4 ++ 10 files changed, 415 insertions(+), 15 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 0b83298..63de232 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -150,6 +150,29 @@ hooks/ ### ToolCall field normalization Sessions store toolCall blocks as `{type:"toolCall", id, name, arguments}` but `ToolCallContent` uses `{toolCallId, toolName, input}`. `normalizeToolCalls()` in `lib/normalize.ts` handles this — called in both `session-reader.ts` (file load) and streaming event handling. +### Live tool execution (`tool_execution_start/update/end`) +omp announces a tool the moment it starts, streams the tool's output while it +runs, and only commits the `toolResult` message at the end. The UI must not +wait for that commit: +- `useAgentSession` keeps a `liveToolResults` map keyed by `toolCallId` + (seeded on `tool_execution_start` with `partial: true`, refreshed on + `tool_execution_update` — omp sends the FULL accumulated partial result per + chunk, latest wins — and released on `_end`/the committed toolResult). + Committed results always win over live entries (`ChatWindow` merges them), so + a reload never shows a stale snapshot. +- `ToolCallBlock` renders a `partial` result as **running** (spinner, and + "Running tool…" instead of the "(no output)" marker when nothing has been + printed yet), and opens the row while it runs when the "Keep tool calls + collapsed" setting is off — that is what that setting means. `AppShell` must + pass `toolCallsDefaultCollapsed` into `ChatWindow`; without it the setting is + inert (the chat then always collapses). +- `tool_execution_update` is coalesced per tool call at display rate in + `lib/message-update-coalescer.ts` (chatty commands emit ~10-100+ frames/s). + `message_end` drops the pending `message_update` (the committed message + supersedes it) but must NOT drop buffered tool updates. +- Live entries are cleared on `agent_start`, terminal `agent_end`, prompt + send/settlement failure — a tool must never leak into the next run. + ### Event protocol differences vs pi omp emits no `prompt_done` / `prompt_error` / `queue_update` / `compaction_start` / `compaction_end` events. Completion is `agent_end` diff --git a/components/AppShell.tsx b/components/AppShell.tsx index b398843..5ce97c9 100644 --- a/components/AppShell.tsx +++ b/components/AppShell.tsx @@ -1963,6 +1963,7 @@ export function AppShell() { onSessionStatsPanelOpen={openSessionStatsPanel} onGenerationSpeedChange={handleGenerationSpeedChange} onOpenProviders={() => setSettingsTab("providers")} + toolCallsDefaultCollapsed={toolCallsDefaultCollapsed} /> ) : initialCwdStatus === "validating" ? (
>(() => { + if (liveToolResults.size === 0) return conversationMeta.toolResultsMap; + const merged = new Map(liveToolResults); + for (const [toolCallId, result] of conversationMeta.toolResultsMap) merged.set(toolCallId, result); + return merged; + }, [liveToolResults, conversationMeta]); + const conversationMetaWithLive = useMemo( + () => (toolResultsWithLive === conversationMeta.toolResultsMap + ? conversationMeta + : { ...conversationMeta, toolResultsMap: toolResultsWithLive }), + [conversationMeta, toolResultsWithLive], + ); // The ref array is sized by the count of user/assistant messages — exactly // what conversationMeta's visibleRefIndexByMessage already tallies, so no // separate filter pass (which would re-run on every streaming frame). @@ -1179,7 +1197,7 @@ export function ChatWindow({ session, newSessionCwd, newSessionWorkspace, toolCa diff --git a/components/MessageView.test.mjs b/components/MessageView.test.mjs index a6404be..e5914d7 100644 --- a/components/MessageView.test.mjs +++ b/components/MessageView.test.mjs @@ -169,6 +169,77 @@ test("advisor custom messages use the localized advisor label", () => { }); +test("a running tool call shows a spinner instead of the no-result marker", () => { + const html = renderToStaticMarkup(React.createElement(MessageView, { + message: { + role: "assistant", + content: [{ type: "toolCall", toolCallId: "call-1", toolName: "bash", input: { command: "long-job" } }], + }, + toolResults: new Map([[ + "call-1", + { role: "toolResult", toolCallId: "call-1", toolName: "bash", content: [], partial: true }, + ]]), + })); + + assert.match(html, /activity-row-spinner/); + assert.doesNotMatch(html, /lucide-check/); + assert.doesNotMatch(html, /lucide-circle-slash/); +}); + +test("a running tool with no output yet says so instead of reporting no output", () => { + const html = renderToStaticMarkup(React.createElement(MessageView, { + toolCallsDefaultCollapsed: false, + message: { + role: "assistant", + content: [{ type: "toolCall", toolCallId: "call-1", toolName: "bash", input: { command: "long-job" } }], + }, + toolResults: new Map([[ + "call-1", + { role: "toolResult", toolCallId: "call-1", toolName: "bash", content: [], partial: true }, + ]]), + })); + + assert.match(html, /data-tool-running="true"/); + assert.doesNotMatch(html, /No output/); +}); + +test("a running tool streams its output before the result is committed", () => { + const html = renderToStaticMarkup(React.createElement(MessageView, { + toolCallsDefaultCollapsed: false, + message: { + role: "assistant", + content: [{ type: "toolCall", toolCallId: "call-1", toolName: "bash", input: { command: "long-job" } }], + }, + toolResults: new Map([[ + "call-1", + { role: "toolResult", toolCallId: "call-1", toolName: "bash", content: [{ type: "text", text: "line-1\nline-2" }], partial: true }, + ]]), + })); + + assert.match(html, /data-tool-output="true"/); + assert.match(html, /line-1/); + assert.match(html, /line-2/); + assert.doesNotMatch(html, /data-tool-running="true"/); +}); + +test("a committed tool result replaces the running affordances", () => { + const html = renderToStaticMarkup(React.createElement(MessageView, { + toolCallsDefaultCollapsed: false, + message: { + role: "assistant", + content: [{ type: "toolCall", toolCallId: "call-1", toolName: "bash", input: { command: "long-job" } }], + }, + toolResults: new Map([[ + "call-1", + { role: "toolResult", toolCallId: "call-1", toolName: "bash", content: [{ type: "text", text: "done" }], timestamp: 3000 }, + ]]), + })); + + assert.doesNotMatch(html, /activity-row-spinner/); + assert.doesNotMatch(html, /data-tool-running="true"/); + assert.match(html, /lucide-check/); +}); + test("expanded edit results with a patch render the split diff view", () => { const html = renderToStaticMarkup(React.createElement(MessageView, { isStreaming: true, diff --git a/components/MessageView.tsx b/components/MessageView.tsx index 11d11f9..b7f5b19 100644 --- a/components/MessageView.tsx +++ b/components/MessageView.tsx @@ -880,7 +880,20 @@ const ToolCallBlock = memo(function ToolCallBlock({ onOpenFile?: (filePath: string) => void; }) { const { t } = useI18n(); - const [expanded, setExpanded] = useState(Boolean(isStreaming) && !defaultCollapsed); + // `partial` results are omp's live snapshots for a tool that is still + // executing (see lib/types.ts); the committed toolResult replaces them. + const isRunning = result?.partial === true; + // A running tool opens its row when the interface keeps tool calls expanded + // ("Keep tool calls collapsed" off) so its output is watchable live. + const [expanded, setExpanded] = useState(Boolean(isStreaming || isRunning) && !defaultCollapsed); + // The row can also mount while the tool is idle and start running later (the + // assistant message commits before `tool_execution_start`). It is never + // auto-collapsed: the output stays where the user was reading it. + const wasRunningRef = useRef(false); + useEffect(() => { + if (isRunning && !wasRunningRef.current && !defaultCollapsed) setExpanded(true); + wasRunningRef.current = isRunning; + }, [isRunning, defaultCollapsed]); const resultText = result ? (typeof result.content === "string" ? result.content @@ -913,9 +926,9 @@ const ToolCallBlock = memo(function ToolCallBlock({ {isError ? ( - ) : result ? ( + ) : result && !isRunning ? ( - ) : isStreaming ? ( + ) : isRunning || isStreaming ? ( ) : ( @@ -979,7 +992,13 @@ const ToolCallBlock = memo(function ToolCallBlock({
)} - {result ? ( + {isRunning && (resultText ?? "").trim() === "" ? ( + // No output yet: say so instead of the "(no output)" marker that + // would claim the tool finished with nothing. +
+ {t("chatWindow.runningTool")} +
+ ) : result ? ( resultDiff ? ( ) : ( @@ -1039,7 +1058,11 @@ const ToolCallGroupBlock = memo(function ToolCallGroupBlock({ const groupSummary = useMemo(() => summarizeToolCallGroup(blocks), [blocks]); const hasError = blocks.some((b) => toolResults?.get(b.toolCallId)?.isError); - const isPending = isStreaming && blocks.some((b) => !toolResults?.has(b.toolCallId)); + // A partial snapshot is a tool still executing, not a settled result. + const isPending = isStreaming && blocks.some((b) => { + const result = toolResults?.get(b.toolCallId); + return !result || result.partial === true; + }); const totalDuration = useMemo(() => { if (!toolCallDurations) return undefined; diff --git a/hooks/useAgentSession.rpc.test.mjs b/hooks/useAgentSession.rpc.test.mjs index a9b403f..5aa0287 100644 --- a/hooks/useAgentSession.rpc.test.mjs +++ b/hooks/useAgentSession.rpc.test.mjs @@ -448,6 +448,76 @@ test("tool activity and turn_end errors do not count as a successful answer", as assert.equal(w.latest.agentRunning, false); }); +test("tool output streams live before the toolResult message lands", async (t) => { + t.after(unmountAll); + resetWorld(); + primeSession("s1", [userMsg("u0", "q")]); + const { w, es } = await startRun(t, "s1", "q1"); + + const toolCallAssistant = { + ...assistantMsg("a1", ""), + content: [{ type: "toolCall", toolCallId: "tc1", toolName: "bash", input: { command: "long-job" } }], + }; + await act(async () => { + es.emit({ type: "agent_start" }); + es.emit({ type: "message_end", message: toolCallAssistant }); + await Promise.resolve(); + }); + assert.equal(w.latest.liveToolResults.size, 0, "nothing is live before the tool starts"); + + // The tool starts: its row must go live immediately, with no output yet. + await act(async () => { + es.emit({ type: "tool_execution_start", toolCallId: "tc1", toolName: "bash", args: { command: "long-job" } }); + await Promise.resolve(); + }); + const started = w.latest.liveToolResults.get("tc1"); + assert.equal(started?.partial, true, "a running tool is a partial result"); + assert.equal(started?.toolName, "bash"); + assert.deepEqual(started?.content, []); + + // omp sends the FULL accumulated output per chunk; only the latest survives + // a display frame. + await act(async () => { + es.emit({ type: "tool_execution_update", toolCallId: "tc1", toolName: "bash", partialResult: { content: [{ type: "text", text: "line-1\n" }] } }); + es.emit({ type: "tool_execution_update", toolCallId: "tc1", toolName: "bash", partialResult: { content: [{ type: "text", text: "line-1\nline-2\n" }] } }); + await Promise.resolve(); + }); + await settle(90); + const streamed = w.latest.liveToolResults.get("tc1"); + assert.equal(streamed?.partial, true); + assert.equal(streamed?.content?.[0]?.text, "line-1\nline-2\n", "latest accumulated snapshot wins"); + + // The tool finishes, then omp commits the toolResult message. The committed + // result supersedes the live snapshot. + await act(async () => { + es.emit({ type: "tool_execution_end", toolCallId: "tc1", toolName: "bash", result: { content: [{ type: "text", text: "line-1\nline-2\n" }] } }); + await Promise.resolve(); + }); + assert.equal(w.latest.liveToolResults.get("tc1")?.partial, undefined, "a finished tool is no longer partial"); + await act(async () => { + es.emit({ + type: "message_end", + message: { role: "toolResult", toolCallId: "tc1", toolName: "bash", content: [{ type: "text", text: "line-1\nline-2\n" }] }, + }); + await Promise.resolve(); + }); + assert.equal(w.latest.liveToolResults.size, 0, "the committed result replaces the live entry"); + assert.equal(w.latest.messages.at(-1)?.role, "toolResult"); + + // Terminal frames clear anything still in flight. + await act(async () => { + es.emit({ type: "tool_execution_start", toolCallId: "tc2", toolName: "bash" }); + await Promise.resolve(); + }); + assert.equal(w.latest.liveToolResults.size, 1); + world.sessions.get("s1").messages = [userMsg("u0", "q"), userMsg("u1", "q1"), toolCallAssistant]; + await act(async () => { + es.emit({ type: "agent_end", isTerminal: true }); + }); + await settle(); + assert.equal(w.latest.liveToolResults.size, 0, "a finished run leaves no live tool state"); +}); + test("late frames after the run finished are ignored (no ghost bubble, no double completion)", async (t) => { t.after(unmountAll); resetWorld(); diff --git a/hooks/useAgentSession.ts b/hooks/useAgentSession.ts index 6434e7c..f34fc3c 100644 --- a/hooks/useAgentSession.ts +++ b/hooks/useAgentSession.ts @@ -8,6 +8,7 @@ import type { ExtensionWidgetItem, SessionInfo, SessionTreeNode, + ToolResultMessage, } from "@/lib/types"; import { normalizeToolCalls } from "@/lib/normalize"; import type { ThinkingModelMeta } from "@/lib/thinking-levels"; @@ -252,6 +253,13 @@ export function useAgentSession(opts: UseAgentSessionOptions) { const [compactError, setCompactError] = useState(null); const [compactResult, setCompactResult] = useState(null); const [agentPhase, setAgentPhase] = useState(null); + // In-flight tool executions, keyed by toolCallId: omp reports a tool's output + // as it runs (`tool_execution_start` / `tool_execution_update` / `_end`), but + // the matching `toolResult` message only lands when the tool finishes. Until + // then these entries keep the tool call row live (running indicator + streamed + // output). The committed toolResult always wins over the live entry, and the + // entry is dropped as soon as it arrives. + const [liveToolResults, setLiveToolResults] = useState>(() => new Map()); const [slashCommands, setSlashCommands] = useState([]); const [slashCommandsLoading, setSlashCommandsLoading] = useState(false); const [noticeState, dispatchNotice] = useReducer(noticeReducer, { visible: [], pending: [] }); @@ -271,6 +279,23 @@ export function useAgentSession(opts: UseAgentSessionOptions) { useEffect(() => () => { clearTerminalReconcileTimer(); }, [clearTerminalReconcileTimer]); + /** Upsert (or with `null`, drop) the live entry for one in-flight tool. */ + const setLiveToolResult = useCallback((toolCallId: string, result: ToolResultMessage | null) => { + setLiveToolResults((prev) => { + if (result === null) { + if (!prev.has(toolCallId)) return prev; + const next = new Map(prev); + next.delete(toolCallId); + return next; + } + const next = new Map(prev); + next.set(toolCallId, result); + return next; + }); + }, []); + const clearLiveToolResults = useCallback(() => { + setLiveToolResults((prev) => (prev.size === 0 ? prev : new Map())); + }, []); // Blocks prompt submission while the initial hydration's live-state fetch is // still pending. Without this, showLoading=false after disk messages lets // a prompt increment runId before the stale pre-prompt state arrives and @@ -1355,6 +1380,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { agentRunningRef.current = false; setAgentRunning(false); setAgentPhase(null); + clearLiveToolResults(); setRetryInfo(null); setSubagents([]); setAdvisorActiveAt(0); @@ -1373,7 +1399,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { slashCommandRunRef.current = false; onAgentEnd?.(); } - }, [addNotice, clearTerminalReconcileTimer, loadSession, onAgentEnd, refreshSubagentHistory, resetSubagentActivityState, surfaceQuotaOnStream]); + }, [addNotice, clearLiveToolResults, clearTerminalReconcileTimer, loadSession, onAgentEnd, refreshSubagentHistory, resetSubagentActivityState, surfaceQuotaOnStream]); const waitForPromptSettlement = useCallback(async (sid: string, runId?: number) => { await delay(PROMPT_SETTLE_INITIAL_DELAY_MS); const startedAt = Date.now(); @@ -1618,6 +1644,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { agentRunningRef.current = true; setAgentRunning(true); setAgentPhase({ kind: "waiting_model" }); + clearLiveToolResults(); dispatch({ type: "start" }); runHadContentRef.current = false; lastQuotaErrorRef.current = null; @@ -1675,6 +1702,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { agentRunningRef.current = false; setAgentRunning(false); setAgentPhase(null); + clearLiveToolResults(); setRetryInfo(null); setSubagents([]); setAdvisorActiveAt(0); @@ -1898,6 +1926,10 @@ export function useAgentSession(opts: UseAgentSessionOptions) { // The advisor model injects its review as a custom message mid-run; // surface it as live advisor activity for the composer thunder icon. if ((completed as CustomMessage).customType === "advisor") setAdvisorActiveAt(Date.now()); + // The committed result supersedes the live snapshot for this call. + if (completed.role === "toolResult" && typeof completed.toolCallId === "string") { + setLiveToolResult(completed.toolCallId, null); + } setMessages((prev) => [...prev, normalizeToolCalls(completed)]); } if (completed?.role === "assistant") { @@ -1930,6 +1962,9 @@ export function useAgentSession(opts: UseAgentSessionOptions) { case "tool_execution_start": { const id = event.toolCallId as string; const name = event.toolName as string; + // Seed the live entry immediately: the row must show a running tool + // from the first frame, before any output exists. + setLiveToolResult(id, { role: "toolResult", toolCallId: id, toolName: name, content: [], partial: true }); setAgentPhase((prev) => { const tools = prev?.kind === "running_tools" ? [...prev.tools] : []; if (!tools.some((t) => t.id === id)) tools.push({ id, name }); @@ -1937,8 +1972,42 @@ export function useAgentSession(opts: UseAgentSessionOptions) { }); break; } + case "tool_execution_update": { + // omp streams the FULL accumulated partial result per output chunk, so + // the newest frame replaces the previous snapshot (the SSE coalescer + // already drops superseded frames within a display frame). + const id = event.toolCallId as string; + const partial = isRecord(event.partialResult) ? event.partialResult : null; + const content = Array.isArray(partial?.content) ? (partial.content as ToolResultMessage["content"]) : []; + const toolName = typeof event.toolName === "string" ? event.toolName : undefined; + setLiveToolResults((prev) => { + const existing = prev.get(id); + const next = new Map(prev); + next.set(id, { + role: "toolResult", + toolCallId: id, + ...(toolName ?? existing?.toolName ? { toolName: toolName ?? existing?.toolName } : {}), + content, + ...(partial?.details !== undefined ? { details: partial.details } : {}), + partial: true, + }); + return next; + }); + break; + } case "tool_execution_end": { const id = event.toolCallId as string; + const finalResult = isRecord(event.result) ? event.result : null; + // Hold the final snapshot (no longer `partial`) until the committed + // toolResult message lands a frame later, so the row never blanks out. + setLiveToolResult(id, { + role: "toolResult", + toolCallId: id, + toolName: typeof event.toolName === "string" ? event.toolName : undefined, + content: Array.isArray(finalResult?.content) ? (finalResult.content as ToolResultMessage["content"]) : [], + isError: finalResult?.isError === true || event.isError === true, + ...(finalResult?.details !== undefined ? { details: finalResult.details } : {}), + }); if (event.toolName === "todo" && sessionIdRef.current) { void reconcileAgentState(sessionIdRef.current); } @@ -2131,7 +2200,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { handleExtensionUiRequest(event as unknown as IncomingExtensionUiRequest); break; } - }, [addNotice, clearTerminalReconcileTimer, consumeQueuedMessage, finishPromptWithoutStream, handleExtensionUiRequest, handleHostToolCall, handleHostUriRequest, loadSession, mergeSubagents, onAgentEnd, reconcileAgentState, resetSubagentActivityState, applyAuthoritativeModel, beginAuthoritativeModelSync, surfaceQuotaOnStream]); + }, [addNotice, clearLiveToolResults, clearTerminalReconcileTimer, consumeQueuedMessage, finishPromptWithoutStream, handleExtensionUiRequest, handleHostToolCall, handleHostUriRequest, loadSession, mergeSubagents, onAgentEnd, reconcileAgentState, resetSubagentActivityState, applyAuthoritativeModel, beginAuthoritativeModelSync, setLiveToolResult, surfaceQuotaOnStream]); handleAgentEventRef.current = handleAgentEvent; const handleSend = useCallback(async (message: string, images?: AttachedImage[]): Promise => { @@ -2171,6 +2240,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { // A new run starts fresh: drop any rescued in-flight reconcile state from // the previous run (its late response is fenced out by the new run id). reconcileGuardRef.current?.reset(); + clearLiveToolResults(); setAgentRunning(true); setAgentPhase(isSlashCommandPrompt ? { kind: "running_command" } : { kind: "waiting_model" }); setAdvisorActiveAt(0); @@ -2263,13 +2333,14 @@ export function useAgentSession(opts: UseAgentSessionOptions) { agentRunningRef.current = false; setAgentRunning(false); setAgentPhase(null); + clearLiveToolResults(); lastQuotaErrorRef.current = null; lastRunErrorRef.current = null; slashCommandRunRef.current = false; dispatch({ type: "end" }); return false; } - }, [isNew, newSessionCwd, newSessionModel, session, ensureNewSession, ensureEventsConnected, promoteNewSession, waitForPromptSettlement, addNotice, opts.chatInputRef, refreshSubagentRoster, registerHostTools, registerHostUriSchemes, clearTerminalReconcileTimer]); + }, [isNew, newSessionCwd, newSessionModel, session, ensureNewSession, ensureEventsConnected, promoteNewSession, waitForPromptSettlement, addNotice, opts.chatInputRef, refreshSubagentRoster, registerHostTools, registerHostUriSchemes, clearLiveToolResults, clearTerminalReconcileTimer]); /** Abort the running agent and send the message as a fresh prompt * (abort_and_prompt). Only valid mid-run; the old turn's agent_end is @@ -2915,8 +2986,18 @@ export function useAgentSession(opts: UseAgentSessionOptions) { }, []); // Load session on mount + // React StrictMode re-invokes this effect for the freshly mounted keyed + // (setup → cleanup → setup), which made every session switch + // fetch and commit the whole transcript twice — in dev, the doubled + // transcript commit is the dominant cost of switching sessions. Latch on the + // session id so one mounted instance loads its session exactly once; the + // surviving first invocation still owns the async continuation (SSE attach, + // roster hydration), which its cleanup does not tear down. + const mountedSessionLoadRef = useRef(null); useEffect(() => { if (session) { + if (mountedSessionLoadRef.current === session.id) return; + mountedSessionLoadRef.current = session.id; sessionIdRef.current = session.id; loadSession(session.id, true, true).then((agentState) => { if (agentState?.running) { @@ -3163,6 +3244,7 @@ export function useAgentSession(opts: UseAgentSessionOptions) { handleToolPresetChange, handleThinkingLevelChange, loadSlashCommands, setActiveLeafId, setData, setMessages, dispatch, setAgentRunning, setForkingEntryId, bashRunning, pendingBash, + liveToolResults, // Subscriptions handleAgentEventRef, }; diff --git a/lib/message-update-coalescer.test.mjs b/lib/message-update-coalescer.test.mjs index 5412908..91b13a1 100644 --- a/lib/message-update-coalescer.test.mjs +++ b/lib/message-update-coalescer.test.mjs @@ -106,3 +106,74 @@ test("reset drops the pending update and cancels the scheduled flush", async () assert.equal(dispatched.length, 0); assert.equal(scheduler.cancelledCount, 1); }); + +test("coalesces tool_execution_update per tool call and keeps each tool's latest", async () => { + const { createMessageUpdateCoalescer } = await loadSubject(); + const dispatched = []; + const scheduler = manualScheduler(); + const coalescer = createMessageUpdateCoalescer((e) => dispatched.push(e), scheduler.schedule); + + coalescer.push({ type: "tool_execution_update", toolCallId: "a", chunk: 1 }); + coalescer.push({ type: "tool_execution_update", toolCallId: "b", chunk: 1 }); + coalescer.push({ type: "tool_execution_update", toolCallId: "a", chunk: 2 }); + assert.equal(dispatched.length, 0); + + scheduler.fire(); + assert.deepEqual(dispatched, [ + { type: "tool_execution_update", toolCallId: "a", chunk: 2 }, + { type: "tool_execution_update", toolCallId: "b", chunk: 1 }, + ]); +}); + +test("flushes buffered tool updates before the next event", async () => { + const { createMessageUpdateCoalescer } = await loadSubject(); + const dispatched = []; + const scheduler = manualScheduler(); + const coalescer = createMessageUpdateCoalescer((e) => dispatched.push(e), scheduler.schedule); + + coalescer.push({ type: "tool_execution_update", toolCallId: "a", chunk: 1 }); + coalescer.push({ type: "tool_execution_end", toolCallId: "a" }); + assert.deepEqual(dispatched.map((e) => e.type), ["tool_execution_update", "tool_execution_end"]); + + scheduler.fire(); + assert.equal(dispatched.length, 2); +}); + +test("message_end does not drop a buffered tool update", async () => { + const { createMessageUpdateCoalescer } = await loadSubject(); + const dispatched = []; + const scheduler = manualScheduler(); + const coalescer = createMessageUpdateCoalescer((e) => dispatched.push(e), scheduler.schedule); + + // The committed assistant message supersedes the partial text, but the tool + // it announced is still executing — its buffered update must survive. + coalescer.push({ type: "message_update", seq: 1 }); + coalescer.push({ type: "tool_execution_update", toolCallId: "a", chunk: 1 }); + coalescer.push({ type: "message_end", seq: 2 }); + assert.deepEqual(dispatched.map((e) => e.type), ["message_end"]); + + scheduler.fire(); + assert.deepEqual(dispatched.map((e) => e.type), ["message_end", "tool_execution_update"]); +}); + +test("a tool update without an id is dispatched immediately", async () => { + const { createMessageUpdateCoalescer } = await loadSubject(); + const dispatched = []; + const scheduler = manualScheduler(); + const coalescer = createMessageUpdateCoalescer((e) => dispatched.push(e), scheduler.schedule); + + coalescer.push({ type: "tool_execution_update", chunk: 1 }); + assert.deepEqual(dispatched.map((e) => e.type), ["tool_execution_update"]); +}); + +test("reset drops buffered tool updates", async () => { + const { createMessageUpdateCoalescer } = await loadSubject(); + const dispatched = []; + const scheduler = manualScheduler(); + const coalescer = createMessageUpdateCoalescer((e) => dispatched.push(e), scheduler.schedule); + + coalescer.push({ type: "tool_execution_update", toolCallId: "a", chunk: 1 }); + coalescer.reset(); + scheduler.fire(); + assert.equal(dispatched.length, 0); +}); diff --git a/lib/message-update-coalescer.ts b/lib/message-update-coalescer.ts index 3931e42..89389d6 100644 --- a/lib/message-update-coalescer.ts +++ b/lib/message-update-coalescer.ts @@ -5,11 +5,17 @@ // one re-renders the whole streaming bubble. Only the latest pending update is // worth showing; buffer it and flush at animation-frame rate. // +// `tool_execution_update` frames are coalesced the same way (omp sends the full +// accumulated partial result roughly per output chunk — measured ~10-100+/s on +// chatty commands), keyed by toolCallId so concurrent tools keep their own +// latest snapshot. +// // Ordering contract: -// - Any non-update event type flushes the pending update synchronously BEFORE +// - Any non-update event type flushes the pending updates synchronously BEFORE // it is dispatched, so no state is applied out of order. // - `message_end` carries the complete message and therefore supersedes -// (drops) any pending partial update. +// (drops) any pending partial update. It does NOT drop pending tool updates: +// those belong to a tool the committed message merely announced. export type CoalescableEvent = { type: string; [key: string]: unknown }; @@ -42,6 +48,9 @@ export function createMessageUpdateCoalescer( schedule: FlushScheduler = defaultScheduler, ): MessageUpdateCoalescer { let pending: CoalescableEvent | null = null; + // Latest snapshot per in-flight tool. Map preserves insertion order, so + // concurrent tools dispatch in the order they first reported progress. + const pendingToolUpdates = new Map(); let cancelScheduled: (() => void) | null = null; const cancel = () => { @@ -54,30 +63,57 @@ export function createMessageUpdateCoalescer( const flush = () => { cancelScheduled = null; const event = pending; + const toolEvents = [...pendingToolUpdates.values()]; pending = null; + pendingToolUpdates.clear(); if (event) dispatch(event); + for (const toolEvent of toolEvents) dispatch(toolEvent); + }; + + const scheduleFlush = () => { + if (!cancelScheduled) cancelScheduled = schedule(flush); }; return { push(event: CoalescableEvent) { if (event.type === "message_update") { pending = event; - if (!cancelScheduled) cancelScheduled = schedule(flush); + scheduleFlush(); + return; + } + if (event.type === "tool_execution_update") { + // A frame without an id cannot be keyed; dispatch it straight away + // rather than letting two tools overwrite each other's slot. + const toolCallId = typeof event.toolCallId === "string" ? event.toolCallId : null; + if (!toolCallId) { + flush(); + dispatch(event); + return; + } + pendingToolUpdates.set(toolCallId, event); + scheduleFlush(); return; } if (event.type === "message_end") { + // The complete message supersedes any pending partial message update. + // A buffered tool update belongs to the tool the committed message + // announced, so it survives and its scheduled flush stays armed. pending = null; - cancel(); - } else if (pending) { + if (pendingToolUpdates.size === 0) cancel(); + } else if (pending || pendingToolUpdates.size > 0) { cancel(); const buffered = pending; + const toolEvents = [...pendingToolUpdates.values()]; pending = null; - dispatch(buffered); + pendingToolUpdates.clear(); + if (buffered) dispatch(buffered); + for (const toolEvent of toolEvents) dispatch(toolEvent); } dispatch(event); }, reset() { pending = null; + pendingToolUpdates.clear(); cancel(); }, }; diff --git a/lib/types.ts b/lib/types.ts index 0347bb8..b0b800f 100644 --- a/lib/types.ts +++ b/lib/types.ts @@ -104,6 +104,10 @@ export interface ToolResultMessage { isError?: boolean; details?: unknown; timestamp?: number; + /** Live snapshot of a tool that is still executing (omp's + * `tool_execution_update`), not a committed result. Rendering treats it as + * "running" and keeps it out of the session's toolResult list. */ + partial?: boolean; } export interface CustomMessage {