diff --git a/src/index.ts b/src/index.ts index 27fd57f..eb4c663 100644 --- a/src/index.ts +++ b/src/index.ts @@ -31,10 +31,6 @@ type PatchChunk = { isEndOfFile: boolean; }; -type BaselineState = { - nonGptToolNames: string[]; -}; - export type FreeformToolFormat = { type: "grammar"; syntax: "lark"; @@ -252,7 +248,7 @@ function normalizeApplyPatchArguments(args: unknown): ApplyPatchParams { return { input: "" }; } -const EDIT_TOOL_NAMES = new Set(["write", "edit"]); +const STANDARD_EDIT_TOOL_NAMES = ["edit", "write"] as const; export const APPLY_PATCH_FREEFORM_DESCRIPTION = "Use the `apply_patch` tool to edit files. This is a FREEFORM tool, so do not wrap the patch in JSON."; export const APPLY_PATCH_LARK_GRAMMAR = `start: begin_patch hunk+ end_patch @@ -1196,28 +1192,19 @@ async function createPendingPatchUpdate( return { text: progress ? title : formatPendingPatchPaths(patchText), details: progress ? { progress } : undefined }; } -function hasEditTools(toolNames: string[]): boolean { - return toolNames.some((toolName) => EDIT_TOOL_NAMES.has(toolName)); -} - -function withoutApplyPatch(toolNames: string[]): string[] { - return toolNames.filter((toolName) => toolName !== "apply_patch"); +function withoutExtensionManagedEditTools(toolNames: string[]): string[] { + return toolNames.filter( + (toolName) => + toolName !== "apply_patch" && !STANDARD_EDIT_TOOL_NAMES.some((editToolName) => editToolName === toolName), + ); } function replaceEditToolsWithApplyPatch(toolNames: string[]): string[] { - const filteredToolNames = withoutApplyPatch(toolNames).filter((toolName) => !EDIT_TOOL_NAMES.has(toolName)); - if (!hasEditTools(toolNames)) { - return filteredToolNames; - } - return [...filteredToolNames, "apply_patch"]; + return [...withoutExtensionManagedEditTools(toolNames), "apply_patch"]; } -function restoreEditToolsFromBaseline(currentToolNames: string[], baselineToolNames: string[]): string[] { - const restoredToolNames = [ - ...withoutApplyPatch(currentToolNames), - ...baselineToolNames.filter((toolName) => EDIT_TOOL_NAMES.has(toolName)), - ]; - return [...new Set(restoredToolNames)]; +function replaceApplyPatchWithEditTools(toolNames: string[]): string[] { + return [...withoutExtensionManagedEditTools(toolNames), ...STANDARD_EDIT_TOOL_NAMES]; } function resolvePatchPath(cwd: string, filePath: string): string { @@ -1227,26 +1214,14 @@ function resolvePatchPath(cwd: string, filePath: string): string { function syncToolset( pi: Pick, model: Model | undefined, - state: BaselineState, ): void { const currentToolNames = pi.getActiveTools(); if (isOpenAIGptModel(model)) { - if (hasEditTools(currentToolNames)) { - state.nonGptToolNames = withoutApplyPatch(currentToolNames); - } pi.setActiveTools(replaceEditToolsWithApplyPatch(currentToolNames)); return; } - if (state.nonGptToolNames.length > 0) { - const restoredToolNames = restoreEditToolsFromBaseline(currentToolNames, state.nonGptToolNames); - state.nonGptToolNames = restoredToolNames; - pi.setActiveTools(restoredToolNames); - return; - } - - state.nonGptToolNames = withoutApplyPatch(currentToolNames); - pi.setActiveTools(state.nonGptToolNames); + pi.setActiveTools(replaceApplyPatchWithEditTools(currentToolNames)); } export function createApplyPatchTool(): ApplyPatchToolDefinition { @@ -1375,18 +1350,18 @@ export function createApplyPatchTool(): ApplyPatchToolDefinition { } export function registerApplyPatchExtension(pi: ApplyPatchExtensionAPI): void { - const state: BaselineState = { - nonGptToolNames: [], - }; - pi.registerTool(createApplyPatchTool()); pi.on("session_start", async (_event, ctx) => { - syncToolset(pi, ctx.model, state); + syncToolset(pi, ctx.model); }); pi.on("model_select", async (event) => { - syncToolset(pi, event.model, state); + syncToolset(pi, event.model); + }); + + pi.on("before_agent_start", async (_event, ctx) => { + syncToolset(pi, ctx.model); }); } diff --git a/test/index.test.ts b/test/index.test.ts index 1bffaed..eeaabd4 100644 --- a/test/index.test.ts +++ b/test/index.test.ts @@ -24,6 +24,62 @@ const identityTheme = { }; type ApplyPatchTool = ReturnType; type ApplyPatchUpdate = Parameters[3]>>[0]; +type ToolsetHandler = ( + event: { model?: { provider: string; id: string } }, + ctx: { model: { provider: string; id: string } | undefined }, +) => void | Promise; + +function isToolsetHandler(value: unknown): value is ToolsetHandler { + return typeof value === "function"; +} + +function createToolsetTestApi(initialActiveTools: string[]): { + api: ApplyPatchExtensionAPI; + trigger: (eventName: string, model: { provider: string; id: string } | undefined) => Promise; + setActiveTools: (toolNames: string[]) => void; + getActiveTools: () => string[]; + getSetActiveToolsCalls: () => string[][]; +} { + let activeTools = [...initialActiveTools]; + const setActiveToolsCalls: string[][] = []; + const handlers = new Map(); + const api: ApplyPatchExtensionAPI = { + registerTool() {}, + on(...args: unknown[]) { + const eventName = args[0]; + const handler = args[1]; + if (typeof eventName !== "string" || !isToolsetHandler(handler)) { + return; + } + handlers.set(eventName, [...(handlers.get(eventName) ?? []), handler]); + }, + getActiveTools() { + return [...activeTools]; + }, + setActiveTools(toolNames: string[]) { + activeTools = [...toolNames]; + setActiveToolsCalls.push([...toolNames]); + }, + }; + + return { + api, + async trigger(eventName, model) { + for (const handler of handlers.get(eventName) ?? []) { + await handler({ model }, { model }); + } + }, + setActiveTools(toolNames) { + activeTools = [...toolNames]; + }, + getActiveTools() { + return [...activeTools]; + }, + getSetActiveToolsCalls() { + return setActiveToolsCalls.map((toolNames) => [...toolNames]); + }, + }; +} async function createTempDirectory(): Promise { const directory = await mkdtemp(path.join(process.cwd(), "test-temp-")); @@ -72,6 +128,70 @@ describe("pi-apply-patch", () => { }); }); + it("#given GPT model after reload with apply_patch already active #when session starts #then keeps apply_patch active", async () => { + // given + const harness = createToolsetTestApi(["read", "bash", "apply_patch"]); + registerApplyPatchExtension(harness.api); + + // when + await harness.trigger("session_start", { provider: "openai", id: "gpt-5" }); + + // then + expect(harness.getActiveTools()).toEqual(["read", "bash", "apply_patch"]); + expect(harness.getSetActiveToolsCalls()).toEqual([["read", "bash", "apply_patch"]]); + }); + + it("#given GPT model with stale edit tools #when session starts #then normalizes to apply_patch only", async () => { + // given + const harness = createToolsetTestApi(["read", "apply_patch", "edit", "write"]); + registerApplyPatchExtension(harness.api); + + // when + await harness.trigger("session_start", { provider: "openai", id: "gpt-5" }); + + // then + expect(harness.getActiveTools()).toEqual(["read", "apply_patch"]); + }); + + it("#given non GPT model and no original edit tools #when session starts #then restores standard edit tools", async () => { + // given + const harness = createToolsetTestApi(["read", "apply_patch"]); + registerApplyPatchExtension(harness.api); + + // when + await harness.trigger("session_start", { provider: "anthropic", id: "claude-sonnet-4" }); + + // then + expect(harness.getActiveTools()).toEqual(["read", "edit", "write"]); + }); + + it("#given external tool change in GPT mode #when agent starts #then reconciles before model request", async () => { + // given + const harness = createToolsetTestApi(["read", "edit", "write"]); + registerApplyPatchExtension(harness.api); + await harness.trigger("session_start", { provider: "openai", id: "gpt-5" }); + harness.setActiveTools(["read", "write", "apply_patch", "edit"]); + + // when + await harness.trigger("before_agent_start", { provider: "openai", id: "gpt-5" }); + + // then + expect(harness.getActiveTools()).toEqual(["read", "apply_patch"]); + }); + + it("#given GPT mode #when model switches to non GPT #then apply_patch is replaced with edit tools", async () => { + // given + const harness = createToolsetTestApi(["read", "edit", "write"]); + registerApplyPatchExtension(harness.api); + await harness.trigger("session_start", { provider: "openai", id: "gpt-5" }); + + // when + await harness.trigger("model_select", { provider: "anthropic", id: "claude-sonnet-4" }); + + // then + expect(harness.getActiveTools()).toEqual(["read", "edit", "write"]); + }); + it("#given raw codex patch #when executed #then applies file update", async () => { // given const directory = await createTempDirectory();