diff --git a/packages/agent/src/agent.ts b/packages/agent/src/agent.ts index 54020435..452b8d76 100644 --- a/packages/agent/src/agent.ts +++ b/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; afterToolCall?: (context: AfterToolCallContext, signal?: AbortSignal) => Promise; prepareNextTurn?: ( + context: PrepareNextTurnContext, signal?: AbortSignal, ) => Promise | AgentLoopTurnUpdate | undefined; steeringMode?: QueueMode; @@ -184,6 +186,7 @@ export class Agent { signal?: AbortSignal, ) => Promise; public prepareNextTurn?: ( + context: PrepareNextTurnContext, signal?: AbortSignal, ) => Promise | 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, diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index c5304a0f..cd25768f 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/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 // ========================================================================= 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 new file mode 100644 index 00000000..74062b63 --- /dev/null +++ b/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(); + } + }); +});