fix(extensions): wire invokeTool to the extension path as same-tool delegation

Addresses PR review. The initial version put invokeTool on AgentToolContext via ToolContextStore, but the extension execute path (RegisteredToolAdapter) builds its own ExtensionContext and never saw it, so the documented registerTool wrapper use case did not work. It also allowed arbitrary cross-tool targets (bypassing the target's approval policy), used a session-global recursion counter that tripped on concurrent independent delegations, and missed discoverable built-ins that xdev partitioning moves out of the tool array. Rework: - Move invokeTool onto ExtensionContext, and bind it in RegisteredToolAdapter to the tool's own name, so a re-registered built-in actually receives it. - Make delegation same-tool only: invokeTool takes just (params, options) and runs the native built-in of the caller's own name. It cannot reach an arbitrary target, so it cannot escalate past the approval already granted for the call, and the native call is not re-gated. - Track recursion depth per call chain (threaded through invokeNativeTool and createContext) instead of session-global state, so concurrent delegations do not interfere. - Seed the native resolver from the xdev registry when present (it retains discoverable built-ins like browser), else the built-in registry. Replaces the ToolContextStore-level unit test with an end-to-end test that registers a built-in wrapper through the extension/session path and asserts the native tool runs the wrapper's delegated input.
2026-07-30 10:35:12 -07:00
parent 6acf957dd8
commit a50bcd5fed
9 changed files with 231 additions and 195 deletions
@@ -1,7 +1,13 @@
 /**
 * Extension runner - executes extensions and manages their lifecycle.
 */
-import type { AgentMessage } from "@oh-my-pi/pi-agent-core";
+import type {
+	AgentMessage,
+	AgentTool,
+	AgentToolContext,
+	AgentToolResult,
+	AgentToolUpdateCallback,
+} from "@oh-my-pi/pi-agent-core";
 import type { CredentialDisabledEvent, ImageContent, Model, ProviderResponseMetadata } from "@oh-my-pi/pi-ai";
 import type { KeyId } from "@oh-my-pi/pi-tui";
 import { logger } from "@oh-my-pi/pi-utils";
@@ -407,6 +413,56 @@ export class ExtensionRunner {
 		return this.#emittedToolCalls.delete(`${toolCallId}:${toolName}`);
 	}

+	/**
+	 * Resolves a tool NAME to its native built-in implementation (the pre-extension-override,
+	 * unwrapped tool) plus a factory for the `AgentToolContext` that native tool expects, or
+	 * undefined when no native built-in of that name exists. Set by the SDK; backs same-tool
+	 * `invokeTool`. The context factory is the same one the agent loop uses for tool execution, so a
+	 * delegated native call sees the ordinary session tool context (ui, cwd, snapshot state, etc.).
+	 */
+	#nativeToolResolver?: (name: string) => { tool: AgentTool; makeContext: () => AgentToolContext } | undefined;
+
+	/** Wires the native-tool resolver used by {@link invokeNativeTool}. */
+	setNativeToolResolver(
+		resolve: (name: string) => { tool: AgentTool; makeContext: () => AgentToolContext } | undefined,
+	): void {
+		this.#nativeToolResolver = resolve;
+	}
+
+	/** Whether a native built-in of `name` is available to delegate to. */
+	hasNativeTool(name: string): boolean {
+		return this.#nativeToolResolver?.(name) !== undefined;
+	}
+
+	/**
+	 * Run the native built-in of `name` with `params` and return its result — the delegation target
+	 * of a same-tool `ctx.invokeTool`. Calls the unwrapped native `execute` directly with the loop's
+	 * ordinary tool context, so it inherits the caller's already-granted approval (the caller is the
+	 * same tool) rather than re-running the gate. `depth` guards a wrapper that recurses into itself;
+	 * it is per call chain (threaded from the caller), not session-global, so concurrent independent
+	 * delegations do not interfere.
+	 */
+	async invokeNativeTool<TDetails = unknown>(
+		name: string,
+		params: Record<string, unknown>,
+		options?: { signal?: AbortSignal; onUpdate?: AgentToolUpdateCallback<TDetails>; depth?: number },
+	): Promise<AgentToolResult<TDetails>> {
+		const resolved = this.#nativeToolResolver?.(name);
+		if (!resolved) throw new Error(`invokeTool: no native built-in named "${name}" to delegate to`);
+		const depth = options?.depth ?? 0;
+		if (depth >= 8) {
+			throw new Error(`invokeTool: delegation depth exceeded 8 (recursive invokeTool for "${name}"?)`);
+		}
+		const toolCallId = `invoke-${name}-${Date.now().toString(36)}-${depth}`;
+		return (await resolved.tool.execute(
+			toolCallId,
+			params as never,
+			options?.signal,
+			options?.onUpdate as never,
+			resolved.makeContext(),
+		)) as AgentToolResult<TDetails>;
+	}
+
 	constructor(
 		private readonly extensions: Extension[],
 		private readonly runtime: ExtensionRuntime,
@@ -728,8 +784,13 @@ export class ExtensionRunner {
 		return undefined;
 	}

-	/** Creates an extension context, optionally scoped to a provider request model. */
-	createContext(model?: Model): ExtensionContext {
+	/**
+	 * Creates an extension context, optionally scoped to a provider request model. When `toolName` is
+	 * given and a native built-in of that name exists, the context carries a same-tool `invokeTool`
+	 * that delegates to that native implementation (see {@link invokeNativeTool}); `depth` threads the
+	 * per-chain recursion counter so a wrapper that re-invokes itself is bounded.
+	 */
+	createContext(model?: Model, toolName?: string, depth = 0): ExtensionContext {
 		const getModel = model ? () => model : this.#getModel;
 		return {
 			ui: this.#uiContext,
@@ -754,6 +815,15 @@ export class ExtensionRunner {
 			setInterval: (callback, ms, ...args) => this.#managedTimers.setInterval(callback, ms, ...args),
 			setTimeout: (callback, ms, ...args) => this.#managedTimers.setTimeout(callback, ms, ...args),
 			clearTimer: timer => this.#managedTimers.clear(timer),
+			invokeTool:
+				toolName !== undefined && this.hasNativeTool(toolName)
+					? (params, options) =>
+							this.invokeNativeTool(toolName, params, {
+								signal: options?.signal,
+								onUpdate: options?.onUpdate,
+								depth: depth + 1,
+							})
+					: undefined,
 		};
 	}

@@ -464,6 +464,21 @@ export interface ExtensionContext {
 	setTimeout(callback: (...args: unknown[]) => void, ms?: number, ...args: unknown[]): Timer;
 	/** Clear a timer scheduled via {@link setInterval} or {@link setTimeout}. */
 	clearTimer(timer: Timer): void;
+	/**
+	 * Run the NATIVE built-in implementation of the tool this handler re-registered, with `params`,
+	 * and return its result. Lets a tool that re-registers a built-in (e.g. wrapping `write` to add
+	 * logging or a policy check) delegate to the original instead of reimplementing it — the native
+	 * tool performs its own side effects and internal bookkeeping.
+	 *
+	 * Delegation is same-tool only: it invokes the built-in of the SAME name as the registering tool,
+	 * never an arbitrary target, so it cannot escalate past the approval already granted for this
+	 * call. Present only when a native built-in of that name exists (undefined otherwise, e.g. for a
+	 * net-new tool that shadows no built-in). Recursion is depth-guarded per call chain.
+	 */
+	invokeTool?<TDetails = unknown>(
+		params: Record<string, unknown>,
+		options?: { signal?: AbortSignal; onUpdate?: AgentToolUpdateCallback<TDetails> },
+	): Promise<AgentToolResult<TDetails>>;
 }

 /**
@@ -66,7 +66,15 @@ export class RegisteredToolAdapter implements AgentTool<any, any, any> {
 		onUpdate?: AgentToolUpdateCallback<any>,
 		_context?: AgentToolContext,
 	) {
-		return this.registeredTool.definition.execute(toolCallId, params, signal, onUpdate, this.runner.createContext());
+		// Bind the extension context to this tool's own name so `ctx.invokeTool` delegates to the
+		// native built-in of the same name (present only when this tool re-registers a built-in).
+		return this.registeredTool.definition.execute(
+			toolCallId,
+			params,
+			signal,
+			onUpdate,
+			this.runner.createContext(undefined, this.registeredTool.definition.name),
+		);
 	}
 }