From a50bcd5fed89e10c737e47c8d36f082fdba3587e Mon Sep 17 00:00:00 2001 From: Larry Gordon Date: Tue, 28 Jul 2026 11:05:02 -0700 Subject: [PATCH] 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/); - }); -});