Added nono profiles & orca integration - Added some extensions to start working on

extensions/shared/activity-status.ts

@@ -0,0 +1,31 @@
+import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
+
+type Theme = ExtensionContext["ui"]["theme"];
+
+interface ActivityCounts {
+  running: number;
+  done: number;
+  failed: number;
+}
+
+const SQUARE = "■";
+
+export function formatActivityStatus(
+  theme: Theme,
+  label: "subagents" | "workflows",
+  counts: ActivityCounts,
+) {
+  const parts: string[] = [];
+  if (counts.running > 0) {
+    parts.push(theme.fg("warning", `${SQUARE} ${counts.running} running`));
+  }
+  if (counts.done > 0) {
+    parts.push(theme.fg("success", `${SQUARE} ${counts.done} done`));
+  }
+  if (counts.failed > 0) {
+    parts.push(theme.fg("error", `${SQUARE} ${counts.failed} failed`));
+  }
+  parts.push(theme.fg("accent", `/${label}`) + theme.fg("dim", " to view"));
+
+  return `${theme.fg("muted", `${label}:`)} ${parts.join(theme.fg("dim", " · "))}`;
+}

extensions/shared/child-session.test.ts

@@ -0,0 +1,270 @@
+import assert from "node:assert/strict";
+import { mkdir, mkdtemp, rm, writeFile } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import * as path from "node:path";
+import test from "node:test";
+import {
+  createAgentSession,
+  DefaultResourceLoader,
+  defineTool,
+  ProjectTrustStore,
+  SessionManager,
+  SettingsManager,
+  type SessionShutdownEvent,
+} from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+import {
+  bindChildSessionExtensions,
+  CHILD_EXCLUDED_TOOL_NAMES,
+  childToolPolicy,
+  createChildResources,
+  resolveStandaloneChildProjectTrust,
+  shutdownAndDisposeChildSession,
+  type DisposableChildSession,
+} from "./child-session.ts";
+
+async function withTempDir(run: (directory: string) => Promise<void>) {
+  const directory = await mkdtemp(path.join(tmpdir(), "pi-child-policy-"));
+  try {
+    await run(directory);
+  } finally {
+    await rm(directory, { recursive: true, force: true });
+  }
+}
+
+test("child denylist keeps extension and workflow structured tools available", async () => {
+  await withTempDir(async (directory) => {
+    let starts = 0;
+    let shutdowns = 0;
+    const settingsManager = SettingsManager.inMemory(undefined, {
+      projectTrusted: false,
+    });
+    const inlineLoader = new DefaultResourceLoader({
+      cwd: directory,
+      agentDir: path.join(directory, "inline-agent"),
+      settingsManager,
+      extensionFactories: [
+        (pi) => {
+          pi.on("session_start", () => {
+            starts++;
+          });
+          pi.on("session_shutdown", () => {
+            shutdowns++;
+          });
+          for (const name of [
+            "fixture_extension_tool",
+            ...CHILD_EXCLUDED_TOOL_NAMES,
+          ]) {
+            pi.registerTool({
+              name,
+              label: name,
+              description: name,
+              parameters: Type.Object({}),
+              async execute() {
+                return {
+                  content: [{ type: "text", text: "ok" }],
+                  details: {},
+                };
+              },
+            });
+          }
+        },
+      ],
+    });
+    await inlineLoader.reload();
+
+    const structuredOutput = defineTool({
+      name: "structured_output",
+      label: "Structured Output",
+      description: "fixture structured result",
+      parameters: Type.Object({ value: Type.String() }),
+      async execute(_id, params) {
+        return {
+          content: [{ type: "text", text: params.value }],
+          details: {},
+        };
+      },
+    });
+    const { session } = await createAgentSession({
+      cwd: directory,
+      agentDir: path.join(directory, "inline-agent"),
+      resourceLoader: inlineLoader,
+      settingsManager,
+      sessionManager: SessionManager.inMemory(directory),
+      customTools: [structuredOutput],
+      ...childToolPolicy(),
+    });
+    await bindChildSessionExtensions(session);
+
+    assert.deepEqual(
+      [...CHILD_EXCLUDED_TOOL_NAMES],
+      [
+        "subagent_spawn",
+        "subagent_wait",
+        "subagent_cancel",
+        "subagent_check",
+        "subagent_list",
+        "workflow",
+        "ask_user",
+      ],
+    );
+    const allTools = new Set(session.getAllTools().map((tool) => tool.name));
+    const activeTools = new Set(session.getActiveToolNames());
+    assert.equal(starts, 1);
+    assert.equal(allTools.has("fixture_extension_tool"), true);
+    assert.equal(activeTools.has("fixture_extension_tool"), true);
+    assert.equal(allTools.has("structured_output"), true);
+    assert.equal(activeTools.has("structured_output"), true);
+    for (const denied of CHILD_EXCLUDED_TOOL_NAMES) {
+      assert.equal(allTools.has(denied), false, `${denied} should be denied`);
+      assert.equal(
+        activeTools.has(denied),
+        false,
+        `${denied} should be inactive`,
+      );
+    }
+    for (const builtin of ["read", "bash", "edit", "write"]) {
+      assert.equal(
+        activeTools.has(builtin),
+        true,
+        `${builtin} should stay active`,
+      );
+    }
+
+    await Promise.all([
+      shutdownAndDisposeChildSession(session),
+      shutdownAndDisposeChildSession(session),
+    ]);
+    assert.equal(shutdowns, 1);
+  });
+});
+
+test("resource loading gates project extensions but retains global extensions", async () => {
+  await withTempDir(async (directory) => {
+    const cwd = path.join(directory, "project");
+    const agentDir = path.join(directory, "agent");
+    await mkdir(path.join(cwd, ".pi", "extensions"), { recursive: true });
+    await mkdir(path.join(agentDir, "extensions"), { recursive: true });
+    const extensionSource = (name: string) => `
+      export default function (pi) {
+        pi.registerTool({
+          name: ${JSON.stringify(name)}, label: ${JSON.stringify(name)},
+          description: "fixture", parameters: { type: "object", properties: {} },
+          async execute() { return { content: [{ type: "text", text: "ok" }] }; }
+        });
+      }
+    `;
+    await writeFile(
+      path.join(agentDir, "extensions", "global.ts"),
+      extensionSource("global_fixture"),
+    );
+    await writeFile(
+      path.join(cwd, ".pi", "extensions", "project.ts"),
+      extensionSource("project_fixture"),
+    );
+
+    const untrusted = await createChildResources({
+      cwd,
+      agentDir,
+      projectTrusted: false,
+    });
+    const untrustedTools = untrusted.loader
+      .getExtensions()
+      .extensions.flatMap((extension) => [...extension.tools.keys()]);
+    assert.equal(untrustedTools.includes("global_fixture"), true);
+    assert.equal(untrustedTools.includes("project_fixture"), false);
+
+    const trusted = await createChildResources({
+      cwd,
+      agentDir,
+      projectTrusted: true,
+    });
+    const trustedTools = trusted.loader
+      .getExtensions()
+      .extensions.flatMap((extension) => [...extension.tools.keys()]);
+    assert.equal(trustedTools.includes("global_fixture"), true);
+    assert.equal(trustedTools.includes("project_fixture"), true);
+  });
+});
+
+test("alternate standalone cwd only uses explicit saved trust", async () => {
+  await withTempDir(async (directory) => {
+    const parentCwd = path.join(directory, "parent");
+    const childCwd = path.join(directory, "alternate");
+    const agentDir = path.join(directory, "agent");
+    await mkdir(parentCwd, { recursive: true });
+    await mkdir(childCwd, { recursive: true });
+
+    assert.equal(
+      resolveStandaloneChildProjectTrust({
+        parentCwd,
+        childCwd: parentCwd,
+        parentTrusted: true,
+        agentDir,
+      }),
+      true,
+    );
+    assert.equal(
+      resolveStandaloneChildProjectTrust({
+        parentCwd,
+        childCwd,
+        parentTrusted: true,
+        agentDir,
+      }),
+      false,
+    );
+
+    new ProjectTrustStore(agentDir).set(childCwd, true);
+    assert.equal(
+      resolveStandaloneChildProjectTrust({
+        parentCwd,
+        childCwd,
+        parentTrusted: false,
+        agentDir,
+      }),
+      true,
+    );
+  });
+});
+
+test("shutdown helper balances hooks and disposal despite errors", async () => {
+  let emits = 0;
+  let disposals = 0;
+  const session: DisposableChildSession = {
+    extensionRunner: {
+      hasHandlers: () => true,
+      async emit(event: SessionShutdownEvent) {
+        emits++;
+        assert.deepEqual(event, { type: "session_shutdown", reason: "quit" });
+        throw new Error("fixture shutdown failure");
+      },
+    },
+    dispose() {
+      disposals++;
+    },
+  };
+
+  await Promise.all([
+    shutdownAndDisposeChildSession(session),
+    shutdownAndDisposeChildSession(session),
+    shutdownAndDisposeChildSession(session),
+  ]);
+  assert.equal(emits, 1);
+  assert.equal(disposals, 1);
+});
+
+test("shutdown helper bounds a stuck hook before disposal", async () => {
+  let disposals = 0;
+  const session: DisposableChildSession = {
+    extensionRunner: {
+      hasHandlers: () => true,
+      emit: () => new Promise(() => {}),
+    },
+    dispose() {
+      disposals++;
+    },
+  };
+
+  await shutdownAndDisposeChildSession(session, { timeoutMs: 10 });
+  assert.equal(disposals, 1);
+});

