diff --git a/docs/extensions.md b/docs/extensions.md index 5bbe0e5ab..cafb561ed 100644 --- a/docs/extensions.md +++ b/docs/extensions.md @@ -312,6 +312,26 @@ execute( ): Promise ``` +### Delegating to a native built-in (`ctx.invokeTool`) + +A tool that re-registers a built-in name (e.g. wrapping `write` to add logging or a policy check) can +run the original instead of reimplementing it. When your registered tool shadows a built-in, the `ctx` +passed to `execute` carries: + +```ts +ctx.invokeTool?( + params: Record, + options?: { signal?: AbortSignal; onUpdate?: AgentToolUpdateCallback }, +): Promise> +``` + +It runs the **native** built-in of the same name as your tool (delegation is same-tool only, so it +cannot reach an arbitrary target or escalate past the approval already granted for this call) and +returns its result, including the native tool's own side effects and internal bookkeeping. It is +present only when a native built-in of that name exists — `ctx.invokeTool` is `undefined` for a +net-new tool that shadows no built-in. The native call is not re-gated, since it is the same tool you +are already approved as, and delegation depth is guarded against accidental self-recursion. + Template: ```ts diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index f65966890..fe49cc599 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -44,6 +44,7 @@ - Fixed the Python RPC client dropping current context, compaction, OAuth URL, and terminal-settlement fields, and made additive notification variants observable without stopping the stdout reader. - Fixed the browser tool silently ignoring `url` when opening a new tab on an attached browser (`app.cdp_url` or `app.path`), so the tab now navigates on open exactly as it already did on reuse and in headless mode. - Fixed browser automation disrupting a browser it attached to over `app.cdp_url`: the tool now adopts the tab the user actually has in the foreground and no longer raises its own tab when taking a screenshot. Owned and headless browsers keep activating the target before capture. +- Added `ctx.invokeTool(params, options?)` to a re-registered built-in's extension context, letting a wrapper run the native tool of the same name instead of reimplementing it, and inheriting the caller's context, abort signal, and progress updates. ## [17.2.1] - 2026-07-30 diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index 3685ade92..af1f931c7 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -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,67 @@ 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( + name: string, + params: Record, + options?: { + signal?: AbortSignal; + onUpdate?: AgentToolUpdateCallback; + depth?: number; + /** + * The caller tool's own context. Reused for the native call so metadata the native tool + * reads — `toolCall` (write/edit LSP batch flushing) and provider metadata / + * `providerSafetyApproved` (computer) — is preserved. Falls back to a fresh session tool + * context only when the caller had none. + */ + callerContext?: AgentToolContext; + }, + ): Promise> { + 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, + options?.callerContext ?? resolved.makeContext(), + )) as AgentToolResult; + } + constructor( private readonly extensions: Extension[], private readonly runtime: ExtensionRuntime, @@ -747,8 +814,27 @@ 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. + * + * `delegation` wires the same-tool `ctx.invokeTool` for a re-registered built-in: when `toolName` + * names an existing native built-in, the context carries an `invokeTool` that runs it (see + * {@link invokeNativeTool}). The rest inherits the wrapper's own call so a bare + * `ctx.invokeTool(params)` behaves like the outer call — `context` preserves `toolCall`/provider + * metadata, `signal`/`onUpdate` default to the wrapper's own channels so aborting the outer tool + * call stops the native one and native progress still streams, and `depth` bounds recursion per + * call chain. Explicit options passed to `invokeTool` override the inherited `signal`/`onUpdate`. + */ + createContext( + model?: Model, + delegation?: { + toolName: string; + depth?: number; + context?: AgentToolContext; + signal?: AbortSignal; + onUpdate?: AgentToolUpdateCallback; + }, + ): ExtensionContext { const getModel = model ? () => model : this.#getModel; return { ui: this.#uiContext, @@ -773,6 +859,18 @@ 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: + delegation !== undefined && this.hasNativeTool(delegation.toolName) + ? (params, options) => + this.invokeNativeTool(delegation.toolName, params, { + // Inherit the wrapper's own channels so a bare `ctx.invokeTool(params)` aborts + // and streams with the outer call. Explicit options win. + signal: options?.signal ?? delegation.signal, + onUpdate: options?.onUpdate ?? delegation.onUpdate, + depth: (delegation.depth ?? 0) + 1, + callerContext: delegation.context, + }) + : undefined, }; } diff --git a/packages/coding-agent/src/extensibility/extensions/types.ts b/packages/coding-agent/src/extensibility/extensions/types.ts index 5ef2cb872..232291225 100644 --- a/packages/coding-agent/src/extensibility/extensions/types.ts +++ b/packages/coding-agent/src/extensibility/extensions/types.ts @@ -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?( + params: Record, + options?: { signal?: AbortSignal; onUpdate?: AgentToolUpdateCallback }, + ): Promise>; } /** diff --git a/packages/coding-agent/src/extensibility/extensions/wrapper.ts b/packages/coding-agent/src/extensibility/extensions/wrapper.ts index 093abdf1f..8ce2305f2 100644 --- a/packages/coding-agent/src/extensibility/extensions/wrapper.ts +++ b/packages/coding-agent/src/extensibility/extensions/wrapper.ts @@ -64,9 +64,26 @@ export class RegisteredToolAdapter implements AgentTool { params: any, signal?: AbortSignal, onUpdate?: AgentToolUpdateCallback, - _context?: AgentToolContext, + 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). The + // wrapper's own context, abort signal, and progress callback are inherited by the delegated + // call, so a bare `ctx.invokeTool(params)` keeps the caller's `toolCall`/provider metadata + // (write/edit LSP batching, computer safety acknowledgement), stops when the outer call is + // aborted, and still streams native progress. + return this.registeredTool.definition.execute( + toolCallId, + params, + signal, + onUpdate, + this.runner.createContext(undefined, { + toolName: this.registeredTool.definition.name, + context, + signal, + onUpdate, + }), + ); } } diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index 4528581f6..393044ead 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -2572,6 +2572,14 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} autoApprove: options.autoApprove ?? false, }); const toolContextStore = new ToolContextStore(getSessionContext); + // 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 + // re-running the gate. Seeded from the xdev registry when present (it retains discoverable + // built-ins like `browser` that xdev partitioning removes from the active tool array), else + // from the built-in registry; captured before the ExtensionToolWrapper pass so the natives + // stay unwrapped. The extension runner exposes it to re-registered tools via createContext. + const nativeToolsByName = new Map(toolSession.xdev?.tools ?? undefined); const registeredTools = restrictToolNames ? [] : extensionRunner.getAllRegisteredTools(); const sdkCustomTools = @@ -2595,17 +2603,32 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} // All built-in tools are active (conditional tools like git/ask return null from factory if disabled) const builtInRegistryToolNames = toolSession.xdev?.builtInNames ?? new Set(toolRegistry.keys()); + // Capture the native built-in implementations before extension re-registration replaces registry + // entries and before the ExtensionToolWrapper pass below, so `ctx.invokeTool` reaches the + // unwrapped native execute (inheriting the caller's already-granted approval, not re-gating). + for (const [name, tool] of toolRegistry) { + nativeToolsByName.set(name, tool); + } if (!restrictToolNames && !toolRegistry.has("goal") && settings.get("goal.enabled")) { const goalTool = await logger.time("createTools:goal:session", HIDDEN_TOOLS.goal, toolSession); if (goalTool) { - toolRegistry.set(goalTool.name, wrapToolWithMetaNotice(goalTool)); + const wrapped = wrapToolWithMetaNotice(goalTool); + toolRegistry.set(goalTool.name, wrapped); builtInRegistryToolNames.add(goalTool.name); + nativeToolsByName.set(goalTool.name, wrapped); } } for (const tool of wrappedExtensionTools) { toolRegistry.set(tool.name, tool); builtInRegistryToolNames.delete(tool.name); } + // Expose the native built-ins to same-tool `ctx.invokeTool` on re-registered tools. Set after + // the override loop so the map holds the natives, not the extension replacements. The context + // factory is the loop's own tool context, so a delegated native call sees ordinary session state. + extensionRunner.setNativeToolResolver(name => { + const tool = nativeToolsByName.get(name); + return tool ? { tool, makeContext: () => toolContextStore.getContext() } : undefined; + }); if (deferMCPDiscoveryForUI && mcpManager) { for (const name of collectPendingMCPToolNames(options.toolNames)) { if (!toolRegistry.has(name)) { @@ -2668,11 +2691,10 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} writeRegistration ??= (async () => { const writeTool = await logger.time("createTools:write:session", BUILTIN_TOOLS.write, toolSession); if (!writeTool || toolRegistry.has("write")) return builtInRegistryToolNames.has("write"); - toolRegistry.set( - writeTool.name, - new ExtensionToolWrapper(wrapToolWithMetaNotice(writeTool), extensionRunner) as Tool, - ); + const nativeWrite = wrapToolWithMetaNotice(writeTool); + toolRegistry.set(writeTool.name, new ExtensionToolWrapper(nativeWrite, extensionRunner) as Tool); builtInRegistryToolNames.add(writeTool.name); + nativeToolsByName.set(writeTool.name, nativeWrite); return true; })().finally(() => { writeRegistration = undefined; diff --git a/packages/coding-agent/test/agent-session-message-pipeline.test.ts b/packages/coding-agent/test/agent-session-message-pipeline.test.ts index b7662a9a7..70b215059 100644 --- a/packages/coding-agent/test/agent-session-message-pipeline.test.ts +++ b/packages/coding-agent/test/agent-session-message-pipeline.test.ts @@ -25,7 +25,7 @@ import { Settings } from "@oh-my-pi/pi-coding-agent/config/settings"; import * as memoryBackend from "@oh-my-pi/pi-coding-agent/memory-backend"; import type { MemoryBackend } from "@oh-my-pi/pi-coding-agent/memory-backend/types"; import { type MnemopiSessionState, setMnemopiSessionState } from "@oh-my-pi/pi-coding-agent/mnemopi/state"; -import { createAgentSession, type ExtensionFactory } from "@oh-my-pi/pi-coding-agent/sdk"; +import { createAgentSession, type ExtensionContext, type ExtensionFactory } from "@oh-my-pi/pi-coding-agent/sdk"; import { obfuscateProviderContext, SecretObfuscator } from "@oh-my-pi/pi-coding-agent/secrets"; import { AgentSession, type AgentSessionEvent } from "@oh-my-pi/pi-coding-agent/session/agent-session"; import { AuthStorage } from "@oh-my-pi/pi-coding-agent/session/auth-storage"; @@ -945,6 +945,111 @@ describe("AgentSession message pipeline", () => { authStorage.close(); } }); + it("exposes ctx.invokeTool to a re-registered built-in so it can delegate to the native tool", async () => { + // End-to-end for the extension path: a tool that re-registers `bash` receives ctx.invokeTool + // (bound to its own name), delegates to the native bash, and the native output flows back. + using tempDir = TempDir.createSync("@pi-invoke-tool-"); + const api = "test-invoke-tool"; + let requests = 0; + registerCustomApi(api, () => { + requests++; + const stream = new AssistantMessageEventStream(); + queueMicrotask(() => { + if (requests === 1) { + const message = createAssistantMessage(""); + const toolCall = { + type: "toolCall", + id: "call-invoke-1", + name: "bash", + arguments: { command: "echo from-model" }, + } as const; + message.content = [toolCall]; + message.stopReason = "toolUse"; + stream.push({ type: "toolcall_start", contentIndex: 0, partial: message }); + stream.push({ type: "toolcall_end", contentIndex: 0, toolCall: toolCall as never, partial: message }); + stream.push({ type: "done", reason: "toolUse", message }); + } else { + const message = createAssistantMessage("done"); + stream.push({ type: "done", reason: "stop", message }); + } + }); + return stream; + }); + const model = buildModel({ + id: "local-invoke-model", + name: "Local Invoke Model", + api, + provider: "ollama", + baseUrl: "http://127.0.0.1:11434", + reasoning: false, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 4096, + maxTokens: 1024, + } as ModelSpec) as Model; + let invokeToolPresent = false; + let delegatedText = ""; + // Re-register `bash`: the wrapper ignores the model's args, delegates to the native bash with + // its own command via ctx.invokeTool, and returns the native result. + const wrapBash: ExtensionFactory = pi => { + pi.registerTool({ + name: "bash", + label: "Bash", + description: "wrapped bash", + parameters: pi.zod.object({ command: pi.zod.string() }), + async execute( + _toolCallId: string, + _params: unknown, + _signal: unknown, + _onUpdate: unknown, + ctx: ExtensionContext, + ) { + invokeToolPresent = typeof ctx.invokeTool === "function"; + const native = await ctx.invokeTool?.({ command: "echo from-wrapper" }); + const textBlock = native?.content.find(b => b.type === "text"); + delegatedText = textBlock?.type === "text" ? textBlock.text : ""; + return native ?? { content: [{ type: "text" as const, text: "no invokeTool" }], details: {} }; + }, + }); + }; + const authStorage = await AuthStorage.create(tempDir.join("auth.db")); + const modelRegistry = new ModelRegistry(authStorage, tempDir.join("models.yml")); + const { session } = await createAgentSession({ + cwd: tempDir.path(), + agentDir: tempDir.path(), + sessionManager: SessionManager.inMemory(tempDir.path()), + authStorage, + modelRegistry, + settings: Settings.isolated({ + "compaction.enabled": false, + "bash.autoBackground.enabled": false, + "bashInterceptor.enabled": false, + "tools.xdev": false, + }), + model, + disableExtensionDiscovery: true, + extensions: [wrapBash], + skills: [], + contextFiles: [], + promptTemplates: [], + slashCommands: [], + enableMCP: false, + enableLsp: false, + skipPythonPreflight: true, + toolNames: ["bash"], + }); + try { + await session.sendUserMessage("run it"); + + expect(invokeToolPresent).toBe(true); + // The native bash actually ran the wrapper's command, not the model's. + expect(delegatedText).toContain("from-wrapper"); + expect(delegatedText).not.toContain("from-model"); + } finally { + await session.dispose(); + authStorage.close(); + } + }); it("clears promoted memory from the base prompt when switching sessions", async () => { using tempDir = TempDir.createSync("@pi-injected-memory-switch-"); diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index 81fc506a4..a8cdd8324 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -3126,4 +3126,91 @@ describe("ExtensionRunner", () => { } }); }); + + describe("invokeTool same-tool delegation", () => { + // Records what the native tool actually received, so the inherited abort/progress channels and + // the caller context are observable. + function nativeProbe(seen: { signal?: AbortSignal; onUpdate?: unknown; params?: unknown }): AgentTool { + return { + name: "bash", + label: "Bash", + description: "native bash", + parameters: Type.Object({ command: Type.String() }), + execute: async (_id: string, params: unknown, signal?: AbortSignal, onUpdate?: unknown) => { + seen.params = params; + seen.signal = signal; + seen.onUpdate = onUpdate; + return { content: [{ type: "text", text: "native ran" }], details: {} }; + }, + } as AgentTool; + } + + const runnerWithNative = async (native: AgentTool) => { + const result = await loadTestExtensions(); + const runner = new ExtensionRunner( + result.extensions, + result.runtime, + tempDir.path(), + sessionManager, + modelRegistry, + ); + runner.setNativeToolResolver(name => + name === native.name ? { tool: native, makeContext: () => ({}) as never } : undefined, + ); + return runner; + }; + + it("inherits the wrapper call's signal and onUpdate for a bare invokeTool", async () => { + const seen: { signal?: AbortSignal; onUpdate?: unknown; params?: unknown } = {}; + const runner = await runnerWithNative(nativeProbe(seen)); + const controller = new AbortController(); + const onUpdate = () => {}; + + const ctx = runner.createContext(undefined, { + toolName: "bash", + signal: controller.signal, + onUpdate, + }); + await ctx.invokeTool?.({ command: "echo hi" }); + + // Aborting the outer tool call must reach the native one, and native progress must stream. + expect(seen.signal).toBe(controller.signal); + expect(seen.onUpdate).toBe(onUpdate); + expect(seen.params).toEqual({ command: "echo hi" }); + }); + + it("lets explicit invokeTool options override the inherited channels", async () => { + const seen: { signal?: AbortSignal; onUpdate?: unknown; params?: unknown } = {}; + const runner = await runnerWithNative(nativeProbe(seen)); + const outer = new AbortController(); + const inner = new AbortController(); + const innerOnUpdate = () => {}; + + const ctx = runner.createContext(undefined, { + toolName: "bash", + signal: outer.signal, + onUpdate: () => {}, + }); + await ctx.invokeTool?.({ command: "echo hi" }, { signal: inner.signal, onUpdate: innerOnUpdate }); + + expect(seen.signal).toBe(inner.signal); + expect(seen.onUpdate).toBe(innerOnUpdate); + }); + + it("omits invokeTool when no native built-in of that name exists", async () => { + const runner = await runnerWithNative(nativeProbe({})); + expect(runner.createContext(undefined, { toolName: "not_a_builtin" }).invokeTool).toBeUndefined(); + // Also absent when the context is not scoped to a tool at all. + expect(runner.createContext().invokeTool).toBeUndefined(); + }); + + it("bounds recursion per call chain", async () => { + const runner = await runnerWithNative(nativeProbe({})); + await expect(runner.invokeNativeTool("bash", { command: "echo hi" }, { depth: 8 })).rejects.toThrow( + /delegation depth exceeded/, + ); + // A fresh chain at depth 0 is unaffected by another chain's depth. + await expect(runner.invokeNativeTool("bash", { command: "echo hi" }, { depth: 0 })).resolves.toBeDefined(); + }); + }); });