diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index c1c02515c..287f07c30 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed late tool progress callbacks after tool settlement to be ignored instead of emitting stale `tool_execution_update` events ([#5573](https://github.com/earendil-works/pi/issues/5573)). + ## [0.79.1] - 2026-06-09 ## [0.79.0] - 2026-06-08 diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 28f037f5d..77bd9568a 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -631,6 +631,7 @@ async function executePreparedToolCall( emit: AgentEventSink, ): Promise { const updateEvents: Promise[] = []; + let acceptingUpdates = true; try { const result = await prepared.tool.execute( @@ -638,6 +639,7 @@ async function executePreparedToolCall( prepared.args as never, signal, (partialResult) => { + if (!acceptingUpdates) return; updateEvents.push( Promise.resolve( emit({ @@ -651,14 +653,18 @@ async function executePreparedToolCall( ); }, ); + acceptingUpdates = false; await Promise.all(updateEvents); return { result, isError: false }; } catch (error) { + acceptingUpdates = false; await Promise.all(updateEvents); return { result: createErrorToolResult(error instanceof Error ? error.message : String(error)), isError: true, }; + } finally { + acceptingUpdates = false; } } diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index f24d24963..cb99a79a8 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -354,7 +354,12 @@ export interface AgentToolResult { terminate?: boolean; } -/** Callback used by tools to stream partial execution updates. */ +/** + * Callback used by tools to stream partial execution updates. + * + * The callback is scoped to the current `execute()` invocation. Calls made after + * the tool promise settles are ignored. + */ export type AgentToolUpdateCallback = (partialResult: AgentToolResult) => void; /** Tool definition used by the agent runtime. */ diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index 82cc58dee..4cc51f744 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -1,6 +1,7 @@ import { type AssistantMessage, type AssistantMessageEvent, EventStream, getModel } from "@earendil-works/pi-ai"; +import { Type } from "typebox"; import { describe, expect, it } from "vitest"; -import { Agent } from "../src/index.ts"; +import { Agent, type AgentEvent, type AgentTool, type AgentToolUpdateCallback } from "../src/index.ts"; // Mock stream that mimics AssistantMessageEventStream class MockAssistantStream extends EventStream { @@ -36,6 +37,28 @@ function createAssistantMessage(text: string): AssistantMessage { }; } +type ToolCallContent = Extract; + +function createAssistantToolUseMessage(content: ToolCallContent[]): AssistantMessage { + return { + role: "assistant", + content, + api: "openai-responses", + provider: "openai", + model: "mock", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "toolUse", + timestamp: Date.now(), + }; +} + function createDeferred(): { promise: Promise; resolve: () => void; @@ -242,6 +265,147 @@ describe("Agent", () => { expect(receivedSignal?.aborted).toBe(true); }); + it("should ignore tool updates after the tool execution settles", async () => { + const toolSchema = Type.Object({}); + let delayedUpdate: AgentToolUpdateCallback<{ status: string }> | undefined; + const events: AgentEvent[] = []; + const unhandledRejections: unknown[] = []; + const onUnhandledRejection = (error: unknown) => { + unhandledRejections.push(error); + }; + const tool: AgentTool = { + name: "delayed_tool", + label: "Delayed Tool", + description: "Captures progress callbacks", + parameters: toolSchema, + async execute(_toolCallId, _params, _signal, onUpdate) { + delayedUpdate = onUpdate; + onUpdate?.({ + content: [{ type: "text", text: "running" }], + details: { status: "running" }, + }); + return { + content: [{ type: "text", text: "ok" }], + details: { status: "done" }, + terminate: true, + }; + }, + }; + const agent = new Agent({ + initialState: { tools: [tool] }, + streamFn: () => { + const stream = new MockAssistantStream(); + queueMicrotask(() => { + stream.push({ + type: "done", + reason: "toolUse", + message: createAssistantToolUseMessage([ + { type: "toolCall", id: "call-1", name: "delayed_tool", arguments: {} }, + ]), + }); + }); + return stream; + }, + }); + agent.subscribe((event) => { + events.push(event); + }); + + process.on("unhandledRejection", onUnhandledRejection); + try { + await agent.prompt("run tool"); + const eventCountAfterPrompt = events.length; + + delayedUpdate?.({ + content: [{ type: "text", text: "late" }], + details: { status: "late" }, + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + + expect(events.filter((event) => event.type === "tool_execution_update")).toHaveLength(1); + expect(events).toHaveLength(eventCountAfterPrompt); + expect(unhandledRejections).toEqual([]); + } finally { + process.off("unhandledRejection", onUnhandledRejection); + } + }); + + it("should ignore a settled parallel tool update while another tool is still running", async () => { + const toolSchema = Type.Object({}); + const slowStarted = createDeferred(); + const settledToolEnded = createDeferred(); + const releaseSlow = createDeferred(); + let settledToolUpdate: AgentToolUpdateCallback<{ status: string }> | undefined; + const events: AgentEvent[] = []; + const settledTool: AgentTool = { + name: "settled_tool", + label: "Settled Tool", + description: "Captures progress callbacks", + parameters: toolSchema, + async execute(_toolCallId, _params, _signal, onUpdate) { + settledToolUpdate = onUpdate; + return { + content: [{ type: "text", text: "done" }], + details: { status: "done" }, + terminate: true, + }; + }, + }; + const slowTool: AgentTool = { + name: "slow_tool", + label: "Slow Tool", + description: "Keeps the agent run active", + parameters: toolSchema, + async execute() { + slowStarted.resolve(); + await releaseSlow.promise; + return { + content: [{ type: "text", text: "done" }], + details: { status: "done" }, + terminate: true, + }; + }, + }; + const agent = new Agent({ + initialState: { tools: [settledTool, slowTool] }, + streamFn: () => { + const stream = new MockAssistantStream(); + queueMicrotask(() => { + stream.push({ + type: "done", + reason: "toolUse", + message: createAssistantToolUseMessage([ + { type: "toolCall", id: "call-1", name: "settled_tool", arguments: {} }, + { type: "toolCall", id: "call-2", name: "slow_tool", arguments: {} }, + ]), + }); + }); + return stream; + }, + }); + agent.subscribe((event) => { + events.push(event); + if (event.type === "tool_execution_end" && event.toolCallId === "call-1") { + settledToolEnded.resolve(); + } + }); + + const promptPromise = agent.prompt("run tools"); + await Promise.all([slowStarted.promise, settledToolEnded.promise]); + const eventCountBeforeLateUpdate = events.length; + + settledToolUpdate?.({ + content: [{ type: "text", text: "late" }], + details: { status: "late" }, + }); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(events).toHaveLength(eventCountBeforeLateUpdate); + + releaseSlow.resolve(); + await promptPromise; + expect(events.filter((event) => event.type === "tool_execution_update")).toHaveLength(0); + }); + it("should update state with mutators", () => { const agent = new Agent();