fix(coding-agent): refresh session state before next turn

packages/agent/src/agent.ts

@@ -21,6 +21,7 @@ import type {
 	AgentTool,
 	BeforeToolCallContext,
 	BeforeToolCallResult,
+	PrepareNextTurnContext,
 	QueueMode,
 	StreamFn,
 	ToolExecutionMode,
@@ -104,6 +105,7 @@ export interface AgentOptions {
 	beforeToolCall?: (context: BeforeToolCallContext, signal?: AbortSignal) => Promise<BeforeToolCallResult | undefined>;
 	afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise<AfterToolCallResult | undefined>;
 	prepareNextTurn?: (
+		context: PrepareNextTurnContext,
 		signal?: AbortSignal,
 	) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;
 	steeringMode?: QueueMode;
@@ -184,6 +186,7 @@ export class Agent {
 		signal?: AbortSignal,
 	) => Promise<AfterToolCallResult | undefined>;
 	public prepareNextTurn?: (
+		context: PrepareNextTurnContext,
 		signal?: AbortSignal,
 	) => Promise<AgentLoopTurnUpdate | undefined> | AgentLoopTurnUpdate | undefined;
 	private activeRun?: ActiveRun;
@@ -433,7 +436,9 @@ export class Agent {
 			toolExecution: this.toolExecution,
 			beforeToolCall: this.beforeToolCall,
 			afterToolCall: this.afterToolCall,
-			prepareNextTurn: this.prepareNextTurn ? async () => await this.prepareNextTurn?.(this.signal) : undefined,
+			prepareNextTurn: this.prepareNextTurn
+				? async (context) => await this.prepareNextTurn?.(context, this.signal)
+				: undefined,
 			convertToLlm: this.convertToLlm,
 			transformContext: this.transformContext,
 			getApiKey: this.getApiKey,

packages/coding-agent/src/core/agent-session.ts

@@ -352,6 +352,7 @@ export class AgentSession {
 		// (session persistence, extensions, auto-compaction, retry logic)
 		this._unsubscribeAgent = this.agent.subscribe(this._handleAgentEvent);
 		this._installAgentToolHooks();
+		this._installAgentNextTurnRefresh();
 
 		this._buildRuntime({
 			activeToolNames: this._initialActiveToolNames,
@@ -462,6 +463,25 @@ export class AgentSession {
 		};
 	}
 
+	private _installAgentNextTurnRefresh(): void {
+		const previousPrepareNextTurn = this.agent.prepareNextTurn;
+		this.agent.prepareNextTurn = async (turn, signal) => {
+			const previousSnapshot = await previousPrepareNextTurn?.(turn, signal);
+			const previousContext = previousSnapshot?.context ?? turn.context;
+
+			return {
+				...previousSnapshot,
+				context: {
+					...previousContext,
+					systemPrompt: this.agent.state.systemPrompt,
+					tools: this.agent.state.tools.slice(),
+				},
+				model: this.agent.state.model,
+				thinkingLevel: this.agent.state.thinkingLevel,
+			};
+		};
+	}
+
 	// =========================================================================
 	// Event Subscription
 	// =========================================================================

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

@@ -0,0 +1,68 @@
+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();
+		}
+	});
+});