From a447ba11f3e7760fcff91e924bda8c20033953ea Mon Sep 17 00:00:00 2001 From: David Brailovsky Date: Mon, 12 Oct 2026 09:53:24 +0200 Subject: [PATCH] fix(coding-agent): emit before_agent_start for runs started by custom messages (#10739) Runs started by pi.sendMessage(..., { triggerTurn: true }) skipped before_agent_start. After their first tool call, the next-turn refresh fell back to the base prompt options and patched away the sections handlers had added, changing the system prompt mid-run and breaking the provider prompt cache. These runs now go through the same setup as prompt(): compaction check, before_agent_start, tool loadout, nextTurn messages, handler messages, and the system patch. BeforeAgentStartEvent is a union on messageType ("user" or "custom"); custom events carry the starting message. The session counts as busy from the start of setup, so messages sent while handlers run go into the starting run instead of starting a second one, and an abort requested meanwhile aborts the run once it starts. closes #10267, closes #5581, closes #10639, closes #10289 partly fixes #10606 --- packages/coding-agent/CHANGELOG.md | 3 + packages/coding-agent/docs/extensions.md | 3 +- .../examples/extensions/plan-mode/index.ts | 8 +- .../coding-agent/src/core/agent-session.ts | 158 +++++++++------- .../coding-agent/src/core/extensions/index.ts | 2 + .../src/core/extensions/runner.ts | 13 +- .../coding-agent/src/core/extensions/types.ts | 26 ++- packages/coding-agent/src/index.ts | 2 + .../test/agent-session-concurrent.test.ts | 10 +- .../test/extensions-runner.test.ts | 11 +- .../test/suite/agent-session-prompt.test.ts | 168 +++++++++++++++++- 11 files changed, 312 insertions(+), 92 deletions(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index afa6243f0..a029fa70d 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -5,6 +5,8 @@ ### Breaking Changes - Published generated configuration schemas for models, settings, keybindings, and themes at `https://pi.dev/schemas/*.schema.json`. The theme schema moved to `https://pi.dev/schemas/theme.schema.json`; themes loaded during startup now receive the same schema validation as later reloads. Update theme `$schema` references, define reusable colors under `vars`, and remove unsupported metadata. +- Changed `before_agent_start` to also fire for runs started by `pi.sendMessage(..., { triggerTurn: true })` while idle. `BeforeAgentStartEvent` is now a union on `messageType`: `"user"` events carry `prompt` and `images`, and `"custom"` events carry the starting `message` and no `prompt`. Messages returned by handlers are added to both kinds of run, and pending `deliverAs: "nextTurn"` messages are delivered with the custom-message run. Handlers that should act only on user prompts must check `event.messageType === "user"` ([#10267](https://github.com/earendil-works/pi/issues/10267), [#5581](https://github.com/earendil-works/pi/issues/5581)) +- Changed the session to count as busy while a run starts, including during `before_agent_start` handlers: `ctx.isIdle()` is false, `pi.sendMessage()` queues into the starting run instead of starting a second run that failed with "Agent is already processing", `pi.sendUserMessage()` requires `deliverAs`, and an abort requested meanwhile aborts the run once it starts ### Added @@ -23,6 +25,7 @@ - Fixed Mistral streams being aborted mid-generation (e.g. during long thinking) after the request timeout; the timeout now applies only to waiting for response headers ([#10609](https://github.com/earendil-works/pi/issues/10609)) - Fixed a symlinked `AGENTS.md` in a git worktree nested inside its main repo causing no project context file to load ([#10681](https://github.com/earendil-works/pi/issues/10681)) - Fixed the code mode tool description claiming every failed tool call rejects; a tool that declares an output schema resolves with its value on failure too (e.g. `bash` with a nonzero exit) +- Fixed runs started by `pi.sendMessage(..., { triggerTurn: true })` removing system prompt sections added in `before_agent_start` after their first tool call, which changed the prompt mid-run and broke the provider prompt cache twice per run ([#10267](https://github.com/earendil-works/pi/issues/10267)) ## [1.1.0] - 2026-10-07 diff --git a/packages/coding-agent/docs/extensions.md b/packages/coding-agent/docs/extensions.md index e43654555..c0c7fd23c 100644 --- a/packages/coding-agent/docs/extensions.md +++ b/packages/coding-agent/docs/extensions.md @@ -60,6 +60,7 @@ Start long-lived resources from `session_start` or from the command or tool that Close session-scoped resources from an idempotent `session_shutdown` handler. A run proceeds from input and `before_agent_start`, through model, message, and tool events, to `agent_end`. +Pi counts as busy while a run starts, including the compaction check and `before_agent_start`, so messages sent meanwhile go into the starting run. Automatic retries, recovery, compaction, or queued work can continue afterward. @@ -100,7 +101,7 @@ Use each event’s declared result type rather than assuming every return value Events cover resource discovery, sessions, agent and message lifecycle, providers, tools, and raw input. -`before_agent_start` exposes both the current prompt and its structured `systemPromptOptions`. Prefer changing prompt sections, selected tools, or guidelines so Pi can append a transcript delta. Returning `systemPrompt`, or setting `forceSystemPrompt`, replaces the whole prompt for that run while the transcript continues recording the structured sections. Providers receive the forced text as their leading system prompt. +`before_agent_start` exposes both the current prompt and its structured `systemPromptOptions`. `event.messageType` says what started the run: `"user"` for a prompt or `pi.sendUserMessage()`, with `prompt` and `images`; `"custom"` for `pi.sendMessage(..., { triggerTurn: true })`, with the starting `message` and no `prompt`. A returned `message` is added to either kind of run; check `messageType` to skip runs started by custom messages. Prefer changing prompt sections, selected tools, or guidelines so Pi can append a transcript delta. Returning `systemPrompt`, or setting `forceSystemPrompt`, replaces the whole prompt for that run while the transcript continues recording the structured sections. Providers receive the forced text as their leading system prompt. `message_end` can replace a finalized message while preserving its role. `tool_call` can mutate input or block execution. `tool_result` handlers compose, with each handler seeing prior changes. diff --git a/packages/coding-agent/examples/extensions/plan-mode/index.ts b/packages/coding-agent/examples/extensions/plan-mode/index.ts index 737ce56ac..bbeb7ab13 100644 --- a/packages/coding-agent/examples/extensions/plan-mode/index.ts +++ b/packages/coding-agent/examples/extensions/plan-mode/index.ts @@ -197,8 +197,12 @@ export default function planModeExtension(pi: ExtensionAPI): void { }; }); - // Inject plan/execution context before agent starts - pi.on("before_agent_start", async () => { + // Inject plan/execution context before a user prompt. Runs started by custom messages, such as + // background task notifications, skip it: the context from the last prompt is still in the + // conversation, and adding it again on every notification would only repeat it. + pi.on("before_agent_start", async (event) => { + if (event.messageType !== "user") return; + if (planModeEnabled) { return { message: { diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index c16121c8e..ca21a5473 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -112,7 +112,7 @@ import { type TurnStartEvent, wrapRegisteredTools, } from "./extensions/index.ts"; -import { emitSessionShutdownEvent } from "./extensions/runner.ts"; +import { type BeforeAgentStartTrigger, emitSessionShutdownEvent } from "./extensions/runner.ts"; import { createToolNameMatcher, isMcpToolName } from "./mcp-servers.ts"; import { type BashExecutionMessage, type CustomMessage, convertToLlm } from "./messages.ts"; import { ModelRegistry } from "./model-registry.ts"; @@ -1101,6 +1101,9 @@ export class AgentSession { /** Internal handler for agent events - shared by subscribe and reconnect */ private _handleAgentEvent = async (event: AgentEvent): Promise => { + // An abort requested while the run was starting, e.g. during before_agent_start, stops it + // here. agent_start is awaited before the first request. + if (event.type === "agent_start" && this._agentRunAbortRequested) this.agent.abort(); // Record the calls a tool made through ctx.executeTool() and their usage on its result message. if (this._nestedToolCalls) { if (event.type === "message_start" && event.message.role === "toolResult") { @@ -1820,15 +1823,37 @@ export class AgentSession { // Prompting // ========================================================================= - private async _runAgentPrompt(messages: AgentMessage | AgentMessage[]): Promise { + /** + * Start a run from a user or custom message. The session counts as busy from here on, so + * messages sent while before_agent_start handlers run are queued into this run instead of + * starting another, and an abort requested meanwhile aborts the run as soon as it exists. + * @param onStarted Called once the run's messages are prepared, before the first request. + */ + private async _startRun(trigger: BeforeAgentStartTrigger, onStarted?: () => void): Promise { + this._isAgentRunActive = true; this._agentRunAbortRequested = false; + let messages: AgentMessage[]; + try { + messages = await this._prepareRun(trigger); + } catch (error) { + this._isAgentRunActive = false; + this._flushPendingBashMessages(); + this._flushPendingCustomMessages(); + this._resolveIdleWaitIfIdle(); + throw error; + } + // Messages held while the run was starting go before its starting messages, as they would + // have while idle. No run exists yet, so they cannot split a tool call from its result. + this._flushPendingBashMessages(); + this._flushPendingCustomMessages(); + onStarted?.(); + // Compaction before the prompt may have scheduled a retry; the new prompt replaces it. this._failedResponse = undefined; this._recordSelection(); // The run records the loadout in the transcript; restored tools that did not register by now // are dropped, so a tool that never registers does not stay pending. this._pendingToolNames.clear(); - this._isAgentRunActive = true; try { await this.agent.prompt(messages); while (!this._agentRunAbortRequested) { @@ -1851,6 +1876,65 @@ export class AgentSession { } } + /** + * Prepare a run and return its starting messages: compact first when needed, emit + * before_agent_start, apply the run's prompt options and tool loadout, and build the starting + * message followed by pending nextTurn messages (consumed here) and handler messages, with a + * system patch first when the prompt changed. + */ + private async _prepareRun(trigger: BeforeAgentStartTrigger): Promise { + // Check if we need to compact before sending (catches aborted responses). + // The starting message is sent below, so do not call agent.continue() here. + const lastAssistant = this._findLastAssistantMessage(); + if (lastAssistant) { + await this._checkCompaction(lastAssistant, false); + } + + const selectedToolsBefore = this._baseSystemPromptOptions.selectedTools; + const result = await this._extensionRunner.emitBeforeAgentStart(trigger, this._baseSystemPromptOptions); + // Handlers may edit event.systemPromptOptions.selectedTools or call setActiveTools(), + // which updates the live loadout instead. An explicit edit wins; otherwise the live + // loadout is authoritative, so a setActiveTools() call is not undone here. + const handlerEditedTools = + result.systemPromptOptions.selectedTools.length !== selectedToolsBefore.length || + result.systemPromptOptions.selectedTools.some((name, index) => name !== selectedToolsBefore[index]); + if (!handlerEditedTools) result.systemPromptOptions.selectedTools = this.getActiveToolNames(); + + const messages: AgentMessage[] = []; + if (trigger.messageType === "custom") { + messages.push(trigger.message); + } else { + // Normalize images after before_agent_start so extension-driven model selection + // determines the resize profile used for the request and history. + const normalized = await this._normalizePromptImages(trigger.images); + const text = + normalized.hints.length > 0 ? `${trigger.prompt}\n\n${normalized.hints.join("\n")}` : trigger.prompt; + messages.push({ + role: "user", + content: [{ type: "text", text }, ...normalized.images], + timestamp: Date.now(), + }); + } + // Inject any pending "nextTurn" messages as context alongside the starting message + messages.push(...this._pendingNextTurnMessages); + this._pendingNextTurnMessages = []; + for (const msg of result.messages) { + messages.push({ + role: "custom", + customType: msg.customType, + // Untyped extensions can pass null/missing content; normalize at ingestion. + content: msg.content ?? [], + display: msg.display, + details: msg.details, + timestamp: Date.now(), + }); + } + const updateMessage = this._preparePromptAndToolLoadout(result.systemPromptOptions); + this._runSystemPromptOptions = result.systemPromptOptions; + if (updateMessage) messages.unshift(updateMessage); + return messages; + } + private async _handlePostAgentRun(): Promise { const message = this._lastAssistantMessage; const toolResults = this._lastAssistantToolResults; @@ -2026,10 +2110,6 @@ export class AgentSession { return; } - // Flush any pending bash and custom messages before the new prompt - this._flushPendingBashMessages(); - this._flushPendingCustomMessages(); - // Validate model if (!this.model) { throw new Error(formatNoModelSelectedMessage()); @@ -2050,65 +2130,9 @@ export class AgentSession { throw new Error(formatNoApiKeyFoundMessage(this.model.provider)); } - // Check if we need to compact before sending (catches aborted responses). - // The user's new prompt is sent below, so do not call agent.continue() here. - const lastAssistant = this._findLastAssistantMessage(); - if (lastAssistant) { - await this._checkCompaction(lastAssistant, false); - } - - // Emit before_agent_start before normalizing images so extension-driven model - // selection determines the resize profile used for the request and history. - const selectedToolsBefore = this._baseSystemPromptOptions.selectedTools; - const result = await this._extensionRunner.emitBeforeAgentStart( - expandedText, - currentImages, - this._baseSystemPromptOptions, + await this._startRun({ messageType: "user", prompt: expandedText, images: currentImages }, () => + preflightResult?.("started"), ); - // Handlers may edit event.systemPromptOptions.selectedTools or call setActiveTools(), - // which updates the live loadout instead. An explicit edit wins; otherwise the live - // loadout is authoritative, so a setActiveTools() call is not undone here. - const handlerEditedTools = - result.systemPromptOptions.selectedTools.length !== selectedToolsBefore.length || - result.systemPromptOptions.selectedTools.some((name, index) => name !== selectedToolsBefore[index]); - if (!handlerEditedTools) result.systemPromptOptions.selectedTools = this.getActiveToolNames(); - - const normalized = await this._normalizePromptImages(currentImages); - const userText = normalized.hints.length > 0 ? `${expandedText}\n\n${normalized.hints.join("\n")}` : expandedText; - - // Build messages only after hooks and image normalization have completed. - const messages: AgentMessage[] = []; - const userContent: (TextContent | ImageContent)[] = [{ type: "text", text: userText }]; - userContent.push(...normalized.images); - messages.push({ - role: "user", - content: userContent, - timestamp: Date.now(), - }); - - // Inject any pending "nextTurn" messages as context alongside the user message - for (const msg of this._pendingNextTurnMessages) { - messages.push(msg); - } - this._pendingNextTurnMessages = []; - - for (const msg of result.messages) { - messages.push({ - role: "custom", - customType: msg.customType, - // Untyped extensions can pass null/missing content; normalize at ingestion. - content: msg.content ?? [], - display: msg.display, - details: msg.details, - timestamp: Date.now(), - }); - } - const updateMessage = this._preparePromptAndToolLoadout(result.systemPromptOptions); - this._runSystemPromptOptions = result.systemPromptOptions; - if (updateMessage) messages.unshift(updateMessage); - - preflightResult?.("started"); - await this._runAgentPrompt(messages); } /** @@ -2314,10 +2338,10 @@ export class AgentSession { } } else if (options?.triggerTurn) { if (this._isEmittingAgentSettled) { - this._deferredSettledActions.push(async () => await this._runAgentPrompt(appMessage)); + this._deferredSettledActions.push(async () => await this.sendCustomMessage(message, options)); return; } - await this._runAgentPrompt(appMessage); + await this._startRun({ messageType: "custom", message: 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 diff --git a/packages/coding-agent/src/core/extensions/index.ts b/packages/coding-agent/src/core/extensions/index.ts index aebd5c60d..1465e8267 100644 --- a/packages/coding-agent/src/core/extensions/index.ts +++ b/packages/coding-agent/src/core/extensions/index.ts @@ -37,8 +37,10 @@ export type { // Events - Tool (ToolCallEvent types) BashToolCallEvent, BashToolResultEvent, + BeforeAgentStartCustomEvent, BeforeAgentStartEvent, BeforeAgentStartEventResult, + BeforeAgentStartUserEvent, BeforeProviderHeadersEvent, BeforeProviderRequestEvent, BeforeProviderRequestEventResult, diff --git a/packages/coding-agent/src/core/extensions/runner.ts b/packages/coding-agent/src/core/extensions/runner.ts index 11ffaca7e..77712e677 100644 --- a/packages/coding-agent/src/core/extensions/runner.ts +++ b/packages/coding-agent/src/core/extensions/runner.ts @@ -27,8 +27,10 @@ import { import type { VirtualModelDefinition } from "../virtual-models.ts"; import type { AgentBeforeSettleEvent, + BeforeAgentStartCustomEvent, BeforeAgentStartEvent, BeforeAgentStartEventResult, + BeforeAgentStartUserEvent, BeforeProviderHeadersEvent, BeforeProviderRequestEvent, BoundaryContextPreview, @@ -163,6 +165,11 @@ function isUserBashEventResult(value: unknown): value is UserBashEventResult { ); } +/** The message that starts a run: the before_agent_start event without its prompt fields. */ +export type BeforeAgentStartTrigger = + | Pick + | Pick; + /** Combined result from all before_agent_start handlers. */ interface BeforeAgentStartCombinedResult { messages: NonNullable[]; @@ -1418,8 +1425,7 @@ export class ExtensionRunner { } async emitBeforeAgentStart( - prompt: string, - images: ImageContent[] | undefined, + trigger: BeforeAgentStartTrigger, systemPromptOptions: BuildSystemPromptOptions, ): Promise { const currentOptions = normalizeBuildSystemPromptOptions(systemPromptOptions); @@ -1438,9 +1444,8 @@ export class ExtensionRunner { for (const handler of handlers) { try { const event: BeforeAgentStartEvent = { + ...trigger, type: "before_agent_start", - prompt, - images, get systemPrompt() { return renderCurrentSystemPrompt(); }, diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index fd24fbccc..60388c0eb 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -929,19 +929,33 @@ export interface ProviderStreamEvent { data: unknown; } -/** Fired after user submits prompt but before agent loop. */ -export interface BeforeAgentStartEvent { +interface BeforeAgentStartEventBase { type: "before_agent_start"; - /** The raw user prompt text (after expansion). */ - prompt: string; - /** Images attached to the user prompt, if any. */ - images?: ImageContent[]; /** The current system prompt, rendered from systemPromptOptions and earlier handler changes. */ readonly systemPrompt: string; /** Mutable prompt sections. Later handlers observe mutations made by earlier handlers. */ systemPromptOptions: NormalizedBuildSystemPromptOptions; } +/** A run started by a submitted prompt, RPC prompt, or `pi.sendUserMessage()`. */ +export interface BeforeAgentStartUserEvent extends BeforeAgentStartEventBase { + messageType: "user"; + /** The raw user prompt text (after expansion). */ + prompt: string; + /** Images attached to the user prompt, if any. */ + images?: ImageContent[]; +} + +/** A run started by `pi.sendMessage(..., { triggerTurn: true })` while idle. */ +export interface BeforeAgentStartCustomEvent extends BeforeAgentStartEventBase { + messageType: "custom"; + /** The custom message that starts the run. */ + message: CustomMessage; +} + +/** Fired before a run starts from a user message or an idle custom message with `triggerTurn`. */ +export type BeforeAgentStartEvent = BeforeAgentStartUserEvent | BeforeAgentStartCustomEvent; + /** Fired when an agent loop starts */ export interface AgentStartEvent { type: "agent_start"; diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 3a25f85c6..bd48684cf 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -64,8 +64,10 @@ export type { AppKeybinding, AutocompleteProviderFactory, BashToolCallEvent, + BeforeAgentStartCustomEvent, BeforeAgentStartEvent, BeforeAgentStartEventResult, + BeforeAgentStartUserEvent, BeforeProviderHeadersEvent, BeforeProviderRequestEvent, BeforeProviderRequestEventResult, diff --git a/packages/coding-agent/test/agent-session-concurrent.test.ts b/packages/coding-agent/test/agent-session-concurrent.test.ts index 3f1dc9723..58c144716 100644 --- a/packages/coding-agent/test/agent-session-concurrent.test.ts +++ b/packages/coding-agent/test/agent-session-concurrent.test.ts @@ -444,8 +444,7 @@ describe("AgentSession concurrent prompt guard", () => { streamingBehavior?: "steer" | "followUp", ) => Promise<{ action: "continue" }>; emitBeforeAgentStart: ( - prompt: string, - images: unknown, + trigger: unknown, systemPromptOptions: BuildSystemPromptOptions, ) => Promise<{ messages: []; systemPromptOptions: NormalizedBuildSystemPromptOptions }>; invalidate: (message?: string) => void; @@ -465,7 +464,7 @@ describe("AgentSession concurrent prompt guard", () => { return undefined; }, emitInput: async () => ({ action: "continue" }), - emitBeforeAgentStart: async (_prompt, _images, systemPromptOptions) => ({ + emitBeforeAgentStart: async (_trigger, systemPromptOptions) => ({ messages: [], systemPromptOptions: normalizeBuildSystemPromptOptions(systemPromptOptions), }), @@ -591,8 +590,7 @@ describe("AgentSession concurrent prompt guard", () => { streamingBehavior?: "steer" | "followUp", ) => Promise<{ action: "continue" }>; emitBeforeAgentStart: ( - prompt: string, - images: unknown, + trigger: unknown, systemPromptOptions: BuildSystemPromptOptions, ) => Promise<{ messages: []; systemPromptOptions: NormalizedBuildSystemPromptOptions }>; invalidate: (message?: string) => void; @@ -608,7 +606,7 @@ describe("AgentSession concurrent prompt guard", () => { return undefined; }, emitInput: async () => ({ action: "continue" }), - emitBeforeAgentStart: async (_prompt, _images, systemPromptOptions) => ({ + emitBeforeAgentStart: async (_trigger, systemPromptOptions) => ({ messages: [], systemPromptOptions: normalizeBuildSystemPromptOptions(systemPromptOptions), }), diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 51d535adc..7c4a5f428 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -877,10 +877,13 @@ describe("ExtensionRunner", () => { runner.onError((error) => errors.push(error.error)); runner.bindCore(extensionActions, extensionContextActions); - const chained = await runner.emitBeforeAgentStart("hello", undefined, { - cwd: tempDir, - customPrompt: "base", - }); + const chained = await runner.emitBeforeAgentStart( + { messageType: "user", prompt: "hello" }, + { + cwd: tempDir, + customPrompt: "base", + }, + ); expect(errors).toEqual([]); expect(chained.messages).toEqual([]); diff --git a/packages/coding-agent/test/suite/agent-session-prompt.test.ts b/packages/coding-agent/test/suite/agent-session-prompt.test.ts index 7361383f3..8052fd646 100644 --- a/packages/coding-agent/test/suite/agent-session-prompt.test.ts +++ b/packages/coding-agent/test/suite/agent-session-prompt.test.ts @@ -2,10 +2,16 @@ import { mkdirSync, rmSync, writeFileSync } from "node:fs"; import { tmpdir } from "node:os"; import { join } from "node:path"; import type { AgentTool } from "@earendil-works/pi-agent-core"; -import { fauxAssistantMessage, fauxToolCall, type Model } from "@earendil-works/pi-ai"; +import { + fauxAssistantMessage, + fauxToolCall, + getCurrentSystemPrompt, + type Model, + type TranscriptContext, +} from "@earendil-works/pi-ai"; import { Type } from "typebox"; import { afterEach, describe, expect, it, vi } from "vitest"; -import type { ExtensionAPI, InputEvent } from "../../src/core/extensions/index.ts"; +import type { BeforeAgentStartEvent, ExtensionAPI, InputEvent } from "../../src/core/extensions/index.ts"; import type { PromptTemplate } from "../../src/core/prompt-templates.ts"; import { createSyntheticSourceInfo } from "../../src/core/source-info.ts"; import { createTestResourceLoader } from "../utilities.ts"; @@ -556,4 +562,162 @@ describe("AgentSession prompt characterization", () => { `No API key found for ${harness.getModel().provider}.`, ); }); + + // Issue #10267, #5581: a run started by sendMessage(..., { triggerTurn: true }) skipped + // before_agent_start, so its second request patched away the sections handlers had added. + describe("runs started by a custom message", () => { + it("keeps the same system prompt for every request of the run", async () => { + const events: BeforeAgentStartEvent[] = []; + const noopTool: AgentTool = { + name: "noop", + label: "Noop", + description: "Does nothing", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {} }), + }; + const harness = await createHarness({ + tools: [noopTool], + extensionFactories: [ + (pi) => { + pi.on("before_agent_start", (event) => { + events.push(event); + event.systemPromptOptions.sections.extension_rules = "Answer in one word."; + }); + }, + ], + }); + harnesses.push(harness); + const prompts: string[] = []; + const record = (response: ReturnType) => (context: TranscriptContext) => { + prompts.push(getCurrentSystemPrompt(context.messages)); + return response; + }; + harness.setResponses([ + record(fauxAssistantMessage("hi")), + record(fauxAssistantMessage([fauxToolCall("noop", {})], { stopReason: "toolUse" })), + record(fauxAssistantMessage("done")), + ]); + + await harness.session.prompt("hello"); + await harness.session.sendCustomMessage( + { customType: "task-done", content: "background task finished", display: true }, + { triggerTurn: true }, + ); + + expect(prompts).toHaveLength(3); + expect(prompts[0]).toContain("Answer in one word."); + expect(new Set(prompts).size).toBe(1); + expect(events.map((event) => event.messageType)).toEqual(["user", "custom"]); + const customEvent = events[1]; + expect(customEvent.messageType === "custom" && customEvent.message.customType).toBe("task-done"); + }); + + it("queues a custom message sent while before_agent_start runs into the starting run", async () => { + let releaseHandler: () => void = () => {}; + const handlerStarted = Promise.withResolvers(); + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("before_agent_start", async (event) => { + if (event.messageType !== "custom" || event.message.customType !== "first") return; + handlerStarted.resolve(); + await new Promise((resolve) => { + releaseHandler = resolve; + }); + }); + }, + ], + }); + harnesses.push(harness); + const requests: TranscriptContext[] = []; + harness.setResponses([ + (context) => { + requests.push(context); + return fauxAssistantMessage("handled both"); + }, + ]); + + const first = harness.session.sendCustomMessage( + { customType: "first", content: "task one finished", display: true }, + { triggerTurn: true }, + ); + await handlerStarted.promise; + expect(harness.session.isIdle).toBe(false); + const second = harness.session.sendCustomMessage( + { customType: "second", content: "task two finished", display: true }, + { triggerTurn: true }, + ); + releaseHandler(); + await Promise.all([first, second]); + + expect(requests).toHaveLength(1); + const texts = requests[0].messages.flatMap((message) => + message.role === "user" && typeof message.content !== "string" + ? message.content.flatMap((part) => (part.type === "text" ? [part.text] : [])) + : [], + ); + expect(texts).toEqual(["task one finished", "task two finished"]); + }); + + it("includes a custom message sent without triggerTurn while the run starts in its first request", async () => { + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("before_agent_start", () => { + pi.sendMessage( + { customType: "note", content: "context from handler", display: true }, + { triggerTurn: false }, + ); + }); + }, + ], + }); + harnesses.push(harness); + const requests: TranscriptContext[] = []; + harness.setResponses([ + (context) => { + requests.push(context); + return fauxAssistantMessage("done"); + }, + ]); + + await harness.session.prompt("hello"); + + expect(requests).toHaveLength(1); + expect(JSON.stringify(requests[0].messages)).toContain("context from handler"); + }); + + it("aborts a run when abort is requested while before_agent_start runs", async () => { + let releaseHandler: () => void = () => {}; + const handlerStarted = Promise.withResolvers(); + const harness = await createHarness({ + extensionFactories: [ + (pi) => { + pi.on("before_agent_start", async () => { + handlerStarted.resolve(); + await new Promise((resolve) => { + releaseHandler = resolve; + }); + }); + }, + ], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("should not complete")]); + + const run = harness.session.sendCustomMessage( + { customType: "task-done", content: "background task finished", display: true }, + { triggerTurn: true }, + ); + await handlerStarted.promise; + const aborted = harness.session.abort(); + releaseHandler(); + await Promise.all([run, aborted]); + + const last = harness.session.messages.at(-1); + expect(last?.role).toBe("assistant"); + expect(last?.role === "assistant" && last.stopReason).toBe("aborted"); + expect(harness.session.isIdle).toBe(true); + }); + }); });