From 759d5515278389195e2d48122b1ca1bcfedab776 Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Wed, 22 Apr 2026 00:15:56 +0200 Subject: [PATCH] fix(agent): emit parallel tool completion eagerly\n\ncloses #3503 --- packages/agent/CHANGELOG.md | 4 + packages/agent/README.md | 4 +- packages/agent/src/agent-loop.ts | 124 ++++++++++++++--------- packages/agent/src/types.ts | 8 +- packages/agent/test/agent-loop.test.ts | 16 ++- packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/docs/extensions.md | 5 +- 7 files changed, 105 insertions(+), 57 deletions(-) diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index 74cf2afb..d158ba19 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed parallel tool execution to emit `tool_execution_end` as soon as each tool is finalized, while still emitting persisted tool-result messages in assistant source order ([#3503](https://github.com/badlogic/pi-mono/issues/3503)) + ## [0.68.0] - 2026-04-20 ### Changed diff --git a/packages/agent/README.md b/packages/agent/README.md index 21f50a73..daf122d6 100644 --- a/packages/agent/README.md +++ b/packages/agent/README.md @@ -101,10 +101,10 @@ prompt("Read config.json") Tool execution mode is configurable: -- `parallel` (default): preflight tool calls sequentially, execute allowed tools concurrently, and emit `tool_execution_end`, toolResult messages, and `turn_end.toolResults` in the order tool calls finish +- `parallel` (default): preflight tool calls sequentially, execute allowed tools concurrently, emit `tool_execution_end` as soon as each tool is finalized, then emit toolResult messages and `turn_end.toolResults` in assistant source order - `sequential`: execute tool calls one by one, matching the historical behavior -In parallel mode, final tool lifecycle and tool-result artifacts follow tool completion order, not assistant source order. A later tool call may therefore finish before an earlier one, including when it is blocked during preflight. +In parallel mode, tool completion events follow tool completion order, but persisted toolResult messages still follow assistant source order. The mode can be set globally via `toolExecution` in the agent config, or per-tool via `executionMode` on `AgentTool`. If any tool call in a batch targets a tool with `executionMode: "sequential"`, the entire batch executes sequentially regardless of the global setting. diff --git a/packages/agent/src/agent-loop.ts b/packages/agent/src/agent-loop.ts index 35097b00..0c940a1c 100644 --- a/packages/agent/src/agent-loop.ts +++ b/packages/agent/src/agent-loop.ts @@ -369,22 +369,29 @@ async function executeToolCallsSequential( }); const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal); + let finalized: FinalizedToolCallOutcome; if (preparation.kind === "immediate") { - results.push(await emitToolCallOutcome(toolCall, preparation.result, preparation.isError, emit)); + finalized = { + toolCall, + result: preparation.result, + isError: preparation.isError, + }; } else { const executed = await executePreparedToolCall(preparation, signal, emit); - results.push( - await finalizeExecutedToolCall( - currentContext, - assistantMessage, - preparation, - executed, - config, - signal, - emit, - ), + finalized = await finalizeExecutedToolCall( + currentContext, + assistantMessage, + preparation, + executed, + config, + signal, ); } + + await emitToolExecutionEnd(finalized, emit); + const toolResultMessage = createToolResultMessage(finalized); + await emitToolResultMessage(toolResultMessage, emit); + results.push(toolResultMessage); } return results; @@ -398,8 +405,7 @@ async function executeToolCallsParallel( signal: AbortSignal | undefined, emit: AgentEventSink, ): Promise { - const results: ToolResultMessage[] = []; - const runnableCalls: PreparedToolCall[] = []; + const finalizedCalls: FinalizedToolCallEntry[] = []; for (const toolCall of toolCalls) { await emit({ @@ -411,30 +417,39 @@ async function executeToolCallsParallel( const preparation = await prepareToolCall(currentContext, assistantMessage, toolCall, config, signal); if (preparation.kind === "immediate") { - results.push(await emitToolCallOutcome(toolCall, preparation.result, preparation.isError, emit)); - } else { - runnableCalls.push(preparation); + const finalized = { + toolCall, + result: preparation.result, + isError: preparation.isError, + } satisfies FinalizedToolCallOutcome; + await emitToolExecutionEnd(finalized, emit); + finalizedCalls.push(finalized); + continue; } - } - const runningCalls = runnableCalls.map((prepared) => ({ - prepared, - execution: executePreparedToolCall(prepared, signal, emit), - })); - - for (const running of runningCalls) { - const executed = await running.execution; - results.push( - await finalizeExecutedToolCall( + finalizedCalls.push(async () => { + const executed = await executePreparedToolCall(preparation, signal, emit); + const finalized = await finalizeExecutedToolCall( currentContext, assistantMessage, - running.prepared, + preparation, executed, config, signal, - emit, - ), - ); + ); + await emitToolExecutionEnd(finalized, emit); + return finalized; + }); + } + + const orderedFinalizedCalls = await Promise.all( + finalizedCalls.map((entry) => (typeof entry === "function" ? entry() : Promise.resolve(entry))), + ); + const results: ToolResultMessage[] = []; + for (const finalized of orderedFinalizedCalls) { + const toolResultMessage = createToolResultMessage(finalized); + await emitToolResultMessage(toolResultMessage, emit); + results.push(toolResultMessage); } return results; @@ -458,6 +473,14 @@ type ExecutedToolCallOutcome = { isError: boolean; }; +type FinalizedToolCallOutcome = { + toolCall: AgentToolCall; + result: AgentToolResult; + isError: boolean; +}; + +type FinalizedToolCallEntry = FinalizedToolCallOutcome | (() => Promise); + function prepareToolCallArguments(tool: AgentTool, toolCall: AgentToolCall): AgentToolCall { if (!tool.prepareArguments) { return toolCall; @@ -568,8 +591,7 @@ async function finalizeExecutedToolCall( executed: ExecutedToolCallOutcome, config: AgentLoopConfig, signal: AbortSignal | undefined, - emit: AgentEventSink, -): Promise { +): Promise { let result = executed.result; let isError = executed.isError; @@ -599,7 +621,11 @@ async function finalizeExecutedToolCall( } } - return await emitToolCallOutcome(prepared.toolCall, result, isError, emit); + return { + toolCall: prepared.toolCall, + result, + isError, + }; } function createErrorToolResult(message: string): AgentToolResult { @@ -609,31 +635,29 @@ function createErrorToolResult(message: string): AgentToolResult { }; } -async function emitToolCallOutcome( - toolCall: AgentToolCall, - result: AgentToolResult, - isError: boolean, - emit: AgentEventSink, -): Promise { +async function emitToolExecutionEnd(finalized: FinalizedToolCallOutcome, emit: AgentEventSink): Promise { await emit({ type: "tool_execution_end", - toolCallId: toolCall.id, - toolName: toolCall.name, - result, - isError, + toolCallId: finalized.toolCall.id, + toolName: finalized.toolCall.name, + result: finalized.result, + isError: finalized.isError, }); +} - const toolResultMessage: ToolResultMessage = { +function createToolResultMessage(finalized: FinalizedToolCallOutcome): ToolResultMessage { + return { role: "toolResult", - toolCallId: toolCall.id, - toolName: toolCall.name, - content: result.content, - details: result.details, - isError, + toolCallId: finalized.toolCall.id, + toolName: finalized.toolCall.name, + content: finalized.result.content, + details: finalized.result.details, + isError: finalized.isError, timestamp: Date.now(), }; +} +async function emitToolResultMessage(toolResultMessage: ToolResultMessage, emit: AgentEventSink): Promise { await emit({ type: "message_start", message: toolResultMessage }); await emit({ type: "message_end", message: toolResultMessage }); - return toolResultMessage; } diff --git a/packages/agent/src/types.ts b/packages/agent/src/types.ts index 511f15ae..c93134b3 100644 --- a/packages/agent/src/types.ts +++ b/packages/agent/src/types.ts @@ -30,7 +30,8 @@ export type StreamFn = ( * * - "sequential": each tool call is prepared, executed, and finalized before the next one starts. * - "parallel": tool calls are prepared sequentially, then allowed tools execute concurrently. - * Final tool lifecycle and tool-result artifacts are emitted in tool completion order. + * `tool_execution_end` is emitted in tool completion order after each tool is finalized, + * while tool-result message artifacts are emitted later in assistant source order. */ export type ToolExecutionMode = "sequential" | "parallel"; @@ -186,7 +187,8 @@ export interface AgentLoopConfig extends SimpleStreamOptions { * Tool execution mode. * - "sequential": execute tool calls one by one * - "parallel": preflight tool calls sequentially, then execute allowed tools concurrently; - * final tool lifecycle and tool-result artifacts are emitted in tool completion order + * emit `tool_execution_end` in tool completion order after each tool is finalized, + * then emit tool-result message artifacts later in assistant source order * * Default: "parallel" */ @@ -201,7 +203,7 @@ export interface AgentLoopConfig extends SimpleStreamOptions { beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise; /** - * Called after a tool finishes executing, before final tool events are emitted. + * Called after a tool finishes executing, before `tool_execution_end` and tool-result message events are emitted. * * Return an `AfterToolCallResult` to override parts of the executed tool result: * - `content` replaces the full content array diff --git a/packages/agent/test/agent-loop.test.ts b/packages/agent/test/agent-loop.test.ts index 1e20329f..8f5773dc 100644 --- a/packages/agent/test/agent-loop.test.ts +++ b/packages/agent/test/agent-loop.test.ts @@ -449,7 +449,7 @@ describe("agentLoop with AgentMessage", () => { expect(executed).toEqual([[{ oldText: "before", newText: "after" }]]); }); - it("should execute tool calls in parallel and emit tool results in source order", async () => { + it("should emit tool_execution_end in completion order but persist tool results in source order", async () => { const toolSchema = Type.Object({ value: Type.String() }); let firstResolved = false; let parallelObserved = false; @@ -519,15 +519,29 @@ describe("agentLoop with AgentMessage", () => { events.push(event); } + const toolExecutionEndIds = events.flatMap((event) => { + if (event.type !== "tool_execution_end") { + return []; + } + return [event.toolCallId]; + }); const toolResultIds = events.flatMap((event) => { if (event.type !== "message_end" || event.message.role !== "toolResult") { return []; } return [event.message.toolCallId]; }); + const turnToolResultIds = events.flatMap((event) => { + if (event.type !== "turn_end") { + return []; + } + return event.toolResults.map((toolResult) => toolResult.toolCallId); + }); expect(parallelObserved).toBe(true); + expect(toolExecutionEndIds).toEqual(["tool-2", "tool-1"]); expect(toolResultIds).toEqual(["tool-1", "tool-2"]); + expect(turnToolResultIds).toEqual(["tool-1", "tool-2"]); }); it("should inject queued messages after all tool calls complete", async () => { diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index a758cca0..4eeb7988 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixed +- Fixed parallel tool-call rows to leave the pending state as soon as each tool is finalized, while still appending persisted tool results in assistant source order ([#3503](https://github.com/badlogic/pi-mono/issues/3503)) - Fixed exported session markdown to render Markdown while showing HTML-like message content such as `...` verbatim, so shared sessions match the TUI instead of letting the browser interpret message text ([#3484](https://github.com/badlogic/pi-mono/issues/3484)) ## [0.68.0] - 2026-04-20 diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index 66d89a57..de771026 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -547,7 +547,8 @@ Fired for tool execution lifecycle updates. In parallel tool mode: - `tool_execution_start` is emitted in assistant source order during the preflight phase - `tool_execution_update` events may interleave across tools -- `tool_execution_end` is emitted in assistant source order, matching final tool result message order +- `tool_execution_end` is emitted in tool completion order after each tool is finalized +- final `toolResult` message events are still emitted later in assistant source order ```typescript pi.on("tool_execution_start", async (event, ctx) => { @@ -698,6 +699,8 @@ pi.on("tool_call", (event) => { Fired after tool execution finishes and before `tool_execution_end` plus the final tool result message events are emitted. **Can modify result.** +In parallel tool mode, `tool_result` and `tool_execution_end` may interleave in tool completion order, while final `toolResult` message events are still emitted later in assistant source order. + `tool_result` handlers chain like middleware: - Handlers run in extension load order - Each handler sees the latest result after previous handler changes