From 6acf957dd825460f5b4faa33caff08937f389c19 Mon Sep 17 00:00:00 2001 From: Larry Gordon Date: Mon, 27 Jul 2026 15:12:28 -0700 Subject: [PATCH 1/5] feat(extensions): add ctx.invokeTool for native built-in delegation A tool's execute context now carries invokeTool(name, params, options?), which runs the native built-in of `name` and returns its result. A tool that re-registers a built-in (e.g. wrapping write to add logging or a policy check) can delegate to the original instead of reimplementing it. The native implementation is captured before extension re-registration replaces the registry entry and before the ExtensionToolWrapper pass, so invokeTool reaches the unwrapped native execute: it does not recurse into the caller's own wrapper, and it inherits the caller's already-granted approval rather than re-running the gate. Delegation depth is guarded against accidental self-recursion, and it resolves to undefined when no native tool of that name exists. Wired through ToolContextStore with a lazy native-tool resolver, so it is coding-agent-only (no agent-loop change) and sees the fully-assembled built-in set at call time. --- docs/extensions.md | 19 ++++ packages/coding-agent/CHANGELOG.md | 1 + packages/coding-agent/src/sdk.ts | 25 +++-- packages/coding-agent/src/tools/context.ts | 78 ++++++++++++++- .../test/tool-context-invoke.test.ts | 98 +++++++++++++++++++ 5 files changed, 213 insertions(+), 8 deletions(-) create mode 100644 packages/coding-agent/test/tool-context-invoke.test.ts diff --git a/docs/extensions.md b/docs/extensions.md index 33c0af111..fd1e10954 100644 --- a/docs/extensions.md +++ b/docs/extensions.md @@ -312,6 +312,25 @@ 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. The `ctx` passed to `execute` carries: + +```ts +ctx.invokeTool?( + name: string, + params: Record, + options?: { signal?: AbortSignal; onUpdate?: AgentToolUpdateCallback }, +): Promise | undefined> +``` + +It runs the **native** built-in of `name` (bypassing your own re-registration, so it does not recurse +into your wrapper) and returns its result, including the native tool's own side effects and internal +bookkeeping. It resolves to `undefined` when there is no native tool of that name. The invoked tool's +approval gate is not re-run — your call already passed approval — 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 25c809467..0d18ada49 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -158,6 +158,7 @@ - Added the bundled `ts-no-local-is-record` TTSR rule, which catches local `isRecord` function and lambda definitions and directs agents to shared guards plus explicit shape validation. - A `tool_call` handler (extension or hook) can now return `input` to revise the arguments a tool executes with, not just `block` it. The returned object is the raw execution input passed to the tool (ignored when `block` is set, and not applied to `computer` tool calls), enabling wrappers that normalize or rewrite a built-in's arguments without reimplementing the tool. For model-issued calls the event fires at arg-prep time in the agent loop, so a revision is revalidated against the tool schema and is what concurrency scheduling, `tool_execution_start`/transcripts, the persisted assistant message, and the approval gate all observe — the user approves exactly what runs, and a revision that changes a tool's functional concurrency (e.g. bash `pty`) schedules correctly. A revised nested `write xd://` device dispatch forfeits the outer write gate's approval and faces the full prompt again ([#6681](https://github.com/can1357/oh-my-pi/pull/6681) by [@psyrendust](https://github.com/psyrendust)). +- A tool's `execute` context now carries `ctx.invokeTool(name, params, options?)`, which runs the native built-in of `name` and returns its result. A tool that re-registers a built-in (e.g. wrapping `write`) can delegate to the original instead of reimplementing it: the native tool performs its own side effects and internal bookkeeping, the call bypasses the caller's own re-registration so it does not recurse, the invoked tool's approval gate is not re-run (the caller already passed approval), and delegation depth is guarded. Resolves to `undefined` when no native tool of that name exists. - Added a parser for macOS `sample`(1) call-tree reports to the read tool: `*.sample.txt` reads now return a compact bottleneck summary — per-thread hot paths with on-CPU sample counts (blocked syscall time excluded), demangled Rust v0/legacy symbols, flattened direct recursion, merged call-site siblings, idle-thread classification, and a process-wide top-functions-by-self-samples table. `:raw` still reads the original report, and files that merely carry the extension fall back to plain text. - Added V8 `.cpuprofile` support to the read tool (Node/Bun `--cpu-prof`, Chrome DevTools, CDP `Profiler.stop` output): reads now return a compact bottleneck summary — hot-path call tree with on-CPU milliseconds (`(idle)` time excluded), collapsed pass-through chains, flattened direct recursion, shortened file URLs, and a top-functions-by-self-time table. `:raw` still reads the original JSON, and files that merely carry the extension fall back to plain text. diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index 53ac1bb6e..618b3c477 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -2563,7 +2563,13 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} localProtocolOptions, autoApprove: options.autoApprove ?? false, }); - const toolContextStore = new ToolContextStore(getSessionContext); + // Native built-in implementations, captured before extension re-registration replaces registry + // entries and before the ExtensionToolWrapper pass. Backs `ctx.invokeTool`, so a tool that + // wraps a built-in (e.g. a re-registered `write`) can delegate to the original — reaching the + // unwrapped native execute, which correctly inherits the caller's already-granted approval + // rather than re-running the gate. Populated at the built-in registry loop below. + const nativeToolsByName = new Map(); + const toolContextStore = new ToolContextStore(getSessionContext, name => nativeToolsByName.get(name)); const registeredTools = restrictToolNames ? [] : extensionRunner.getAllRegisteredTools(); const sdkCustomTools = @@ -2587,11 +2593,19 @@ 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) { @@ -2660,11 +2674,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/src/tools/context.ts b/packages/coding-agent/src/tools/context.ts index e8e5d1560..534f789e9 100644 --- a/packages/coding-agent/src/tools/context.ts +++ b/packages/coding-agent/src/tools/context.ts @@ -1,7 +1,21 @@ -import type { AgentToolContext, ToolCallContext } from "@oh-my-pi/pi-agent-core"; +import type { + AgentTool, + AgentToolContext, + AgentToolResult, + AgentToolUpdateCallback, + ToolCallContext, +} from "@oh-my-pi/pi-agent-core"; import type { CustomToolContext } from "../extensibility/custom-tools/types"; import type { ExtensionUIContext } from "../extensibility/extensions/types"; +/** Options a tool passes when delegating to another tool's native implementation via {@link AgentToolContext.invokeTool}. */ +export interface InvokeToolOptions { + /** Abort signal forwarded to the invoked tool's `execute`. */ + signal?: AbortSignal; + /** Progress callback forwarded to the invoked tool's `execute`. */ + onUpdate?: AgentToolUpdateCallback; +} + declare module "@oh-my-pi/pi-agent-core" { interface AgentToolContext extends CustomToolContext { ui?: ExtensionUIContext; @@ -15,15 +29,41 @@ declare module "@oh-my-pi/pi-agent-core" { xdevApproved?: boolean; /** Set only after an interactive prompt approves provider computer safety checks. */ providerSafetyApproved?: boolean; + /** + * Run the NATIVE built-in implementation of `name` with `params` and return its result, + * bypassing any extension re-registration of that name. Lets a tool that re-registers a + * built-in (e.g. wrapping `write`) delegate to the original instead of reimplementing it — + * the native tool performs its own side effects and internal bookkeeping. Resolves to + * `undefined` when no native tool of that name exists. The invoked tool's approval gate is + * NOT re-run: the caller is itself an already-approved tool call. Recursion is depth-guarded. + */ + invokeTool?( + name: string, + params: Record, + options?: InvokeToolOptions, + ): Promise | undefined>; } } +/** Max depth for `invokeTool` delegation chains — guards a re-registered tool that recurses into itself. */ +const MAX_INVOKE_DEPTH = 8; + export class ToolContextStore { #uiContext: ExtensionUIContext | undefined; #hasUI = false; #toolNames: string[] = []; + #invokeDepth = 0; - constructor(private readonly getBaseContext: () => CustomToolContext) {} + /** + * @param getBaseContext builds the per-call base tool context. + * @param resolveNativeTool resolves a tool NAME to its native built-in implementation (the + * pre-extension-override tool), or undefined if there is no native tool of that name. Used to + * back `ctx.invokeTool`. Lazy: called at invoke time, after the registry is fully assembled. + */ + constructor( + private readonly getBaseContext: () => CustomToolContext, + private readonly resolveNativeTool?: (name: string) => AgentTool | undefined, + ) {} getContext(toolCall?: ToolCallContext): AgentToolContext { return { @@ -32,9 +72,43 @@ export class ToolContextStore { hasUI: this.#hasUI, toolNames: this.#toolNames, toolCall, + invokeTool: this.resolveNativeTool + ? (name, params, options) => this.#invokeTool(name, params, options) + : undefined, }; } + async #invokeTool( + name: string, + params: Record, + options?: InvokeToolOptions, + ): Promise | undefined> { + const native = this.resolveNativeTool?.(name); + if (!native) return undefined; + if (this.#invokeDepth >= MAX_INVOKE_DEPTH) { + throw new Error( + `invokeTool: delegation depth exceeded ${MAX_INVOKE_DEPTH} (recursive invokeTool for "${name}"?)`, + ); + } + // Nested context so the invoked tool sees the same session state (ui, cwd, etc.) and can itself + // delegate. Its approval gate is NOT re-run: it is reached via the native execute directly, not + // the ExtensionToolWrapper, so the caller's already-granted approval covers it. + const nestedContext = this.getContext(undefined); + this.#invokeDepth++; + try { + const toolCallId = `invoke-${name}-${Date.now().toString(36)}-${this.#invokeDepth}`; + return (await native.execute( + toolCallId, + params as never, + options?.signal, + options?.onUpdate as never, + nestedContext, + )) as AgentToolResult; + } finally { + this.#invokeDepth--; + } + } + setUIContext(uiContext: ExtensionUIContext, hasUI: boolean): void { this.#uiContext = uiContext; this.#hasUI = hasUI; diff --git a/packages/coding-agent/test/tool-context-invoke.test.ts b/packages/coding-agent/test/tool-context-invoke.test.ts new file mode 100644 index 000000000..597c50956 --- /dev/null +++ b/packages/coding-agent/test/tool-context-invoke.test.ts @@ -0,0 +1,98 @@ +/** + * Tests for `ctx.invokeTool` — a tool delegating to another tool's native built-in implementation + * via the ToolContextStore-built AgentToolContext. + */ + +import { describe, expect, it } from "bun:test"; +import type { AgentTool, AgentToolResult } from "@oh-my-pi/pi-agent-core"; +import { Type } from "@oh-my-pi/pi-coding-agent/extensibility/typebox"; +import { ToolContextStore } from "@oh-my-pi/pi-coding-agent/tools/context"; + +const baseContext = () => ({}) as never; + +function makeTool(name: string, exec: AgentTool["execute"]): AgentTool { + return { + name, + label: name, + description: `${name} tool`, + parameters: Type.Object({}), + execute: exec, + } as AgentTool; +} + +describe("ctx.invokeTool native-tool delegation", () => { + it("is absent when the store has no native-tool resolver", () => { + const store = new ToolContextStore(baseContext); + expect(store.getContext().invokeTool).toBeUndefined(); + }); + + it("runs the native tool and returns its result", async () => { + const seen: unknown[] = []; + const native = makeTool("write", async (_id, params) => { + seen.push(params); + return { content: [{ type: "text", text: "native wrote" }], details: { ok: true } } as AgentToolResult; + }); + const store = new ToolContextStore(baseContext, name => (name === "write" ? native : undefined)); + + const result = await store.getContext().invokeTool?.("write", { path: "/tmp/x", content: "hi" }); + + expect(seen).toEqual([{ path: "/tmp/x", content: "hi" }]); + expect(result?.content).toEqual([{ type: "text", text: "native wrote" }]); + expect(result?.details).toEqual({ ok: true }); + }); + + it("resolves to undefined for an unknown tool name", async () => { + const store = new ToolContextStore(baseContext, () => undefined); + expect(await store.getContext().invokeTool?.("nope", {})).toBeUndefined(); + }); + + it("forwards signal and onUpdate to the native tool", async () => { + const controller = new AbortController(); + let receivedSignal: AbortSignal | undefined; + let receivedOnUpdate: unknown; + const native = makeTool("read", async (_id, _params, signal, onUpdate) => { + receivedSignal = signal; + receivedOnUpdate = onUpdate; + return { content: [{ type: "text", text: "ok" }], details: {} } as AgentToolResult; + }); + const store = new ToolContextStore(baseContext, () => native); + const onUpdate = () => {}; + + await store.getContext().invokeTool?.("read", {}, { signal: controller.signal, onUpdate }); + + expect(receivedSignal).toBe(controller.signal); + expect(receivedOnUpdate).toBe(onUpdate); + }); + + it("passes a context whose invokeTool lets the native tool delegate further", async () => { + const inner = makeTool( + "inner", + async () => ({ content: [{ type: "text", text: "inner ran" }], details: {} }) as AgentToolResult, + ); + let innerResult: AgentToolResult | undefined; + const outer = makeTool("outer", async (_id, _params, _signal, _onUpdate, ctx) => { + innerResult = await ctx?.invokeTool?.("inner", {}); + return { content: [{ type: "text", text: "outer ran" }], details: {} } as AgentToolResult; + }); + const store = new ToolContextStore(baseContext, name => + name === "outer" ? outer : name === "inner" ? inner : undefined, + ); + + await store.getContext().invokeTool?.("outer", {}); + + expect(innerResult?.content).toEqual([{ type: "text", text: "inner ran" }]); + }); + + it("guards against runaway recursion when a tool invokes itself", async () => { + // A pathological tool that always re-invokes its own native name. The depth guard must throw + // rather than recurse until the stack blows. + const loop = makeTool("loop", async (_id, _params, _signal, _onUpdate, ctx) => { + return (await ctx?.invokeTool?.("loop", {})) as AgentToolResult; + }); + const store = new ToolContextStore(baseContext, name => (name === "loop" ? loop : undefined)); + const invokeTool = store.getContext().invokeTool; + if (!invokeTool) throw new Error("invokeTool should be present when a resolver is set"); + + await expect(invokeTool("loop", {})).rejects.toThrow(/delegation depth exceeded/); + }); +}); From a50bcd5fed89e10c737e47c8d36f082fdba3587e Mon Sep 17 00:00:00 2001 From: Larry Gordon Date: Tue, 28 Jul 2026 11:05:02 -0700 Subject: [PATCH 2/5] 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. --- docs/extensions.md | 17 +-- packages/coding-agent/CHANGELOG.md | 2 +- .../src/extensibility/extensions/runner.ts | 76 ++++++++++++- .../src/extensibility/extensions/types.ts | 15 +++ .../src/extensibility/extensions/wrapper.ts | 10 +- packages/coding-agent/src/sdk.ts | 23 ++-- packages/coding-agent/src/tools/context.ts | 78 +------------ .../agent-session-message-pipeline.test.ts | 107 +++++++++++++++++- .../test/tool-context-invoke.test.ts | 98 ---------------- 9 files changed, 231 insertions(+), 195 deletions(-) delete mode 100644 packages/coding-agent/test/tool-context-invoke.test.ts diff --git a/docs/extensions.md b/docs/extensions.md index fd1e10954..b5d691f6f 100644 --- a/docs/extensions.md +++ b/docs/extensions.md @@ -315,21 +315,22 @@ execute( ### 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. The `ctx` passed to `execute` carries: +run the original instead of reimplementing it. When your registered tool shadows a built-in, the `ctx` +passed to `execute` carries: ```ts ctx.invokeTool?( - name: string, params: Record, options?: { signal?: AbortSignal; onUpdate?: AgentToolUpdateCallback }, -): Promise | undefined> +): Promise> ``` -It runs the **native** built-in of `name` (bypassing your own re-registration, so it does not recurse -into your wrapper) and returns its result, including the native tool's own side effects and internal -bookkeeping. It resolves to `undefined` when there is no native tool of that name. The invoked tool's -approval gate is not re-run — your call already passed approval — and delegation depth is guarded -against accidental self-recursion. +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: diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 0d18ada49..d44fb245f 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -158,7 +158,7 @@ - Added the bundled `ts-no-local-is-record` TTSR rule, which catches local `isRecord` function and lambda definitions and directs agents to shared guards plus explicit shape validation. - A `tool_call` handler (extension or hook) can now return `input` to revise the arguments a tool executes with, not just `block` it. The returned object is the raw execution input passed to the tool (ignored when `block` is set, and not applied to `computer` tool calls), enabling wrappers that normalize or rewrite a built-in's arguments without reimplementing the tool. For model-issued calls the event fires at arg-prep time in the agent loop, so a revision is revalidated against the tool schema and is what concurrency scheduling, `tool_execution_start`/transcripts, the persisted assistant message, and the approval gate all observe — the user approves exactly what runs, and a revision that changes a tool's functional concurrency (e.g. bash `pty`) schedules correctly. A revised nested `write xd://` device dispatch forfeits the outer write gate's approval and faces the full prompt again ([#6681](https://github.com/can1357/oh-my-pi/pull/6681) by [@psyrendust](https://github.com/psyrendust)). -- A tool's `execute` context now carries `ctx.invokeTool(name, params, options?)`, which runs the native built-in of `name` and returns its result. A tool that re-registers a built-in (e.g. wrapping `write`) can delegate to the original instead of reimplementing it: the native tool performs its own side effects and internal bookkeeping, the call bypasses the caller's own re-registration so it does not recurse, the invoked tool's approval gate is not re-run (the caller already passed approval), and delegation depth is guarded. Resolves to `undefined` when no native tool of that name exists. +- A re-registered built-in tool's `execute` context now carries `ctx.invokeTool(params, options?)`, which runs the native built-in of the same name and returns its result. A tool that wraps a built-in (e.g. re-registering `write` to add logging or a policy check) can delegate to the original instead of reimplementing it: the native tool performs its own side effects and internal bookkeeping. Delegation is same-tool only, so it cannot reach an arbitrary target or escalate past the approval already granted for the call, and the native call is not re-gated. It is present only when a native built-in of that name exists, and delegation depth is guarded against self-recursion. - Added a parser for macOS `sample`(1) call-tree reports to the read tool: `*.sample.txt` reads now return a compact bottleneck summary — per-thread hot paths with on-CPU sample counts (blocked syscall time excluded), demangled Rust v0/legacy symbols, flattened direct recursion, merged call-site siblings, idle-thread classification, and a process-wide top-functions-by-self-samples table. `:raw` still reads the original report, and files that merely carry the extension fall back to plain text. - Added V8 `.cpuprofile` support to the read tool (Node/Bun `--cpu-prof`, Chrome DevTools, CDP `Profiler.stop` output): reads now return a compact bottleneck summary — hot-path call tree with on-CPU milliseconds (`(idle)` time excluded), collapsed pass-through chains, flattened direct recursion, shortened file URLs, and a top-functions-by-self-time table. `:raw` still reads the original JSON, and files that merely carry the extension fall back to plain text. diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index 4b4354f87..e847733bb 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,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( + name: string, + params: Record, + options?: { signal?: AbortSignal; onUpdate?: AgentToolUpdateCallback; depth?: number }, + ): 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, + resolved.makeContext(), + )) as AgentToolResult; + } + 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, }; } 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..f2a4651d4 100644 --- a/packages/coding-agent/src/extensibility/extensions/wrapper.ts +++ b/packages/coding-agent/src/extensibility/extensions/wrapper.ts @@ -66,7 +66,15 @@ export class RegisteredToolAdapter implements AgentTool { onUpdate?: AgentToolUpdateCallback, _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), + ); } } diff --git a/packages/coding-agent/src/sdk.ts b/packages/coding-agent/src/sdk.ts index 618b3c477..da5ae6b21 100644 --- a/packages/coding-agent/src/sdk.ts +++ b/packages/coding-agent/src/sdk.ts @@ -2563,13 +2563,15 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} localProtocolOptions, autoApprove: options.autoApprove ?? false, }); - // Native built-in implementations, captured before extension re-registration replaces registry - // entries and before the ExtensionToolWrapper pass. Backs `ctx.invokeTool`, so a tool that - // wraps a built-in (e.g. a re-registered `write`) can delegate to the original — reaching the - // unwrapped native execute, which correctly inherits the caller's already-granted approval - // rather than re-running the gate. Populated at the built-in registry loop below. - const nativeToolsByName = new Map(); - const toolContextStore = new ToolContextStore(getSessionContext, name => nativeToolsByName.get(name)); + 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 = @@ -2612,6 +2614,13 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} 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)) { diff --git a/packages/coding-agent/src/tools/context.ts b/packages/coding-agent/src/tools/context.ts index 534f789e9..e8e5d1560 100644 --- a/packages/coding-agent/src/tools/context.ts +++ b/packages/coding-agent/src/tools/context.ts @@ -1,21 +1,7 @@ -import type { - AgentTool, - AgentToolContext, - AgentToolResult, - AgentToolUpdateCallback, - ToolCallContext, -} from "@oh-my-pi/pi-agent-core"; +import type { AgentToolContext, ToolCallContext } from "@oh-my-pi/pi-agent-core"; import type { CustomToolContext } from "../extensibility/custom-tools/types"; import type { ExtensionUIContext } from "../extensibility/extensions/types"; -/** Options a tool passes when delegating to another tool's native implementation via {@link AgentToolContext.invokeTool}. */ -export interface InvokeToolOptions { - /** Abort signal forwarded to the invoked tool's `execute`. */ - signal?: AbortSignal; - /** Progress callback forwarded to the invoked tool's `execute`. */ - onUpdate?: AgentToolUpdateCallback; -} - declare module "@oh-my-pi/pi-agent-core" { interface AgentToolContext extends CustomToolContext { ui?: ExtensionUIContext; @@ -29,41 +15,15 @@ declare module "@oh-my-pi/pi-agent-core" { xdevApproved?: boolean; /** Set only after an interactive prompt approves provider computer safety checks. */ providerSafetyApproved?: boolean; - /** - * Run the NATIVE built-in implementation of `name` with `params` and return its result, - * bypassing any extension re-registration of that name. Lets a tool that re-registers a - * built-in (e.g. wrapping `write`) delegate to the original instead of reimplementing it — - * the native tool performs its own side effects and internal bookkeeping. Resolves to - * `undefined` when no native tool of that name exists. The invoked tool's approval gate is - * NOT re-run: the caller is itself an already-approved tool call. Recursion is depth-guarded. - */ - invokeTool?( - name: string, - params: Record, - options?: InvokeToolOptions, - ): Promise | undefined>; } } -/** Max depth for `invokeTool` delegation chains — guards a re-registered tool that recurses into itself. */ -const MAX_INVOKE_DEPTH = 8; - export class ToolContextStore { #uiContext: ExtensionUIContext | undefined; #hasUI = false; #toolNames: string[] = []; - #invokeDepth = 0; - /** - * @param getBaseContext builds the per-call base tool context. - * @param resolveNativeTool resolves a tool NAME to its native built-in implementation (the - * pre-extension-override tool), or undefined if there is no native tool of that name. Used to - * back `ctx.invokeTool`. Lazy: called at invoke time, after the registry is fully assembled. - */ - constructor( - private readonly getBaseContext: () => CustomToolContext, - private readonly resolveNativeTool?: (name: string) => AgentTool | undefined, - ) {} + constructor(private readonly getBaseContext: () => CustomToolContext) {} getContext(toolCall?: ToolCallContext): AgentToolContext { return { @@ -72,43 +32,9 @@ export class ToolContextStore { hasUI: this.#hasUI, toolNames: this.#toolNames, toolCall, - invokeTool: this.resolveNativeTool - ? (name, params, options) => this.#invokeTool(name, params, options) - : undefined, }; } - async #invokeTool( - name: string, - params: Record, - options?: InvokeToolOptions, - ): Promise | undefined> { - const native = this.resolveNativeTool?.(name); - if (!native) return undefined; - if (this.#invokeDepth >= MAX_INVOKE_DEPTH) { - throw new Error( - `invokeTool: delegation depth exceeded ${MAX_INVOKE_DEPTH} (recursive invokeTool for "${name}"?)`, - ); - } - // Nested context so the invoked tool sees the same session state (ui, cwd, etc.) and can itself - // delegate. Its approval gate is NOT re-run: it is reached via the native execute directly, not - // the ExtensionToolWrapper, so the caller's already-granted approval covers it. - const nestedContext = this.getContext(undefined); - this.#invokeDepth++; - try { - const toolCallId = `invoke-${name}-${Date.now().toString(36)}-${this.#invokeDepth}`; - return (await native.execute( - toolCallId, - params as never, - options?.signal, - options?.onUpdate as never, - nestedContext, - )) as AgentToolResult; - } finally { - this.#invokeDepth--; - } - } - setUIContext(uiContext: ExtensionUIContext, hasUI: boolean): void { this.#uiContext = uiContext; this.#hasUI = hasUI; 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/tool-context-invoke.test.ts b/packages/coding-agent/test/tool-context-invoke.test.ts deleted file mode 100644 index 597c50956..000000000 --- a/packages/coding-agent/test/tool-context-invoke.test.ts +++ /dev/null @@ -1,98 +0,0 @@ -/** - * Tests for `ctx.invokeTool` — a tool delegating to another tool's native built-in implementation - * via the ToolContextStore-built AgentToolContext. - */ - -import { describe, expect, it } from "bun:test"; -import type { AgentTool, AgentToolResult } from "@oh-my-pi/pi-agent-core"; -import { Type } from "@oh-my-pi/pi-coding-agent/extensibility/typebox"; -import { ToolContextStore } from "@oh-my-pi/pi-coding-agent/tools/context"; - -const baseContext = () => ({}) as never; - -function makeTool(name: string, exec: AgentTool["execute"]): AgentTool { - return { - name, - label: name, - description: `${name} tool`, - parameters: Type.Object({}), - execute: exec, - } as AgentTool; -} - -describe("ctx.invokeTool native-tool delegation", () => { - it("is absent when the store has no native-tool resolver", () => { - const store = new ToolContextStore(baseContext); - expect(store.getContext().invokeTool).toBeUndefined(); - }); - - it("runs the native tool and returns its result", async () => { - const seen: unknown[] = []; - const native = makeTool("write", async (_id, params) => { - seen.push(params); - return { content: [{ type: "text", text: "native wrote" }], details: { ok: true } } as AgentToolResult; - }); - const store = new ToolContextStore(baseContext, name => (name === "write" ? native : undefined)); - - const result = await store.getContext().invokeTool?.("write", { path: "/tmp/x", content: "hi" }); - - expect(seen).toEqual([{ path: "/tmp/x", content: "hi" }]); - expect(result?.content).toEqual([{ type: "text", text: "native wrote" }]); - expect(result?.details).toEqual({ ok: true }); - }); - - it("resolves to undefined for an unknown tool name", async () => { - const store = new ToolContextStore(baseContext, () => undefined); - expect(await store.getContext().invokeTool?.("nope", {})).toBeUndefined(); - }); - - it("forwards signal and onUpdate to the native tool", async () => { - const controller = new AbortController(); - let receivedSignal: AbortSignal | undefined; - let receivedOnUpdate: unknown; - const native = makeTool("read", async (_id, _params, signal, onUpdate) => { - receivedSignal = signal; - receivedOnUpdate = onUpdate; - return { content: [{ type: "text", text: "ok" }], details: {} } as AgentToolResult; - }); - const store = new ToolContextStore(baseContext, () => native); - const onUpdate = () => {}; - - await store.getContext().invokeTool?.("read", {}, { signal: controller.signal, onUpdate }); - - expect(receivedSignal).toBe(controller.signal); - expect(receivedOnUpdate).toBe(onUpdate); - }); - - it("passes a context whose invokeTool lets the native tool delegate further", async () => { - const inner = makeTool( - "inner", - async () => ({ content: [{ type: "text", text: "inner ran" }], details: {} }) as AgentToolResult, - ); - let innerResult: AgentToolResult | undefined; - const outer = makeTool("outer", async (_id, _params, _signal, _onUpdate, ctx) => { - innerResult = await ctx?.invokeTool?.("inner", {}); - return { content: [{ type: "text", text: "outer ran" }], details: {} } as AgentToolResult; - }); - const store = new ToolContextStore(baseContext, name => - name === "outer" ? outer : name === "inner" ? inner : undefined, - ); - - await store.getContext().invokeTool?.("outer", {}); - - expect(innerResult?.content).toEqual([{ type: "text", text: "inner ran" }]); - }); - - it("guards against runaway recursion when a tool invokes itself", async () => { - // A pathological tool that always re-invokes its own native name. The depth guard must throw - // rather than recurse until the stack blows. - const loop = makeTool("loop", async (_id, _params, _signal, _onUpdate, ctx) => { - return (await ctx?.invokeTool?.("loop", {})) as AgentToolResult; - }); - const store = new ToolContextStore(baseContext, name => (name === "loop" ? loop : undefined)); - const invokeTool = store.getContext().invokeTool; - if (!invokeTool) throw new Error("invokeTool should be present when a resolver is set"); - - await expect(invokeTool("loop", {})).rejects.toThrow(/delegation depth exceeded/); - }); -}); From efaf4f5d495da0674cbaf06a7782dd4084a765fe Mon Sep 17 00:00:00 2001 From: Larry Gordon Date: Tue, 28 Jul 2026 14:04:50 -0700 Subject: [PATCH 3/5] fix(extensions): preserve the caller tool context when delegating via invokeTool The delegated native call built a fresh AgentToolContext with no toolCall metadata. Native tools read that: write/edit derive LSP batch flushing from context.toolCall, and computer uses provider metadata plus providerSafetyApproved for the required screenshot/safety acknowledgement, so wrapping those tools lost batching and dropped provider result metadata. Thread the caller's own context (the one the re-registered tool received) through RegisteredToolAdapter.execute and createContext into invokeNativeTool, and reuse it for the native call instead of a bare getContext(), falling back to a fresh session context only when the caller had none. --- .../src/extensibility/extensions/runner.ts | 22 +++++++++++++++---- .../src/extensibility/extensions/wrapper.ts | 8 ++++--- 2 files changed, 23 insertions(+), 7 deletions(-) diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index e847733bb..a48e7ab43 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -445,7 +445,18 @@ export class ExtensionRunner { async invokeNativeTool( name: string, params: Record, - options?: { signal?: AbortSignal; onUpdate?: AgentToolUpdateCallback; depth?: number }, + 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`); @@ -459,7 +470,7 @@ export class ExtensionRunner { params as never, options?.signal, options?.onUpdate as never, - resolved.makeContext(), + options?.callerContext ?? resolved.makeContext(), )) as AgentToolResult; } @@ -788,9 +799,11 @@ export class ExtensionRunner { * 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. + * per-chain recursion counter so a wrapper that re-invokes itself is bounded, and `callerContext` + * (the context the re-registered tool itself received) is reused for the native call so its + * `toolCall`/provider metadata is preserved. */ - createContext(model?: Model, toolName?: string, depth = 0): ExtensionContext { + createContext(model?: Model, toolName?: string, depth = 0, callerContext?: AgentToolContext): ExtensionContext { const getModel = model ? () => model : this.#getModel; return { ui: this.#uiContext, @@ -822,6 +835,7 @@ export class ExtensionRunner { signal: options?.signal, onUpdate: options?.onUpdate, depth: depth + 1, + callerContext, }) : undefined, }; diff --git a/packages/coding-agent/src/extensibility/extensions/wrapper.ts b/packages/coding-agent/src/extensibility/extensions/wrapper.ts index f2a4651d4..b2b2d98dd 100644 --- a/packages/coding-agent/src/extensibility/extensions/wrapper.ts +++ b/packages/coding-agent/src/extensibility/extensions/wrapper.ts @@ -64,16 +64,18 @@ export class RegisteredToolAdapter implements AgentTool { params: any, signal?: AbortSignal, onUpdate?: AgentToolUpdateCallback, - _context?: AgentToolContext, + context?: AgentToolContext, ) { // 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). + // native built-in of the same name (present only when this tool re-registers a built-in). The + // incoming loop context is threaded through so a delegated native call keeps the caller's + // `toolCall`/provider metadata (write/edit LSP batching, computer safety acknowledgement). return this.registeredTool.definition.execute( toolCallId, params, signal, onUpdate, - this.runner.createContext(undefined, this.registeredTool.definition.name), + this.runner.createContext(undefined, this.registeredTool.definition.name, 0, context), ); } } From e8d9f15e9d191b17031e20df1522df8c5b1cbb07 Mon Sep 17 00:00:00 2001 From: Larry Gordon Date: Tue, 28 Jul 2026 15:33:29 -0700 Subject: [PATCH 4/5] fix(extensions): inherit the wrapper call's abort and progress channels A bare ctx.invokeTool(params) passed undefined for both signal and onUpdate, so a wrapper that simply delegates did not stop the native tool when the outer call was aborted, and native progress updates were dropped unless every wrapper forwarded them by hand. createContext now takes the delegation wiring as one named object and binds the wrapper's own signal and onUpdate as defaults for the delegated call, with explicit invokeTool options still taking precedence. Grouping toolName, depth, context, signal, and onUpdate together also keeps the signature readable now that delegation carries five inputs. Tests: the delegated native call receives the outer signal and onUpdate, explicit options override them, invokeTool is absent when no native built-in of that name exists, and recursion stays bounded per call chain. --- .../src/extensibility/extensions/runner.ts | 40 ++++++--- .../src/extensibility/extensions/wrapper.ts | 13 ++- .../test/extensions-runner.test.ts | 87 +++++++++++++++++++ 3 files changed, 124 insertions(+), 16 deletions(-) diff --git a/packages/coding-agent/src/extensibility/extensions/runner.ts b/packages/coding-agent/src/extensibility/extensions/runner.ts index a48e7ab43..a32af1f49 100644 --- a/packages/coding-agent/src/extensibility/extensions/runner.ts +++ b/packages/coding-agent/src/extensibility/extensions/runner.ts @@ -796,14 +796,26 @@ export class ExtensionRunner { } /** - * 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, and `callerContext` - * (the context the re-registered tool itself received) is reused for the native call so its - * `toolCall`/provider metadata is preserved. + * 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, toolName?: string, depth = 0, callerContext?: AgentToolContext): ExtensionContext { + 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, @@ -829,13 +841,15 @@ export class ExtensionRunner { setTimeout: (callback, ms, ...args) => this.#managedTimers.setTimeout(callback, ms, ...args), clearTimer: timer => this.#managedTimers.clear(timer), invokeTool: - toolName !== undefined && this.hasNativeTool(toolName) + delegation !== undefined && this.hasNativeTool(delegation.toolName) ? (params, options) => - this.invokeNativeTool(toolName, params, { - signal: options?.signal, - onUpdate: options?.onUpdate, - depth: depth + 1, - callerContext, + 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/wrapper.ts b/packages/coding-agent/src/extensibility/extensions/wrapper.ts index b2b2d98dd..8ce2305f2 100644 --- a/packages/coding-agent/src/extensibility/extensions/wrapper.ts +++ b/packages/coding-agent/src/extensibility/extensions/wrapper.ts @@ -68,14 +68,21 @@ export class RegisteredToolAdapter implements AgentTool { ) { // 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 - // incoming loop context is threaded through so a delegated native call keeps the caller's - // `toolCall`/provider metadata (write/edit LSP batching, computer safety acknowledgement). + // 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, this.registeredTool.definition.name, 0, context), + this.runner.createContext(undefined, { + toolName: this.registeredTool.definition.name, + context, + signal, + onUpdate, + }), ); } } diff --git a/packages/coding-agent/test/extensions-runner.test.ts b/packages/coding-agent/test/extensions-runner.test.ts index b25dd020f..ead93b19b 100644 --- a/packages/coding-agent/test/extensions-runner.test.ts +++ b/packages/coding-agent/test/extensions-runner.test.ts @@ -3107,4 +3107,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(); + }); + }); }); From 5352d5b51b50c92ee6858676cef97087acd492a2 Mon Sep 17 00:00:00 2001 From: Larry Gordon Date: Tue, 28 Jul 2026 15:44:41 -0700 Subject: [PATCH 5/5] docs(changelog): add the invokeTool entry under Unreleased Upstream rewrote the package changelogs in 93eb95b3b, which condensed every entry and reset Unreleased, so the entry this branch carried no longer had a home and its earlier placement had drifted into a released section. Restated it as a single Added line under Unreleased, matching the shorter house style the rewrite established. --- packages/coding-agent/CHANGELOG.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index d44fb245f..c22870628 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -2,6 +2,10 @@ ## [Unreleased] +### Added + +- 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 ### Added @@ -158,7 +162,6 @@ - Added the bundled `ts-no-local-is-record` TTSR rule, which catches local `isRecord` function and lambda definitions and directs agents to shared guards plus explicit shape validation. - A `tool_call` handler (extension or hook) can now return `input` to revise the arguments a tool executes with, not just `block` it. The returned object is the raw execution input passed to the tool (ignored when `block` is set, and not applied to `computer` tool calls), enabling wrappers that normalize or rewrite a built-in's arguments without reimplementing the tool. For model-issued calls the event fires at arg-prep time in the agent loop, so a revision is revalidated against the tool schema and is what concurrency scheduling, `tool_execution_start`/transcripts, the persisted assistant message, and the approval gate all observe — the user approves exactly what runs, and a revision that changes a tool's functional concurrency (e.g. bash `pty`) schedules correctly. A revised nested `write xd://` device dispatch forfeits the outer write gate's approval and faces the full prompt again ([#6681](https://github.com/can1357/oh-my-pi/pull/6681) by [@psyrendust](https://github.com/psyrendust)). -- A re-registered built-in tool's `execute` context now carries `ctx.invokeTool(params, options?)`, which runs the native built-in of the same name and returns its result. A tool that wraps a built-in (e.g. re-registering `write` to add logging or a policy check) can delegate to the original instead of reimplementing it: the native tool performs its own side effects and internal bookkeeping. Delegation is same-tool only, so it cannot reach an arbitrary target or escalate past the approval already granted for the call, and the native call is not re-gated. It is present only when a native built-in of that name exists, and delegation depth is guarded against self-recursion. - Added a parser for macOS `sample`(1) call-tree reports to the read tool: `*.sample.txt` reads now return a compact bottleneck summary — per-thread hot paths with on-CPU sample counts (blocked syscall time excluded), demangled Rust v0/legacy symbols, flattened direct recursion, merged call-site siblings, idle-thread classification, and a process-wide top-functions-by-self-samples table. `:raw` still reads the original report, and files that merely carry the extension fall back to plain text. - Added V8 `.cpuprofile` support to the read tool (Node/Bun `--cpu-prof`, Chrome DevTools, CDP `Profiler.stop` output): reads now return a compact bottleneck summary — hot-path call tree with on-CPU milliseconds (`(idle)` time excluded), collapsed pass-through chains, flattened direct recursion, shortened file URLs, and a top-functions-by-self-time table. `:raw` still reads the original JSON, and files that merely carry the extension fall back to plain text.