agent,prime-agent-runtime: coalesce tool progress updates and kernel output frames
283 unmodified linesExpand all2842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303312360 unmodified linesExpand all283 unmodified linesExpand all self.assertLess(events.index(raw), done_index) tagged = next(e for e in events if e.get("event") == "stdout" and "tagged" in e["text"]) self.assertEqual(tagged["id"], "rawfd")
def test_print_burst_ships_few_frames_in_order(self): # A print-per-line cell measured ~2.5M frames and froze the host's client; the # tagged channel batches inside a live cell instead of shipping per print(). lines = 4000 events = self.repl.execute("burst", f"for i in range({lines}):\n print('line-%05d' % i)") frames = [e for e in events if e.get("event") == "stdout" and e.get("id") == "burst"] self.assertEqual(stream_text(events, "stdout"), "".join(f"line-{i:05d}\n" for i in range(lines))) # ~11 bytes per line against an 8 KiB budget: a handful of frames, not 4000. self.assertLess(len(frames), 50) self.assertLess(events.index(frames[-1]), events.index(one(events, "done")))
def test_quiet_cell_output_ships_before_done_without_further_writes(self): # One print followed by silent work must still stream: the settled-batch # flusher ships it, not a later write and not the cell's `done`. with tempfile.TemporaryDirectory() as tmp: gate = os.path.join(tmp, "go") code = "\n".join( [ "import os, time", "print('early-line')", f"while not os.path.exists({gate!r}):", " time.sleep(0.01)", ] ) self.repl.send({"type": "execute", "id": "live", "code": code}) seen: list[dict] = [] deadline = time.monotonic() + 5 while time.monotonic() < deadline: if any(e.get("event") == "stdout" and "early-line" in e["text"] for e in seen): break try: seen.append(self.repl.read_event(timeout=0.5)) except TimeoutError: continue self.assertIn("early-line", stream_text(seen, "stdout")) self.assertIsNone(one(seen, "done")) with open(gate, "w"): pass self.assertEqual(one(seen + self.repl.until_done("live"), "done")["status"], "ok")
def test_drain_finalizes_incomplete_raw_utf8_before_done(self): first = self.repl.execute("partial-utf8", "import os\nos.write(1, b'\\xe2\\x82')") replacement = next(e for e in first if e.get("event") == "stdout") self.assertEqual(replacement["text"], "\ufffd")2360 unmodified linesExpand all12- Changed tool progress updates to coalesce and rate limit, so one chatty tool call streams a bounded number of events instead of one per output chunk.- Fixed kernel stdout/stderr shipping one protocol frame per write, which froze the client on print-heavy cells; the runtime now batches output while a cell is live.25 unmodified lines2627282930313233341 unmodified line36373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109797 unmodified linesExpand all9079089099108448459119129139149159169179189199209219229239248518529259269278548558568578588598608618628638649289299308689319328708718728739339349359369371 unmodified line93994094194288388488588688794394494594694794888 unmodified lines25 unmodified lines
export type AgentEventSink = (event: AgentEvent) => Promise<void> | void;
const ABORT_ERROR_MESSAGE = "Request was aborted";const TOOL_UPDATE_MIN_INTERVAL_MS = 50;const EMPTY_USAGE: AssistantMessage["usage"] = { input: 0, output: 0, cacheRead: 0,1 unmodified line totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },};
class ToolUpdateRelay { private accepting = true; private latest: AgentToolResult<unknown> | undefined; private inFlight: Promise<void> | undefined; private failure: { error: unknown } | undefined; private lastEmitAtMs = Number.NEGATIVE_INFINITY;
constructor(private readonly emitUpdate: (partialResult: AgentToolResult<unknown>) => Promise<void> | void) {}
publish(partialResult: AgentToolResult<unknown>): void { if (!this.accepting) { return; } this.latest = partialResult; this.emitLatest(false); }
close(): void { this.accepting = false; }
async flush(): Promise<void> { await this.inFlight; this.throwIfFailed(); this.emitLatest(true); await this.inFlight; this.throwIfFailed(); }
private emitLatest(force: boolean): void { if (this.inFlight || this.failure || !this.latest) { return; } const nowMs = performance.now(); if (!force && nowMs - this.lastEmitAtMs < TOOL_UPDATE_MIN_INTERVAL_MS) { return; }
const partialResult = this.latest; this.latest = undefined; this.lastEmitAtMs = nowMs; let emission: Promise<void> | void; try { emission = this.emitUpdate(partialResult); } catch (error) { this.failure = { error }; throw error; } this.inFlight = Promise.resolve(emission).then( () => { this.inFlight = undefined; }, (error: unknown) => { this.failure = { error }; this.inFlight = undefined; }, ); }
private throwIfFailed(): void { if (this.failure) { throw this.failure.error; } }}
function createAbortError(): Error { return new Error(ABORT_ERROR_MESSAGE);}
797 unmodified linesExpand all prepared: PreparedToolCall, signal: AbortSignal | undefined, emit: AgentEventSink,): Promise<ExecutedToolCallOutcome> { const updateEvents: Promise<void>[] = []; let acceptingUpdates = true; const updates = new ToolUpdateRelay((partialResult) => emit({ type: "tool_execution_update", toolCallId: prepared.toolCall.id, toolName: prepared.toolCall.name, args: prepared.toolCall.arguments, partialResult, }), );
try { throwIfAborted(signal); const result = await raceWithAbort( prepared.tool.execute(prepared.toolCall.id, prepared.args as never, signal, (partialResult) => { if (!acceptingUpdates || signal?.aborted) { return; if (!signal?.aborted) { updates.publish(partialResult); } updateEvents.push( Promise.resolve( emit({ type: "tool_execution_update", toolCallId: prepared.toolCall.id, toolName: prepared.toolCall.name, args: prepared.toolCall.arguments, partialResult, }), ), ); }), signal, ); acceptingUpdates = false; updates.close(); try { await raceWithAbort( Promise.all(updateEvents).then(() => undefined), signal, ); await raceWithAbort(updates.flush(), signal); } catch (error) { if (!signal?.aborted || !isAbortError(error)) { throw error; }1 unmodified line // A tool that returns instead of throwing still gets to declare failure; without // this, a tool's own isError never reaches the tool result or the event. return { result, isError: result.isError === true }; } catch (error) { acceptingUpdates = false; await raceWithAbort( Promise.all(updateEvents).then(() => undefined), signal, ).catch(() => undefined); updates.close(); await raceWithAbort(updates.flush(), signal).catch(() => undefined); return { result: createErrorToolResult( signal?.aborted ? "Tool execution aborted" : error instanceof Error ? error.message : String(error), ),88 unmodified lines354 unmodified linesExpand all35535635735835935936036136236336436536664 unmodified lines354 unmodified linesExpand all */ terminate?: boolean;}
/** Callback used by tools to publish partial execution updates. *//** * Callback used by tools to publish coalesced, rate-limited execution updates. * Only the newest partial is guaranteed delivery; tools must not rely on every partial reaching clients. */export type AgentToolUpdateCallback<T = any> = (partialResult: AgentToolResult<T>) => void;
/** Tool definition used by the agent runtime. */export interface AgentTool<TParameters extends TSchema = TSchema, TDetails = any> extends Tool<TParameters> {64 unmodified lines426 unmodified linesExpand all4274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964971733 unmodified linesExpand all426 unmodified linesExpand all expect(toolResultIds).toEqual(["tool_1"]); expect(events.some((event) => event.type === "agent_end")).toBe(true); });
it("should coalesce synchronous tool update bursts and flush the newest partial", async () => { const toolSchema = Type.Object({}); const tool: AgentTool<typeof toolSchema, { index: number }> = { name: "work", label: "Work", description: "Work", parameters: toolSchema, execute: async (_toolCallId, _params, _signal, onUpdate) => { for (let index = 0; index < 1000; index++) { onUpdate?.({ content: [{ type: "text", text: `progress:${index}` }], details: { index } }); } return { content: [{ type: "text", text: "done" }], details: { index: 1000 } }; }, }; const context: AgentContext = { systemPrompt: "You are helpful.", messages: [], tools: [tool], }; const config: AgentLoopConfig = { model: createModel(), convertToLlm: identityConverter, }; const assistantMessage = createAssistantMessage( [{ type: "toolCall", id: "tool_1", name: "work", arguments: {} }], "toolUse", ); let streamCall = 0; const streamFn = () => { const stream = new MockAssistantStream(); queueMicrotask(() => { if (streamCall++ === 0) { stream.push({ type: "done", reason: "toolUse", message: assistantMessage }); } else { const done = createAssistantMessage([{ type: "text", text: "done" }]); stream.push({ type: "done", reason: "stop", message: done }); } }); return stream; }; const events: AgentEvent[] = [];
await runAgentLoop( [createUserMessage("Hello")], context, config, (event) => { events.push(event); }, undefined, streamFn, );
const updates = events.flatMap((event) => (event.type === "tool_execution_update" ? [event] : [])); const toolEndIndex = events.findIndex((event) => event.type === "tool_execution_end"); expect(updates).toHaveLength(2); expect(updates.at(-1)?.partialResult).toEqual({ content: [{ type: "text", text: "progress:999" }], details: { index: 999 }, }); expect(events.indexOf(updates.at(-1)!)).toBeLessThan(toolEndIndex); });
it("should preserve a successful tool result when abort fires during update flush", async () => { const controller = new AbortController(); const toolSchema = Type.Object({}); const tool: AgentTool<typeof toolSchema, Record<string, never>> = {1733 unmodified linesExpand all