extensions/shared/child-session.ts

@@ -0,0 +1,148 @@
+import * as path from "node:path";
+import {
+  DefaultResourceLoader,
+  getAgentDir,
+  ProjectTrustStore,
+  SettingsManager,
+  type AgentSession,
+  type SessionShutdownEvent,
+} from "@earendil-works/pi-coding-agent";
+
+const CHILD_SHUTDOWN_TIMEOUT_MS = 5_000;
+
+/** Tools that headless children must not receive. Everything else stays enabled. */
+export const CHILD_EXCLUDED_TOOL_NAMES = [
+  "subagent_spawn",
+  "subagent_wait",
+  "subagent_cancel",
+  "subagent_check",
+  "subagent_list",
+  "workflow",
+  "ask_user",
+] as const;
+
+/** Fresh SDK options avoid turning the denylist into an accidental allowlist. */
+export function childToolPolicy() {
+  return { excludeTools: [...CHILD_EXCLUDED_TOOL_NAMES] };
+}
+
+export interface ChildResourceOptions {
+  cwd: string;
+  projectTrusted: boolean;
+  appendSystemPrompt?: string[];
+  agentDir?: string;
+}
+
+/** Load normal global/package resources and trust-gated project resources. */
+export async function createChildResources(options: ChildResourceOptions) {
+  const agentDir = options.agentDir ?? getAgentDir();
+  const settingsManager = SettingsManager.create(options.cwd, agentDir, {
+    projectTrusted: options.projectTrusted,
+  });
+  const loader = new DefaultResourceLoader({
+    cwd: options.cwd,
+    agentDir,
+    settingsManager,
+    ...(options.appendSystemPrompt
+      ? { appendSystemPrompt: options.appendSystemPrompt }
+      : {}),
+  });
+  await loader.reload();
+  return { loader, settingsManager };
+}
+
+/**
+ * Same-directory children inherit the live parent decision. An alternate cwd
+ * is trusted only when Pi's persisted trust store explicitly trusts it (or a
+ * containing directory); unreadable/invalid trust data fails closed.
+ */
+export function resolveStandaloneChildProjectTrust(options: {
+  parentCwd: string;
+  childCwd: string;
+  parentTrusted: boolean;
+  agentDir?: string;
+}) {
+  if (path.resolve(options.childCwd) === path.resolve(options.parentCwd)) {
+    return options.parentTrusted;
+  }
+  try {
+    const trustStore = new ProjectTrustStore(options.agentDir ?? getAgentDir());
+    return trustStore.get(options.childCwd) === true;
+  } catch {
+    return false;
+  }
+}
+
+/** Start child extension session hooks/resources in headless print mode. */
+export async function bindChildSessionExtensions(
+  session: Pick<AgentSession, "bindExtensions">,
+) {
+  await session.bindExtensions({ mode: "print" });
+}
+
+interface ChildExtensionRunner {
+  hasHandlers(eventType: string): boolean;
+  emit(event: SessionShutdownEvent): Promise<unknown>;
+}
+
+export interface DisposableChildSession {
+  readonly extensionRunner: ChildExtensionRunner;
+  dispose(): void;
+}
+
+const childShutdowns = new WeakMap<object, Promise<void>>();
+
+function waitBounded(operation: Promise<unknown>, timeoutMs: number) {
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  const timeout = new Promise<void>((resolve) => {
+    timer = setTimeout(resolve, timeoutMs);
+  });
+  return Promise.race([
+    operation.then(
+      () => undefined,
+      () => undefined,
+    ),
+    timeout,
+  ])
+    .catch(() => {})
+    .finally(() => {
+      if (timer) clearTimeout(timer);
+    });
+}
+
+/**
+ * Emit child session_shutdown once, then dispose once. Hook failures and a
+ * bounded hook deadline never prevent disposal.
+ */
+export function shutdownAndDisposeChildSession(
+  session: DisposableChildSession,
+  options: { timeoutMs?: number } = {},
+) {
+  const existing = childShutdowns.get(session);
+  if (existing) return existing;
+
+  const shutdown = (async () => {
+    try {
+      if (session.extensionRunner.hasHandlers("session_shutdown")) {
+        await waitBounded(
+          session.extensionRunner.emit({
+            type: "session_shutdown",
+            reason: "quit",
+          }),
+          options.timeoutMs ?? CHILD_SHUTDOWN_TIMEOUT_MS,
+        );
+      }
+    } catch {
+      // Extension runner inspection/emission is best-effort during teardown.
+    } finally {
+      try {
+        session.dispose();
+      } catch {
+        // Disposal is terminal and must remain idempotent for callers.
+      }
+    }
+  })();
+
+  childShutdowns.set(session, shutdown);
+  return shutdown;
+}

