Skip to content

Tool replacement/management robustness #11

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fu5ha wants to merge 1 commit into
base: main
Choose a base branch
from
+136 −41
Open

Tool replacement/management robustness #11

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 16 additions & 41 deletions src/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,6 @@ type PatchChunk = {
isEndOfFile: boolean;
};

type BaselineState = {
nonGptToolNames: string[];
};

export type FreeformToolFormat = {
type: "grammar";
syntax: "lark";
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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 {
Expand All @@ -1227,26 +1214,14 @@ function resolvePatchPath(cwd: string, filePath: string): string {
function syncToolset(
pi: Pick<ExtensionAPI, "getActiveTools" | "setActiveTools">,
model: Model<string> | 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 {
Expand Down Expand Up @@ -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);
});
}

Expand Down
120 changes: 120 additions & 0 deletions test/index.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,62 @@ const identityTheme = {
};
type ApplyPatchTool = ReturnType<typeof createApplyPatchTool>;
type ApplyPatchUpdate = Parameters<NonNullable<Parameters<ApplyPatchTool["execute"]>[3]>>[0];
type ToolsetHandler = (
event: { model?: { provider: string; id: string } },
ctx: { model: { provider: string; id: string } | undefined },
) => void | Promise<void>;

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<void>;
setActiveTools: (toolNames: string[]) => void;
getActiveTools: () => string[];
getSetActiveToolsCalls: () => string[][];
} {
let activeTools = [...initialActiveTools];
const setActiveToolsCalls: string[][] = [];
const handlers = new Map<string, ToolsetHandler[]>();
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<string> {
const directory = await mkdtemp(path.join(process.cwd(), "test-temp-"));
Expand Down Expand Up @@ -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();
Expand Down