fix(coding-agent): preserve run prompt during tool refresh

closes #6162

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

packages/agent/src/agent.ts

@@ -105,6 +105,9 @@ export interface AgentOptions {
 	beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
 	afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
 	prepareNextTurn?: (
+		signal?: AbortSignal,
+	) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;
+	prepareNextTurnWithContext?: (
 		context: PrepareNextTurnContext,
 		signal?: AbortSignal,
 	) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;
@@ -186,6 +189,9 @@ export class Agent {
 		signal?: AbortSignal,
 	) => Promise<AfterToolCallResult | undefined>;
 	public prepareNextTurn?: (
+		signal?: AbortSignal,
+	) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;
+	public prepareNextTurnWithContext?: (
 		context: PrepareNextTurnContext,
 		signal?: AbortSignal,
 	) => Promise<AgentLoopTurnUpdate | undefined> | 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,

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<typeof schema> = {
+			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({

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)).

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) {

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

packages/coding-agent/test/suite/regressions/extension-active-tools-next-turn.test.ts

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