extensions/shared/context-utilization.test.ts

@@ -0,0 +1,48 @@
+import assert from "node:assert/strict";
+import { test } from "node:test";
+import {
+  contextPercent,
+  formatContextUtilization,
+} from "./context-utilization.ts";
+
+test("formats current context occupancy against model capacity", () => {
+  assert.equal(
+    formatContextUtilization({ tokens: 26_040, contextWindow: 372_000 }),
+    "7%/372k",
+  );
+});
+
+test("formats the latest post-compaction usage rather than prior cumulative usage", () => {
+  const latestUsage = { tokens: 18_000, contextWindow: 200_000 };
+  assert.equal(formatContextUtilization(latestUsage), "9%/200k");
+});
+
+test("clamps over-capacity and nonsensical token values", () => {
+  assert.equal(
+    contextPercent({ tokens: 500_000, contextWindow: 200_000 }),
+    100,
+  );
+  assert.equal(
+    formatContextUtilization({
+      tokens: Number.POSITIVE_INFINITY,
+      contextWindow: 200_000,
+    }),
+    "?%/200k",
+  );
+  assert.equal(
+    formatContextUtilization({ tokens: -1, contextWindow: 200_000 }),
+    "?%/200k",
+  );
+});
+
+test("handles missing usage or capacity without NaN or Infinity", () => {
+  assert.equal(
+    formatContextUtilization({ tokens: null, contextWindow: 372_000 }),
+    "?%/372k",
+  );
+  assert.equal(formatContextUtilization({ tokens: 12_000 }), "");
+  assert.equal(
+    formatContextUtilization({ tokens: 12_000, contextWindow: 0 }),
+    "",
+  );
+});

