Add Pi extensions and same-run add_tool to CUA CLI

packages/agent/package.json

@@ -40,8 +40,8 @@
     "typecheck": "tsc -b"
   },
   "dependencies": {
-    "@earendil-works/pi-agent-core": "0.80.6",
-    "@earendil-works/pi-ai": "0.80.6",
+    "@earendil-works/pi-agent-core": "https://raw.githubusercontent.com/kernel/cua/eac28f9f4992c648d0d046dfe62550bce643ff85/vendor/pi/earendil-works-pi-agent-core-0.80.7-anchor.3d8f743.tgz",
+    "@earendil-works/pi-ai": "https://raw.githubusercontent.com/kernel/cua/eac28f9f4992c648d0d046dfe62550bce643ff85/vendor/pi/earendil-works-pi-ai-0.80.7-anchor.3d8f743.tgz",
     "@onkernel/cua-ai": "0.6.0",
     "@onkernel/sdk": "0.49.0",
     "sharp": "^0.34.5"

packages/agent/src/agent.ts

@@ -650,6 +650,7 @@ export class CuaAgentHarness<
 	private requestedActiveToolNames?: string[];
 	private emptyResponseRecoveryAttempts = 0;
 	private hasPendingActiveQueue = false;
+	private toolMutation = Promise.resolve();
 
 	constructor(options: CuaAgentHarnessOptions<TSkill, TPromptTemplate>) {
 		const {
@@ -734,6 +735,16 @@ export class CuaAgentHarness<
 		this.emptyResponseRecoveryAttempts += 1;
 	}
 
+	/** Serialize runtime and extension tool-list changes. */
+	async mutateTools<T>(mutation: () => Promise<T>): Promise<T> {
+		const result = this.toolMutation.then(mutation);
+		this.toolMutation = result.then(
+			() => {},
+			() => {},
+		);
+		return result;
+	}
+
 	/**
 	 * Mirror pi `AgentHarness.setModel()` while accepting CUA model refs.
 	 *
@@ -744,16 +755,20 @@ export class CuaAgentHarness<
 	override async setModel(model: CuaRuntimeInput): Promise<void> {
 		this.runtime.setModel(model);
 		const tools = this.runtime.tools();
-		try {
-			await super.setTools(tools, this.requestedActiveToolNames ?? tools.map((tool) => tool.name));
-		} catch (err) {
-			// The pre-switch tools stay exposed, so restore the runtime they are
-			// bound to — including its still-live translator.
-			this.runtime.rollbackSwitch();
-			throw err;
-		}
-		this.runtime.commitSwitch();
-		await super.setModel(this.runtime.model);
+		const resolvedModel = this.runtime.model;
+		await this.mutateTools(async () => {
+			try {
+				await super.setTools(
+					tools,
+					this.requestedActiveToolNames ?? tools.map((tool) => tool.name),
+				);
+			} catch (err) {
+				this.runtime.rollbackSwitch();
+				throw err;
+			}
+			this.runtime.commitSwitch();
+		});
+		await super.setModel(resolvedModel);
 	}
 
 	override async setActiveTools(toolNames: string[]): Promise<void> {
@@ -767,28 +782,39 @@ export class CuaAgentHarness<
 	 * plane conflicts with the requested mode.
 	 */
 	async setMode(mode: CuaMode): Promise<void> {
-		if (mode === this.runtime.mode) return;
-		const previousNames = new Set(this.getTools().map((tool) => tool.name));
-		this.runtime.setMode(mode);
-		const tools = this.runtime.tools();
-		// Tools that survive the mode switch (extraTools, shared names) keep
-		// their requested activation state; names new in this mode activate.
-		const requested = this.requestedActiveToolNames;
-		const active = requested
-			? tools.map((tool) => tool.name).filter((name) => !previousNames.has(name) || requested.includes(name))
-			: tools.map((tool) => tool.name);
-		try {
-			await super.setTools(tools, active);
-		} catch (err) {
-			// The pre-switch tools stay exposed, so restore the runtime they are
-			// bound to — including its still-live translator.
-			this.runtime.rollbackSwitch();
-			throw err;
-		}
-		this.runtime.commitSwitch();
-		// The requested subset now reflects this mode's toolset; without this a
-		// later setModel would restore the pre-switch names.
-		if (requested) this.requestedActiveToolNames = active;
+		await this.mutateTools(async () => {
+			if (mode === this.runtime.mode) return;
+			const previousTools = this.getTools();
+			const previousNames = new Set(previousTools.map((tool) => tool.name));
+			const previousRuntimeNames = new Set(this.runtime.tools().map((tool) => tool.name));
+			const previousActiveNames = new Set(this.getActiveTools().map((tool) => tool.name));
+			this.runtime.setMode(mode);
+			const runtimeTools = this.runtime.tools();
+			const runtimeNames = new Set(runtimeTools.map((tool) => tool.name));
+			const tools = [
+				...runtimeTools,
+				...previousTools.filter(
+					(tool) => !previousRuntimeNames.has(tool.name) && !runtimeNames.has(tool.name),
+				),
+			];
+			// Tools that survive the mode switch keep their active state; names new
+			// in this mode activate.
+			const active = tools
+				.map((tool) => tool.name)
+				.filter((name) => !previousNames.has(name) || previousActiveNames.has(name));
+			try {
+				await super.setTools(tools, active);
+			} catch (err) {
+				// The pre-switch tools stay exposed, so restore the runtime they are
+				// bound to — including its still-live translator.
+				this.runtime.rollbackSwitch();
+				throw err;
+			}
+			this.runtime.commitSwitch();
+			// The requested subset now reflects this mode's toolset; without this a
+			// later setModel would restore the pre-switch names.
+			if (this.requestedActiveToolNames) this.requestedActiveToolNames = active;
+		});
 	}
 
 	/** The action plane(s) currently exposed to the model. */

packages/agent/test/agent.test.ts

@@ -1820,6 +1820,23 @@ describe("CuaAgentHarness", () => {
 		expect(names).toContain("browser_snapshot");
 	});
 
+	it("setMode keeps tools registered after construction", async () => {
+		const harness = new CuaAgentHarness({
+			...(await createHarnessServices()),
+			browser,
+			client,
+			model: "anthropic:claude-opus-4-5",
+		});
+		const extensionTool = createCustomTool("extension_custom");
+		const tools = [...harness.getTools(), extensionTool];
+		await harness.setTools(tools, tools.map((tool) => tool.name));
+
+		await harness.setMode("browser");
+
+		expect(harness.getTools().map((tool) => tool.name)).toContain("extension_custom");
+		expect(harness.getActiveTools().map((tool) => tool.name)).toContain("extension_custom");
+	});
+
 	it("setMode keeps the requested activation state of surviving tools", async () => {
 		const harness = new CuaAgentHarness({
 			...(await createHarnessServices()),

packages/ai/package.json

@@ -41,7 +41,7 @@
     "test:integration": "vitest --run --config vitest.integration.config.ts"
   },
   "dependencies": {
-    "@earendil-works/pi-ai": "0.80.6",
+    "@earendil-works/pi-ai": "https://raw.githubusercontent.com/kernel/cua/eac28f9f4992c648d0d046dfe62550bce643ff85/vendor/pi/earendil-works-pi-ai-0.80.7-anchor.3d8f743.tgz",
     "@tzafon/lightcone": "^0.7.0",
     "openai": "^6.26.0"
   },

packages/ai/test/dynamic-tool-payload.test.ts

@@ -0,0 +1,215 @@
+import { describe, expect, it, vi } from "vitest";
+import type { AgentTool, Context, Message, Model } from "@earendil-works/pi-ai";
+import { createServer } from "node:http";
+import { cuaModels, getCuaModel, resolveCuaRuntimeSpec } from "../src/index";
+import { ANTHROPIC_NATIVE_BROWSER_MESSAGES_API } from "../src/providers/anthropic/native";
+
+const { openaiCreate } = vi.hoisted(() => ({
+	openaiCreate: vi.fn(),
+}));
+
+vi.mock("openai", () => ({
+	default: class {
+		responses = {
+			create: (...args: unknown[]) => {
+				openaiCreate(...args);
+				return {
+					withResponse: async () => {
+						throw new Error("stop after payload capture");
+					},
+				};
+			},
+		};
+	},
+}));
+
+const dynamicTool: AgentTool = {
+	name: "learned_tool",
+	label: "Learned tool",
+	description: "learned",
+	parameters: { type: "object", properties: {} },
+	execute: async () => ({ content: [], details: {} }),
+};
+const addTool: AgentTool = { ...dynamicTool, name: "add_tool", label: "Add tool" };
+
+function addMessages(api: string, responseId?: string): Message[] {
+	return [
+		{
+			role: "assistant",
+			content: [{ type: "toolCall", id: "add_1", name: "add_tool", arguments: {} }],
+			api,
+			provider: api.startsWith("openai") ? "openai" : "anthropic",
+			model: api.startsWith("openai") ? "gpt-5.5" : "claude-sonnet-4-5",
+			responseId,
+			usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
+			stopReason: "toolUse",
+			timestamp: 0,
+		},
+		{
+			role: "toolResult",
+			toolCallId: "add_1",
+			toolName: "add_tool",
+			content: [{ type: "text", text: "added" }],
+			addedToolNames: ["learned_tool"],
+			isError: false,
+			timestamp: 0,
+		},
+	];
+}
+
+async function capturePayload(
+	model: Model<any>,
+	context: Context,
+	transform?: (payload: unknown, model: Model<any>) => unknown | Promise<unknown>,
+): Promise<Record<string, any>> {
+	let payload: Record<string, any> | undefined;
+	await cuaModels()
+		.streamSimple(model, context, {
+			apiKey: "test",
+			onPayload: async (value, payloadModel) => {
+				const transformed = transform
+					? ((await transform(value, payloadModel)) ?? value)
+					: value;
+				payload = transformed as Record<string, any>;
+				return transformed;
+			},
+		})
+		.result();
+	if (!payload) throw new Error("payload was not captured");
+	return payload;
+}
+
+describe("message-anchored dynamic tool payloads", () => {
+	it("uses OpenAI tool search while preserving response threading", async () => {
+		const model = getCuaModel("openai:gpt-5.5");
+		const payload = await capturePayload(model, {
+			systemPrompt: "test",
+			tools: [addTool, dynamicTool],
+			messages: addMessages(model.api, "resp_add"),
+		});
+		expect(payload.store).toBe(true);
+		expect(payload.previous_response_id).toBe("resp_add");
+		const searchOutput = payload.input.find((item: any) => item.type === "tool_search_output");
+		expect(searchOutput.tools).toEqual(
+			expect.arrayContaining([expect.objectContaining({ name: "learned_tool", defer_loading: true })]),
+		);
+		expect(payload.input).toEqual(
+			expect.arrayContaining([
+				expect.objectContaining({ type: "tool_search_call" }),
+				expect.objectContaining({ type: "tool_search_output" }),
+			]),
+		);
+	});
+
+	it("eagerly declares the tool after the add marker is pruned by the next response anchor", async () => {
+		const model = getCuaModel("openai:gpt-5.5");
+		const messages = addMessages(model.api, "resp_add");
+		messages.push(
+			{
+				role: "assistant",
+				content: [{ type: "toolCall", id: "learn_1", name: "learned_tool", arguments: {} }],
+				api: model.api,
+				provider: "openai",
+				model: model.id,
+				responseId: "resp_learned",
+				usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
+				stopReason: "toolUse",
+				timestamp: 0,
+			},
+			{ role: "toolResult", toolCallId: "learn_1", toolName: "learned_tool", content: [], isError: false, timestamp: 0 },
+		);
+		const payload = await capturePayload(model, { systemPrompt: "test", tools: [addTool, dynamicTool], messages });
+		expect(payload.previous_response_id).toBe("resp_learned");
+		expect(payload.tools).toEqual(
+			expect.arrayContaining([expect.objectContaining({ name: "learned_tool" })]),
+		);
+		expect(payload.tools.find((tool: any) => tool.name === "learned_tool").defer_loading).toBeUndefined();
+	});
+
+	it("uses Anthropic references for supported models and eager tools for Haiku", async () => {
+		const supported = getCuaModel("anthropic:claude-sonnet-4-5");
+		const supportedPayload = await capturePayload(supported, {
+			systemPrompt: "test",
+			tools: [addTool, dynamicTool],
+			messages: addMessages(supported.api),
+		});
+		expect(supportedPayload.tools).toEqual(
+			expect.arrayContaining([expect.objectContaining({ name: "learned_tool", defer_loading: true })]),
+		);
+		expect(JSON.stringify(supportedPayload.messages)).toContain("tool_reference");
+
+		const haiku = {
+			...supported,
+			id: "claude-haiku-4",
+			compat: { ...supported.compat, supportsToolReferences: false },
+		};
+		const haikuPayload = await capturePayload(haiku, {
+			systemPrompt: "test",
+			tools: [addTool, dynamicTool],
+			messages: addMessages(haiku.api),
+		});
+		expect(JSON.stringify(haikuPayload.messages)).not.toContain("tool_reference");
+		expect(haikuPayload.tools).toEqual(
+			expect.arrayContaining([expect.objectContaining({ name: "learned_tool" })]),
+		);
+	});
+
+	it("preserves dynamic tools and beta headers through Anthropic native routing", async () => {
+		let requestHeaders: Record<string, string | string[] | undefined> = {};
+		const server = createServer((request, response) => {
+			requestHeaders = request.headers;
+			response.writeHead(500, { "content-type": "application/json" });
+			response.end('{"error":{"message":"captured"}}');
+		});
+		await new Promise<void>((resolve) =>
+			server.listen(0, "127.0.0.1", resolve),
+		);
+		try {
+			const address = server.address();
+			if (!address || typeof address === "string")
+				throw new Error("test server did not bind");
+			const spec = resolveCuaRuntimeSpec(
+				"anthropic:claude-sonnet-4-5",
+				{ nativeTool: { type: "browser_20260701" } },
+			);
+			const model = {
+				...spec.model,
+				api: ANTHROPIC_NATIVE_BROWSER_MESSAGES_API,
+				baseUrl: `http://127.0.0.1:${address.port}`,
+			} as Model<any>;
+			const nativePlaceholder = {
+				...dynamicTool,
+				...spec.toolDefinitions[0],
+			} as AgentTool;
+			const payload = await capturePayload(
+				model,
+				{
+					systemPrompt: "test",
+					tools: [nativePlaceholder, addTool, dynamicTool],
+					messages: addMessages("anthropic-messages"),
+				},
+				spec.onPayload,
+			);
+			expect(payload.tools).toEqual(
+				expect.arrayContaining([
+					expect.objectContaining({
+						name: "learned_tool",
+						defer_loading: true,
+					}),
+				]),
+			);
+			expect(payload.tools).toEqual(
+				expect.arrayContaining([
+					expect.objectContaining({ type: "browser_20260701" }),
+				]),
+			);
+			expect(requestHeaders["anthropic-beta"]).toContain(
+				"browser-use-2026-07-01",
+			);
+		} finally {
+			await new Promise<void>((resolve, reject) =>
+				server.close((error) => (error ? reject(error) : resolve())),
+			);
+		}
+	});
+});