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

closes #6162

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