extensions/shared/context-utilization.ts

@@ -0,0 +1,47 @@
+/** Compact, defensive context-window utilization formatting for child agents. */
+
+export interface ContextUtilization {
+  /** Current conversation context occupancy; null while it is unknown after compaction. */
+  tokens?: number | null;
+  /** Capacity of the model currently serving the conversation. */
+  contextWindow?: number | null;
+}
+
+function usableTokens(value: number | null | undefined) {
+  return typeof value === "number" && Number.isFinite(value) && value >= 0
+    ? value
+    : undefined;
+}
+
+function usableCapacity(value: number | null | undefined) {
+  return typeof value === "number" && Number.isFinite(value) && value > 0
+    ? value
+    : undefined;
+}
+
+export function contextPercent(usage: ContextUtilization) {
+  const tokens = usableTokens(usage.tokens);
+  const capacity = usableCapacity(usage.contextWindow);
+  if (tokens === undefined || capacity === undefined) return undefined;
+  return Math.round(Math.min(100, Math.max(0, (tokens / capacity) * 100)));
+}
+
+export function formatCompactTokens(count: number) {
+  if (count < 1000) return Math.round(count).toString();
+  if (count < 10000) return `${(count / 1000).toFixed(1)}k`;
+  if (count < 1000000) return `${Math.round(count / 1000)}k`;
+  return `${(count / 1000000).toFixed(1)}M`;
+}
+
+/**
+ * Render `%/capacity`. If occupancy is temporarily unknown (notably directly
+ * after compaction), retain the useful model capacity as `?%/capacity`. If no
+ * valid capacity is available, omit the statistic because no percentage can
+ * be computed.
+ */
+export function formatContextUtilization(usage: ContextUtilization) {
+  const capacity = usableCapacity(usage.contextWindow);
+  if (capacity === undefined) return "";
+  const percent = contextPercent(usage);
+  return `${percent === undefined ? "?" : percent}%/${formatCompactTokens(capacity)}`;
+}

