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.
This commit is contained in:
Khaled Ahmed committed 2026-09-12 06:09:59 +03:00
1 parent cf40f12ab2
commit ae6ca84f26
10 files changed
+415 -15

No files matched your search

+23
View File
@@ -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`
+1
View File
@@ -1963,6 +1963,7 @@ export function AppShell() {
onSessionStatsPanelOpen={openSessionStatsPanel}
onGenerationSpeedChange={handleGenerationSpeedChange}
onOpenProviders={() => setSettingsTab("providers")}
toolCallsDefaultCollapsed={toolCallsDefaultCollapsed}
/>
) : initialCwdStatus === "validating" ? (
<div
+20 -1
View File
@@ -528,6 +528,7 @@ export function ChatWindow({ session, newSessionCwd, newSessionWorkspace, toolCa
notices, dismissNotice, extensionDialog, extensionCustomUi, extensionStatuses, extensionWidgets, respondToExtensionUi, sendExtensionCustomInput,
isAutoModelSelection,
agentPhase, activeGoal, activePlan,
liveToolResults,
subagents, subagentEvents, subagentTranscriptVersions, activeSubagentCount, currentTodoPhase, todoPhases,
isNew,
sessionIdRef, messagesEndRef, scrollContainerRef,
@@ -842,6 +843,23 @@ export function ChatWindow({ session, newSessionCwd, newSessionWorkspace, toolCa
});
return { toolResultsMap, lastAnchorIdx, hasCompaction, visibleRefIndexByMessage };
}, [messages]);
// Runtime tool results span committed messages plus the tool calls omp is
// still executing. A committed result always wins; the live snapshot only
// covers the window between `tool_execution_start` and the toolResult message
// landing, which is what makes the row show a running indicator and streamed
// output instead of a dead "no result" row.
const toolResultsWithLive = useMemo<Map<string, ToolResultMessage>>(() => {
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
<CommittedTranscript
messages={messages}
entryIds={entryIds}
conversationMeta={conversationMeta}
conversationMeta={conversationMetaWithLive}
messageRefs={messageRefs}
isStreaming={streamState.isStreaming}
sessionBusy={sessionBusy}
@@ -1205,6 +1223,7 @@ export function ChatWindow({ session, newSessionCwd, newSessionWorkspace, toolCa
modelNames={modelNames}
cwd={messageCwd}
onOpenFile={onOpenFile}
toolResults={toolResultsWithLive}
toolCallsDefaultCollapsed={toolCallsDefaultCollapsed}
liveTokensPerSecond={tokensPerSecond}
/>
+71
View File
@@ -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,
+28 -5
View File
@@ -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({
<span className={`activity-row-indicator${isError ? " activity-row-indicator-error" : ""}`} aria-hidden>
{isError ? (
<CircleAlert size={12} strokeWidth={1.8} />
) : result ? (
) : result && !isRunning ? (
<Check size={12} strokeWidth={2} />
) : isStreaming ? (
) : isRunning || isStreaming ? (
<LoaderCircle size={12} strokeWidth={1.8} className="activity-row-spinner" />
) : (
<CircleSlash size={12} strokeWidth={1.8} style={{ opacity: 0.5 }} />
@@ -979,7 +992,13 @@ const ToolCallBlock = memo(function ToolCallBlock({
</div>
)}
<TaskResultPanel details={result?.details} />
{result ? (
{isRunning && (resultText ?? "").trim() === "" ? (
// No output yet: say so instead of the "(no output)" marker that
// would claim the tool finished with nothing.
<div data-tool-running="true" style={{ color: "var(--text-dim)", fontSize: 12 }}>
{t("chatWindow.runningTool")}
</div>
) : result ? (
resultDiff ? (
<PairedDiffResult diff={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;
+70
View File
@@ -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();
+85 -3
View File
@@ -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<string | null>(null);
const [compactResult, setCompactResult] = useState<CompactResultInfo | null>(null);
const [agentPhase, setAgentPhase] = useState<AgentPhase>(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<Map<string, ToolResultMessage>>(() => new Map());
const [slashCommands, setSlashCommands] = useState<SlashCommandInfo[]>([]);
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<boolean> => {
@@ -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
// <ChatWindow> (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<string | null>(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,
};
+71
View File
@@ -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);
});
+42 -6
View File
@@ -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<string, CoalescableEvent>();
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();
},
};
+4
View File
@@ -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 {