From fd6659dd5d32d67feaa7ce2ba5eeb87c5705149c Mon Sep 17 00:00:00 2001 From: Mario Zechner Date: Tue, 30 Jun 2026 14:08:48 +0200 Subject: [PATCH] fix(coding-agent): preserve run prompt during tool refresh closes #6162 --- packages/agent/CHANGELOG.md | 8 ++ packages/agent/src/agent.ts | 19 ++- packages/agent/test/agent.test.ts | 41 ++++++ packages/coding-agent/CHANGELOG.md | 1 + .../coding-agent/src/core/agent-session.ts | 21 ++- ...2-extension-active-tools-next-turn.test.ts | 136 ++++++++++++++++++ .../extension-active-tools-next-turn.test.ts | 68 --------- 7 files changed, 217 insertions(+), 77 deletions(-) create mode 100644 packages/coding-agent/test/suite/regressions/6162-extension-active-tools-next-turn.test.ts delete mode 100644 packages/coding-agent/test/suite/regressions/extension-active-tools-next-turn.test.ts diff --git a/packages/agent/CHANGELOG.md b/packages/agent/CHANGELOG.md index bec415948..7b1f8e001 100644 --- a/packages/agent/CHANGELOG.md +++ b/packages/agent/CHANGELOG.md @@ -2,6 +2,14 @@ ## [Unreleased] +### Added + +- Added `prepareNextTurnWithContext` for `Agent` users that need the next-turn loop context. + +### Fixed + +- Fixed `Agent.prepareNextTurn` to keep receiving the run abort signal instead of the next-turn context. + ## [0.80.2] - 2026-06-23 ### Changed diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 452b8d764..ab10d5ac1 100644 --- a/packages/agent/src/agent.ts +++ b/packages/agent/src/agent.ts @@ -105,6 +105,9 @@ export interface AgentOptions { beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise; afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise; prepareNextTurn?: ( + signal?: AbortSignal, + ) => Promise | AgentLoopTurnUpdate | undefined; + prepareNextTurnWithContext?: ( context: PrepareNextTurnContext, signal?: AbortSignal, ) => Promise | AgentLoopTurnUpdate | undefined; @@ -186,6 +189,9 @@ export class Agent { signal?: AbortSignal, ) => Promise; public prepareNextTurn?: ( + signal?: AbortSignal, + ) => Promise | AgentLoopTurnUpdate | undefined; + public prepareNextTurnWithContext?: ( context: PrepareNextTurnContext, signal?: AbortSignal, ) => Promise | AgentLoopTurnUpdate | undefined; @@ -212,6 +218,7 @@ export class Agent { this.beforeToolCall = options.beforeToolCall; this.afterToolCall = options.afterToolCall; this.prepareNextTurn = options.prepareNextTurn; + this.prepareNextTurnWithContext = options.prepareNextTurnWithContext; this.steeringQueue = new PendingMessageQueue(options.steeringMode ?? "one-at-a-time"); this.followUpQueue = new PendingMessageQueue(options.followUpMode ?? "one-at-a-time"); this.sessionId = options.sessionId; @@ -436,9 +443,15 @@ export class Agent { toolExecution: this.toolExecution, beforeToolCall: this.beforeToolCall, afterToolCall: this.afterToolCall, - prepareNextTurn: this.prepareNextTurn - ? async (context) => await this.prepareNextTurn?.(context, this.signal) - : undefined, + prepareNextTurn: + this.prepareNextTurnWithContext || this.prepareNextTurn + ? async (context) => { + if (this.prepareNextTurnWithContext) { + return await this.prepareNextTurnWithContext(context, this.signal); + } + return await this.prepareNextTurn?.(this.signal); + } + : undefined, convertToLlm: this.convertToLlm, transformContext: this.transformContext, getApiKey: this.getApiKey, diff --git a/packages/agent/test/agent.test.ts b/packages/agent/test/agent.test.ts index 5fa27c5aa..f677ee788 100644 --- a/packages/agent/test/agent.test.ts +++ b/packages/agent/test/agent.test.ts @@ -630,6 +630,47 @@ describe("Agent", () => { expect(responseCount).toBe(2); }); + it("keeps legacy prepareNextTurn signal callback behavior", async () => { + const schema = Type.Object({}); + const tool: AgentTool = { + name: "noop", + label: "Noop", + description: "Noop tool", + parameters: schema, + execute: async () => ({ content: [{ type: "text", text: "ok" }], details: {} }), + }; + let requestCount = 0; + let sawAbortSignal = false; + const agent = new Agent({ + initialState: { tools: [tool] }, + prepareNextTurn: async (signal) => { + sawAbortSignal = signal instanceof AbortSignal; + return undefined; + }, + streamFn: () => { + requestCount++; + const stream = new MockAssistantStream(); + queueMicrotask(() => { + if (requestCount === 1) { + const message = createAssistantToolUseMessage([ + { type: "toolCall", id: "tool-1", name: "noop", arguments: {} }, + ]); + stream.push({ type: "done", reason: "toolUse", message }); + return; + } + const message = createAssistantMessage("done"); + stream.push({ type: "done", reason: "stop", message }); + }); + return stream; + }, + }); + + await agent.prompt("start"); + + expect(requestCount).toBe(2); + expect(sawAbortSignal).toBe(true); + }); + it("forwards sessionId to streamFn options", async () => { let receivedSessionId: string | undefined; const agent = new Agent({ diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index b69f609d7..5ea14c7a4 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -9,6 +9,7 @@ ### Fixed +- Fixed extension tool changes to apply before the next provider request in the same agent run without dropping `before_agent_start` system-prompt overrides ([#6162](https://github.com/earendil-works/pi/issues/6162)). - Fixed a crash when undici emits an internal client error while terminating a mid-stream HTTP response ([#6133](https://github.com/earendil-works/pi/issues/6133)). - Fixed the compaction event regression test to cover status indicator cleanup and keep CI passing. - Fixed interactive status indicators so ending work, retry, compaction, or branch-summary indicators no longer shrink the TUI when clear-on-shrink is enabled ([#6026](https://github.com/earendil-works/pi/pull/6026)). diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index cd25768f1..ff99d6949 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -21,6 +21,7 @@ import type { AgentMessage, AgentState, AgentTool, + PrepareNextTurnContext, ThinkingLevel, } from "@earendil-works/pi-agent-core"; import type { AssistantMessage, ImageContent, Message, Model, TextContent } from "@earendil-works/pi-ai/compat"; @@ -331,6 +332,7 @@ export class AgentSession { // Base system prompt (without extension appends) - used to apply fresh appends each turn private _baseSystemPrompt = ""; private _baseSystemPromptOptions!: BuildSystemPromptOptions; + private _systemPromptOverride?: string; constructor(config: AgentSessionConfig) { this.agent = config.agent; @@ -464,16 +466,20 @@ export class AgentSession { } private _installAgentNextTurnRefresh(): void { - const previousPrepareNextTurn = this.agent.prepareNextTurn; - this.agent.prepareNextTurn = async (turn, signal) => { - const previousSnapshot = await previousPrepareNextTurn?.(turn, signal); + const previousPrepareNextTurnWithContext = + this.agent.prepareNextTurnWithContext ?? + (this.agent.prepareNextTurn + ? async (_turn: PrepareNextTurnContext, signal?: AbortSignal) => await this.agent.prepareNextTurn?.(signal) + : undefined); + this.agent.prepareNextTurnWithContext = async (turn, signal) => { + const previousSnapshot = await previousPrepareNextTurnWithContext?.(turn, signal); const previousContext = previousSnapshot?.context ?? turn.context; return { ...previousSnapshot, context: { ...previousContext, - systemPrompt: this.agent.state.systemPrompt, + systemPrompt: this._systemPromptOverride ?? this._baseSystemPrompt, tools: this.agent.state.tools.slice(), }, model: this.agent.state.model, @@ -844,7 +850,7 @@ export class AgentSession { // Rebuild base system prompt with new tool set this._baseSystemPrompt = this._rebuildSystemPrompt(validToolNames); - this.agent.state.systemPrompt = this._baseSystemPrompt; + this.agent.state.systemPrompt = this._systemPromptOverride ?? this._baseSystemPrompt; } /** Whether compaction or branch summarization is currently running */ @@ -972,6 +978,7 @@ export class AgentSession { await this.agent.continue(); } } finally { + this._systemPromptOverride = undefined; this._flushPendingBashMessages(); } } @@ -1143,10 +1150,12 @@ export class AgentSession { } } // Apply extension-modified system prompt, or reset to base - if (result?.systemPrompt) { + if (result?.systemPrompt !== undefined) { + this._systemPromptOverride = result.systemPrompt; this.agent.state.systemPrompt = result.systemPrompt; } else { // Ensure we're using the base prompt (in case previous turn had modifications) + this._systemPromptOverride = undefined; this.agent.state.systemPrompt = this._baseSystemPrompt; } } catch (error) { diff --git a/packages/coding-agent/test/suite/regressions/6162-extension-active-tools-next-turn.test.ts b/packages/coding-agent/test/suite/regressions/6162-extension-active-tools-next-turn.test.ts new file mode 100644 index 000000000..7abe4a5e6 --- /dev/null +++ b/packages/coding-agent/test/suite/regressions/6162-extension-active-tools-next-turn.test.ts @@ -0,0 +1,136 @@ +import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; +import { Type } from "typebox"; +import { describe, expect, it } from "vitest"; +import type { ExtensionFactory } from "../../../src/index.ts"; +import { createHarness } from "../harness.ts"; + +describe("extension active tools next-turn refresh", () => { + it("applies pi.setActiveTools before the next provider request in the same run", async () => { + const extensionFactories: ExtensionFactory[] = [ + (pi) => { + pi.registerTool({ + name: "switch_tools", + label: "Switch Tools", + description: "Switch the active extension tool set", + promptSnippet: "Switch to the next extension tool", + parameters: Type.Object({}), + execute: async () => { + pi.setActiveTools(["after_switch"]); + return { + content: [{ type: "text", text: "switched" }], + details: {}, + }; + }, + }); + + pi.registerTool({ + name: "after_switch", + label: "After Switch", + description: "Tool that should be available after switching", + promptSnippet: "Run after the active tool set changes", + parameters: Type.Object({}), + execute: async () => ({ + content: [{ type: "text", text: "after" }], + details: {}, + }), + }); + }, + ]; + const harness = await createHarness({ + extensionFactories, + }); + + try { + harness.session.setActiveToolsByName(["switch_tools"]); + + const providerToolNames: string[][] = []; + harness.setResponses([ + (context) => { + providerToolNames.push((context.tools ?? []).map((tool) => tool.name).sort()); + return fauxAssistantMessage(fauxToolCall("switch_tools", {}), { stopReason: "toolUse" }); + }, + (context) => { + providerToolNames.push((context.tools ?? []).map((tool) => tool.name).sort()); + return fauxAssistantMessage("done"); + }, + ]); + + expect(harness.session.getActiveToolNames()).toEqual(["switch_tools"]); + + await harness.session.prompt("start"); + + expect(harness.session.getActiveToolNames()).toEqual(["after_switch"]); + expect(providerToolNames).toEqual([["switch_tools"], ["after_switch"]]); + } finally { + harness.cleanup(); + } + }); + + it("preserves before_agent_start system prompt overrides when tools change mid-run", async () => { + const extensionFactories: ExtensionFactory[] = [ + (pi) => { + pi.on("before_agent_start", async (event) => ({ + systemPrompt: `${event.systemPrompt}\n\nkeep this run override`, + })); + + pi.registerTool({ + name: "switch_tools", + label: "Switch Tools", + description: "Switch the active extension tool set", + promptSnippet: "Switch to the next extension tool", + parameters: Type.Object({}), + execute: async () => { + pi.setActiveTools(["after_switch"]); + return { + content: [{ type: "text", text: "switched" }], + details: {}, + }; + }, + }); + + pi.registerTool({ + name: "after_switch", + label: "After Switch", + description: "Tool that should be available after switching", + promptSnippet: "Run after the active tool set changes", + parameters: Type.Object({}), + execute: async () => ({ + content: [{ type: "text", text: "after" }], + details: {}, + }), + }); + }, + ]; + const harness = await createHarness({ + extensionFactories, + }); + + try { + harness.session.setActiveToolsByName(["switch_tools"]); + + const providerSystemPrompts: string[] = []; + const providerToolNames: string[][] = []; + harness.setResponses([ + (context) => { + providerSystemPrompts.push(context.systemPrompt ?? ""); + providerToolNames.push((context.tools ?? []).map((tool) => tool.name).sort()); + return fauxAssistantMessage(fauxToolCall("switch_tools", {}), { stopReason: "toolUse" }); + }, + (context) => { + providerSystemPrompts.push(context.systemPrompt ?? ""); + providerToolNames.push((context.tools ?? []).map((tool) => tool.name).sort()); + return fauxAssistantMessage("done"); + }, + ]); + + await harness.session.prompt("start"); + + expect(providerToolNames).toEqual([["switch_tools"], ["after_switch"]]); + expect(providerSystemPrompts).toHaveLength(2); + expect(providerSystemPrompts[0]).toContain("keep this run override"); + expect(providerSystemPrompts[1]).toContain("keep this run override"); + } finally { + harness.cleanup(); + } + }); +}); diff --git a/packages/coding-agent/test/suite/regressions/extension-active-tools-next-turn.test.ts b/packages/coding-agent/test/suite/regressions/extension-active-tools-next-turn.test.ts deleted file mode 100644 index 74062b639..000000000 --- a/packages/coding-agent/test/suite/regressions/extension-active-tools-next-turn.test.ts +++ /dev/null @@ -1,68 +0,0 @@ -import { fauxAssistantMessage, fauxToolCall } from "@earendil-works/pi-ai"; -import { Type } from "typebox"; -import { describe, expect, it } from "vitest"; -import type { ExtensionFactory } from "../../../src/index.ts"; -import { createHarness } from "../harness.ts"; - -describe("extension active tools next-turn refresh", () => { - it("applies pi.setActiveTools before the next provider request in the same run", async () => { - const extensionFactories: ExtensionFactory[] = [ - (pi) => { - pi.registerTool({ - name: "switch_tools", - label: "Switch Tools", - description: "Switch the active extension tool set", - promptSnippet: "Switch to the next extension tool", - parameters: Type.Object({}), - execute: async () => { - pi.setActiveTools(["after_switch"]); - return { - content: [{ type: "text", text: "switched" }], - details: {}, - }; - }, - }); - - pi.registerTool({ - name: "after_switch", - label: "After Switch", - description: "Tool that should be available after switching", - promptSnippet: "Run after the active tool set changes", - parameters: Type.Object({}), - execute: async () => ({ - content: [{ type: "text", text: "after" }], - details: {}, - }), - }); - }, - ]; - const harness = await createHarness({ - extensionFactories, - }); - - try { - harness.session.setActiveToolsByName(["switch_tools"]); - - const providerToolNames: string[][] = []; - harness.setResponses([ - (context) => { - providerToolNames.push((context.tools ?? []).map((tool) => tool.name).sort()); - return fauxAssistantMessage(fauxToolCall("switch_tools", {}), { stopReason: "toolUse" }); - }, - (context) => { - providerToolNames.push((context.tools ?? []).map((tool) => tool.name).sort()); - return fauxAssistantMessage("done"); - }, - ]); - - expect(harness.session.getActiveToolNames()).toEqual(["switch_tools"]); - - await harness.session.prompt("start"); - - expect(harness.session.getActiveToolNames()).toEqual(["after_switch"]); - expect(providerToolNames).toEqual([["switch_tools"], ["after_switch"]]); - } finally { - harness.cleanup(); - } - }); -});