extensions/shared/tool-call-timeout.test.ts

@@ -0,0 +1,117 @@
+import assert from "node:assert/strict";
+import test from "node:test";
+import type { ToolDefinition } from "@earendil-works/pi-coding-agent";
+import { Type } from "typebox";
+import {
+  createToolCallTimeoutGuard,
+  runWithToolCallTimeout,
+  CHILD_TOOL_CALL_TIMEOUT_MS,
+  ToolCallTimeoutError,
+} from "./tool-call-timeout.ts";
+
+test("the production timeout error names the tool and three-minute limit", () => {
+  assert.equal(
+    new ToolCallTimeoutError("fixture_tool", CHILD_TOOL_CALL_TIMEOUT_MS)
+      .message,
+    'Tool call "fixture_tool" timed out after 3 minutes.',
+  );
+});
+
+test("a hung tool call fails clearly and receives an abort signal", async () => {
+  let executionSignal: AbortSignal | undefined;
+
+  await assert.rejects(
+    runWithToolCallTimeout("hung_fixture", 10, undefined, (signal) => {
+      executionSignal = signal;
+      return new Promise(() => {});
+    }),
+    (error: unknown) => {
+      assert.equal(error instanceof ToolCallTimeoutError, true);
+      assert.equal(
+        error instanceof Error ? error.message : "",
+        'Tool call "hung_fixture" timed out after 10 ms.',
+      );
+      return true;
+    },
+  );
+
+  assert.equal(executionSignal?.aborted, true);
+  assert.equal(executionSignal?.reason instanceof ToolCallTimeoutError, true);
+});
+
+test("parent cancellation still stops the timeout wrapper immediately", async () => {
+  const controller = new AbortController();
+  const reason = new Error("cancelled fixture");
+  const pending = runWithToolCallTimeout(
+    "hung_fixture",
+    60_000,
+    controller.signal,
+    () => new Promise(() => {}),
+  );
+
+  controller.abort(reason);
+
+  await assert.rejects(pending, (error: unknown) => error === reason);
+});
+
+test("the guard wraps each definition once and can discover later tools", () => {
+  const definitions = new Map<string, ToolDefinition>();
+  const createDefinition = (name: string): ToolDefinition => ({
+    name,
+    label: name,
+    description: name,
+    parameters: Type.Object({}),
+    async execute() {
+      return { content: [{ type: "text", text: "done" }], details: {} };
+    },
+  });
+  const first = createDefinition("first");
+  definitions.set(first.name, first);
+  const registry = {
+    getAllTools: () => [...definitions.keys()].map((name) => ({ name })),
+    getToolDefinition: (name: string) => definitions.get(name),
+  };
+  const guard = createToolCallTimeoutGuard(10);
+
+  const firstExecute = first.execute;
+  guard.apply(registry);
+  const firstWrappedExecute = first.execute;
+  assert.notEqual(firstWrappedExecute, firstExecute);
+
+  const second = createDefinition("second");
+  const secondExecute = second.execute;
+  definitions.set(second.name, second);
+  guard.apply(registry);
+
+  assert.equal(first.execute, firstWrappedExecute);
+  assert.notEqual(second.execute, secondExecute);
+});
+
+test("successful and terminating tool results pass through unchanged", async () => {
+  const result = {
+    content: [{ type: "text" as const, text: "recorded" }],
+    details: { value: "fixture" },
+    terminate: true,
+  };
+
+  assert.equal(
+    await runWithToolCallTimeout(
+      "structured_output",
+      10,
+      undefined,
+      async () => result,
+    ),
+    result,
+  );
+});
+
+test("the timeout is fresh for each tool call, not shared across calls", async () => {
+  const execute = () =>
+    runWithToolCallTimeout("slow_fixture", 100, undefined, async () => {
+      await new Promise((resolve) => setTimeout(resolve, 60));
+      return "done";
+    });
+
+  assert.equal(await execute(), "done");
+  assert.equal(await execute(), "done");
+});

