From 240eb29c41b2f5c6e7fcab80bc2b090d871ccaad Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 25 Aug 2026 12:13:50 +0200 Subject: [PATCH] fix(coding-agent): append run-time custom messages after the turn's tool results Extension messages sent with triggerTurn: false while the agent was running were pushed into agent state and the session tree immediately, landing between an assistant tool call and its result. Providers that validate message order reject the replayed history. Queue them during the run and append at turn_end, mirroring recordBashResult. closes #8537 --- packages/coding-agent/CHANGELOG.md | 4 + .../coding-agent/src/core/agent-session.ts | 60 ++++++-- ...ustom-message-tool-result-ordering.test.ts | 144 ++++++++++++++++++ 3 files changed, 197 insertions(+), 11 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/8537-custom-message-tool-result-ordering.test.ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index d3885a628..491aa86b4 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Fixed + +- Fixed extension messages sent with `triggerTurn: false` while the agent is running being inserted between a tool call and its result, which made providers that validate message order reject the replayed history. They are now appended once the turn's tool results are in ([#8537](https://github.com/earendil-works/pi/issues/8537)). + ## [0.84.3] - 2026-08-24 ### New Features diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 488f5e0ae..60e278d34 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -327,6 +327,8 @@ export class AgentSession { private _followUpMessages: string[] = []; /** Messages queued to be included with the next user prompt as context ("asides"). */ private _pendingNextTurnMessages: CustomMessage[] = []; + /** Context-only custom messages queued during a run, flushed once the current turn's tool results are in. */ + private _pendingCustomMessages: CustomMessage[] = []; // Compaction state private _compactionAbortController: AbortController | undefined = undefined; @@ -689,6 +691,15 @@ export class AgentSession { } } } + + // A turn ends after its assistant message and every tool result has been appended, + // so this is the first point in the run where a context-only custom message can be + // inserted without landing between a tool call and its result. Flushing after the + // extension and listener dispatch above also picks up messages that turn_end + // handlers queued. + if (event.type === "turn_end") { + this._flushPendingCustomMessages(); + } }; private _willRetryAfterAgentEnd(event: Extract): boolean { @@ -1081,6 +1092,7 @@ export class AgentSession { } finally { this._systemPromptOverride = undefined; this._flushPendingBashMessages(); + this._flushPendingCustomMessages(); await this._emitAgentSettled(); } } @@ -1190,8 +1202,9 @@ export class AgentSession { return; } - // Flush any pending bash messages before the new prompt + // Flush any pending bash and custom messages before the new prompt this._flushPendingBashMessages(); + this._flushPendingCustomMessages(); // Validate model if (!this.model) { @@ -1436,8 +1449,9 @@ export class AgentSession { /** * Send a custom message to the session. Creates a CustomMessageEntry. * - * Handles three cases: + * Handles four cases: * - Streaming: queues message, processed when loop pulls from queue + * - Streaming + triggerTurn false: appended to state/session once the current turn ends * - Not streaming + triggerTurn: appends to state/session, starts new turn * - Not streaming + no trigger: appends to state/session, no turn * @@ -1468,16 +1482,40 @@ export class AgentSession { } } else if (options?.triggerTurn) { await this._runAgentPrompt(appMessage); + } else if (this.isStreaming) { + // Appending now would put the message between an assistant tool call and its + // result, which providers that validate message order reject on replay. Defer + // to the end of the turn. Nothing is emitted yet: message events must not + // describe messages the session tree does not contain. + this._pendingCustomMessages.push(appMessage); } else { - this.agent.state.messages.push(appMessage); - this.sessionManager.appendCustomMessageEntry( - message.customType, - message.content, - message.display, - message.details, - ); - this._emit({ type: "message_start", message: appMessage }); - this._emit({ type: "message_end", message: appMessage }); + this._appendCustomMessage(appMessage); + } + } + + private _appendCustomMessage(appMessage: CustomMessage): void { + this.agent.state.messages.push(appMessage); + this.sessionManager.appendCustomMessageEntry( + appMessage.customType, + appMessage.content, + appMessage.display, + appMessage.details, + ); + this._emit({ type: "message_start", message: appMessage }); + this._emit({ type: "message_end", message: appMessage }); + } + + /** + * Append custom messages queued while the agent was running. + * Called once the current turn's tool results are in agent state and session history. + */ + private _flushPendingCustomMessages(): void { + if (this._pendingCustomMessages.length === 0) return; + + const pending = this._pendingCustomMessages; + this._pendingCustomMessages = []; + for (const appMessage of pending) { + this._appendCustomMessage(appMessage); } } diff --git a/packages/coding-agent/test/suite/regressions/8537-custom-message-tool-result-ordering.test.ts b/packages/coding-agent/test/suite/regressions/8537-custom-message-tool-result-ordering.test.ts new file mode 100644 index 000000000..161f10da6 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/8537-custom-message-tool-result-ordering.test.ts @@ -0,0 +1,144 @@ +import type { AgentMessage, AgentTool } from "@earendil-works/pi-agent-core"; +import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; +import { Type } from "typebox"; +import { afterEach, describe, expect, it } from "vitest"; +import { convertToLlm } from "../../../src/core/messages.ts"; +import { createHarness, type Harness } from "../harness.ts"; + +function roles(messages: AgentMessage[]): string[] { + return messages.map((message) => message.role); +} + +describe("#8537 custom messages injected during tool execution", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) { + harnesses.pop()?.cleanup(); + } + }); + + it("appends the message after the turn's tool results instead of between call and result", async () => { + let notify: (() => Promise) | undefined; + const slowTool: AgentTool = { + name: "wait", + label: "Wait", + description: "Wait for a background task", + parameters: Type.Object({}), + execute: async () => { + // A background task (e.g. a subagent reply) notifies the session while the + // tool is still running. + await notify?.(); + return { content: [{ type: "text", text: "tool done" }], details: {} }; + }, + }; + + const harness = await createHarness({ tools: [slowTool] }); + harnesses.push(harness); + notify = () => + harness.session.sendCustomMessage( + { customType: "subagent-reply", content: "subagent replied", display: true }, + { triggerTurn: false }, + ); + + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("wait", {})], { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + ]); + + await harness.session.prompt("hi"); + + expect(roles(harness.session.messages)).toEqual(["user", "assistant", "toolResult", "custom", "assistant"]); + }); + + it("keeps session entries and message events in the same order as agent state", async () => { + let notify: (() => Promise) | undefined; + const slowTool: AgentTool = { + name: "wait", + label: "Wait", + description: "Wait for a background task", + parameters: Type.Object({}), + execute: async () => { + await notify?.(); + return { content: [{ type: "text", text: "tool done" }], details: {} }; + }, + }; + + const harness = await createHarness({ tools: [slowTool] }); + harnesses.push(harness); + notify = () => + harness.session.sendCustomMessage( + { customType: "subagent-reply", content: "subagent replied", display: true }, + { triggerTurn: false }, + ); + + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("wait", {})], { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + ]); + + await harness.session.prompt("hi"); + + const entryKinds = harness.sessionManager + .getBranch() + .flatMap((entry) => + entry.type === "message" ? [entry.message.role] : entry.type === "custom_message" ? ["custom"] : [], + ); + expect(entryKinds).toEqual(["user", "assistant", "toolResult", "custom", "assistant"]); + + // message events must never describe a message the session tree does not contain yet + const messageStarts = harness.events.flatMap((event) => + event.type === "message_start" ? [event.message.role] : [], + ); + expect(messageStarts).toEqual(["user", "assistant", "toolResult", "custom", "assistant"]); + }); + + it("produces an llm history where every tool result follows its tool call", async () => { + let notify: (() => Promise) | undefined; + const slowTool: AgentTool = { + name: "wait", + label: "Wait", + description: "Wait for a background task", + parameters: Type.Object({}), + execute: async () => { + await notify?.(); + return { content: [{ type: "text", text: "tool done" }], details: {} }; + }, + }; + + const harness = await createHarness({ tools: [slowTool] }); + harnesses.push(harness); + notify = () => + harness.session.sendCustomMessage( + { customType: "subagent-reply", content: "subagent replied", display: true }, + { triggerTurn: false }, + ); + + harness.setResponses([ + fauxAssistantMessage([fauxToolCall("wait", {})], { stopReason: "toolUse" }), + fauxAssistantMessage("done"), + fauxAssistantMessage("second turn"), + ]); + + await harness.session.prompt("hi"); + await harness.session.prompt("and now?"); + + const llmMessages = convertToLlm(harness.session.messages); + const openToolCallIds = new Set(); + for (const message of llmMessages) { + if (message.role === "assistant") { + openToolCallIds.clear(); + for (const block of message.content) { + if (block.type === "toolCall") openToolCallIds.add(block.id); + } + continue; + } + if (message.role === "toolResult") { + expect(openToolCallIds.has(message.toolCallId)).toBe(true); + openToolCallIds.delete(message.toolCallId); + continue; + } + openToolCallIds.clear(); + } + }); +});