fix(extensions): restore activation rollback context

This commit is contained in:
Duncan Ogilvie
2026-08-10 01:54:36 +02:00
parent 0cf54f428c
commit 77916b6c99
4 changed files with 59 additions and 6 deletions
+7 -3
View File
@@ -2576,6 +2576,11 @@ async function createAgentSessionScoped(options: CreateAgentSessionOptions): Pro
autoApprove: options.autoApprove ?? false,
});
const toolContextStore = new ToolContextStore(getSessionContext);
const setSessionActiveToolNames = (names: Iterable<string>): void => {
const snapshot = Array.from(names);
setActiveToolNames(snapshot);
toolContextStore.setToolNames(snapshot);
};
// Native built-in implementations backing same-tool `ctx.invokeTool`, so a tool that
// re-registers a built-in (e.g. wrapping `write`) can delegate to the original — reaching the
// unwrapped native execute, which inherits the caller's already-granted approval rather than
@@ -2789,7 +2794,6 @@ async function createAgentSessionScoped(options: CreateAgentSessionOptions): Pro
toolNames: string[],
tools: Map<string, AgentTool>,
): Promise<BuildSystemPromptResult> => {
toolContextStore.setToolNames(toolNames);
const promptCwd = sessionManager.getCwd();
const activeRepoContext = hasSession
? await logger.time("resolveActiveRepoContext", resolveRepoContext, promptCwd)
@@ -3042,7 +3046,7 @@ async function createAgentSessionScoped(options: CreateAgentSessionOptions): Pro
if (mountedNames.length > 0 && !initialToolNames.includes("write")) initialToolNames.push("write");
}
setActiveToolNames(initialToolNames);
setSessionActiveToolNames(initialToolNames);
const { systemPrompt } = await logger.time(
"buildSystemPrompt",
rebuildSystemPrompt,
@@ -3395,7 +3399,7 @@ async function createAgentSessionScoped(options: CreateAgentSessionOptions): Pro
getXdevToolEntries: () => (toolSession.xdev ? xdevEntries(toolSession.xdev) : []),
xdev: toolSession.xdev,
presentationPinnedToolNames: explicitlyRequestedToolNameSet,
setActiveToolNames,
setActiveToolNames: setSessionActiveToolNames,
ensureWriteRegistered,
getMcpServerInstructions: mcpManager
? () => {
@@ -1191,12 +1191,14 @@ export class SessionTools {
const reconciled = await this.reconcileInspectImageTool();
const after = this.getEnabledToolNames().includes("inspect_image");
if (!reconciled || before === after) return;
const model = this.#host.model();
const modelName = model ? formatModelString(model) : "the current model";
this.#host.emitNotice(
"info",
after
? "inspect_image is now active for the selected model."
: "inspect_image is unavailable for the selected model.",
"inspect_image",
? `inspect_image is now available: ${modelName} has no native image input.`
: `inspect_image is now hidden: ${modelName} supports image input natively. Override with /vision on.`,
"vision",
);
});
}
@@ -40,6 +40,10 @@ describe("model switch from vision to text-only", () => {
rules: [],
contextFiles: [],
});
const notices: string[] = [];
const unsubscribe = session.subscribe(event => {
if (event.type === "notice") notices.push(`${event.source}:${event.message}`);
});
try {
await session.prompt("see image", { images: [{ type: "image", data: "aaaa", mimeType: "image/png" }] });
await session.setModel(text);
@@ -49,7 +53,12 @@ describe("model switch from vision to text-only", () => {
expect(
messages.flatMap<unknown>(message => (Array.isArray(message.content) ? message.content : [])),
).not.toContainEqual(expect.objectContaining({ type: "image" }));
await session.setModel(vision);
expect(notices.at(-1)).toContain("vision:inspect_image is now hidden:");
expect(notices.at(-1)).toContain("supports image input natively. Override with /vision on.");
} finally {
unsubscribe();
await session.dispose();
}
} finally {
@@ -1083,6 +1083,17 @@ describe("createAgentSession defaultInactive tool activation", () => {
const { session } = await createAgentSession({
...baseOptions(tempDir),
settings: Settings.isolated({
"bashInterceptor.enabled": true,
"bashInterceptor.patterns": [
{
pattern: "^\\s*printf\\s+",
tool: "detached_registration_tool",
message: "Use the detached registration tool.",
},
],
}),
autoApprove: true,
extensions: [detachedRegistrationExtension],
systemPrompt: defaultPrompt => {
if (rejectDetachedPrompt) throw new Error("expected detached registration failure");
@@ -1107,6 +1118,33 @@ describe("createAgentSession defaultInactive tool activation", () => {
error: "expected detached registration failure",
});
expect(session.getToolByName("detached_registration_tool")).toBeUndefined();
rejectDetachedPrompt = false;
const toolCallId = "detached-rollback-bash";
const mock = createMockModel({
responses: [
{
content: [
{
type: "toolCall",
id: toolCallId,
name: "bash",
arguments: { command: "printf rollback-ok" },
},
],
},
{ content: [{ type: "text", text: "done" }] },
],
});
vi.spyOn(session.agent, "streamFn").mockImplementation(mock.stream);
await withProviderAuth(["openai"], async () => {
await session.prompt("verify rollback context");
const bashResult = session.messages.find(
(message): message is ToolResultMessage =>
message.role === "toolResult" && message.toolCallId === toolCallId,
);
expect(bashResult?.isError).toBe(false);
expect(JSON.stringify(bashResult?.content)).toContain("rollback-ok");
});
} finally {
releaseDetachedRegistration.resolve();
await session.dispose();