extensions/shared/tool-call-timeout.ts

@@ -0,0 +1,104 @@
+import type { ToolDefinition } from "@earendil-works/pi-coding-agent";
+
+export const CHILD_TOOL_CALL_TIMEOUT_MS = 3 * 60 * 1_000;
+
+interface ToolRegistry {
+  getAllTools(): Array<{ name: string }>;
+  getToolDefinition(name: string): ToolDefinition | undefined;
+}
+
+function formatTimeout(timeoutMs: number) {
+  if (timeoutMs % 60_000 === 0) {
+    const minutes = timeoutMs / 60_000;
+    return `${minutes} minute${minutes === 1 ? "" : "s"}`;
+  }
+  if (timeoutMs % 1_000 === 0) {
+    const seconds = timeoutMs / 1_000;
+    return `${seconds} second${seconds === 1 ? "" : "s"}`;
+  }
+  return `${timeoutMs} ms`;
+}
+
+export class ToolCallTimeoutError extends Error {
+  constructor(toolName: string, timeoutMs: number) {
+    super(
+      `Tool call "${toolName}" timed out after ${formatTimeout(timeoutMs)}.`,
+    );
+    this.name = "ToolCallTimeoutError";
+  }
+}
+
+export async function runWithToolCallTimeout<T>(
+  toolName: string,
+  timeoutMs: number,
+  signal: AbortSignal | undefined,
+  execute: (signal: AbortSignal) => Promise<T>,
+) {
+  const timeoutController = new AbortController();
+  const executionSignal = signal
+    ? AbortSignal.any([signal, timeoutController.signal])
+    : timeoutController.signal;
+  const timeoutError = new ToolCallTimeoutError(toolName, timeoutMs);
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  const timeout = new Promise<never>((_resolve, reject) => {
+    timer = setTimeout(() => {
+      reject(timeoutError);
+      timeoutController.abort(timeoutError);
+    }, timeoutMs);
+  });
+
+  let removeAbortListener: (() => void) | undefined;
+  const aborted = new Promise<never>((_resolve, reject) => {
+    if (!signal) return;
+    const onAbort = () => {
+      reject(
+        signal.reason instanceof Error
+          ? signal.reason
+          : new Error(`Tool call "${toolName}" was aborted.`),
+      );
+    };
+    if (signal.aborted) {
+      onAbort();
+      return;
+    }
+    signal.addEventListener("abort", onAbort, { once: true });
+    removeAbortListener = () => signal.removeEventListener("abort", onAbort);
+  });
+
+  try {
+    return await Promise.race([execute(executionSignal), timeout, aborted]);
+  } finally {
+    if (timer) clearTimeout(timer);
+    removeAbortListener?.();
+  }
+}
+
+/**
+ * Wrap every currently registered child tool with an independent execution
+ * timeout. Calling apply() again is safe and picks up tools registered later.
+ */
+export function createToolCallTimeoutGuard(
+  timeoutMs = CHILD_TOOL_CALL_TIMEOUT_MS,
+) {
+  const wrapped = new WeakSet<ToolDefinition>();
+
+  const wrap = (definition: ToolDefinition) => {
+    if (wrapped.has(definition)) return;
+    wrapped.add(definition);
+
+    const execute = definition.execute;
+    definition.execute = async (toolCallId, params, signal, onUpdate, ctx) =>
+      runWithToolCallTimeout(definition.name, timeoutMs, signal, (signal) =>
+        execute.call(definition, toolCallId, params, signal, onUpdate, ctx),
+      );
+  };
+
+  return {
+    apply(session: ToolRegistry) {
+      for (const { name } of session.getAllTools()) {
+        const definition = session.getToolDefinition(name);
+        if (definition) wrap(definition);
+      }
+    },
+  };
+}