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(); + } + }); +});