diff --git a/.changeset/remove-enter-with.md b/.changeset/remove-enter-with.md
new file mode 100644
index 000000000..d661de675
--- /dev/null
+++ b/.changeset/remove-enter-with.md
@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+ref: Remove `AsyncLocalStorage.enterWith()` usage
diff --git a/js/src/global-instrumentation-hooks.test.ts b/js/src/global-instrumentation-hooks.test.ts
index f46573e86..687695850 100644
--- a/js/src/global-instrumentation-hooks.test.ts
+++ b/js/src/global-instrumentation-hooks.test.ts
@@ -446,7 +446,6 @@ describe("global instrumentation hooks", () => {
 
     const storeError = new Error("store failed");
     const brokenStore = {
-      enterWith() {},
       getStore() {
         return undefined;
       },
@@ -475,7 +474,6 @@ describe("global instrumentation hooks", () => {
 
     let callback: (() => unknown) | undefined;
     channel.start.bindStore({
-      enterWith() {},
       getStore() {
         return undefined;
       },
diff --git a/js/src/global-instrumentation-hooks.ts b/js/src/global-instrumentation-hooks.ts
index 5dbefc85c..119d1bfd5 100644
--- a/js/src/global-instrumentation-hooks.ts
+++ b/js/src/global-instrumentation-hooks.ts
@@ -19,7 +19,6 @@ const hookBrand = Symbol.for(GLOBAL_INSTRUMENTATION_HOOK_BRAND);
 const invocationHookBrand = Symbol.for(GLOBAL_INVOCATION_HOOK_BRAND);
 
 export interface GlobalHookAsyncLocalStorage<T> {
-  enterWith(store: T): void;
   run<R>(store: T | undefined, callback: () => R): R;
   getStore(): T | undefined;
 }
diff --git a/js/src/instrumentation/auto-instrumentation-suppression.test.ts b/js/src/instrumentation/auto-instrumentation-suppression.test.ts
index d4561deb7..f3439ac2e 100644
--- a/js/src/instrumentation/auto-instrumentation-suppression.test.ts
+++ b/js/src/instrumentation/auto-instrumentation-suppression.test.ts
@@ -1,8 +1,8 @@
 import { beforeAll, describe, expect, it } from "vitest";
 import { configureNode } from "../node/config";
 import {
-  enterAutoInstrumentationAllowed,
   isAutoInstrumentationSuppressed,
+  runWithAutoInstrumentationAllowed,
   runWithAutoInstrumentationSuppressed,
 } from "./auto-instrumentation-suppression";
 
@@ -19,32 +19,34 @@ describe("auto instrumentation suppression context", () => {
       await Promise.resolve();
       expect(isAutoInstrumentationSuppressed()).toBe(true);
 
-      const restoreToolContext = enterAutoInstrumentationAllowed();
-      expect(isAutoInstrumentationSuppressed()).toBe(false);
+      await runWithAutoInstrumentationAllowed(async () => {
+        expect(isAutoInstrumentationSuppressed()).toBe(false);
 
-      await runWithAutoInstrumentationSuppressed(async () => {
-        expect(isAutoInstrumentationSuppressed()).toBe(true);
-        await Promise.resolve();
-        expect(isAutoInstrumentationSuppressed()).toBe(true);
-      });
+        await runWithAutoInstrumentationSuppressed(async () => {
+          expect(isAutoInstrumentationSuppressed()).toBe(true);
+          await Promise.resolve();
+          expect(isAutoInstrumentationSuppressed()).toBe(true);
+        });
 
-      expect(isAutoInstrumentationSuppressed()).toBe(false);
-      restoreToolContext();
+        expect(isAutoInstrumentationSuppressed()).toBe(false);
+      });
       expect(isAutoInstrumentationSuppressed()).toBe(true);
     });
 
     expect(isAutoInstrumentationSuppressed()).toBe(false);
   });
 
-  it("keeps instrumentation allowed until every active allow frame exits", async () => {
+  it("restores nested allow contexts at each callback boundary", async () => {
     await runWithAutoInstrumentationSuppressed(async () => {
-      const restoreFirstTool = enterAutoInstrumentationAllowed();
-      const restoreSecondTool = enterAutoInstrumentationAllowed();
-
-      expect(isAutoInstrumentationSuppressed()).toBe(false);
-      restoreFirstTool();
-      expect(isAutoInstrumentationSuppressed()).toBe(false);
-      restoreSecondTool();
+      await runWithAutoInstrumentationAllowed(async () => {
+        expect(isAutoInstrumentationSuppressed()).toBe(false);
+        await runWithAutoInstrumentationAllowed(async () => {
+          expect(isAutoInstrumentationSuppressed()).toBe(false);
+          await Promise.resolve();
+          expect(isAutoInstrumentationSuppressed()).toBe(false);
+        });
+        expect(isAutoInstrumentationSuppressed()).toBe(false);
+      });
       expect(isAutoInstrumentationSuppressed()).toBe(true);
     });
   });
diff --git a/js/src/instrumentation/auto-instrumentation-suppression.ts b/js/src/instrumentation/auto-instrumentation-suppression.ts
index 2b31044ef..8bb6b4fa7 100644
--- a/js/src/instrumentation/auto-instrumentation-suppression.ts
+++ b/js/src/instrumentation/auto-instrumentation-suppression.ts
@@ -1,89 +1,22 @@
-import iso, {
-  type IsoAsyncLocalStorage,
-  type IsoTracingChannel,
-} from "../isomorph";
-
-type AutoInstrumentationSuppressionFrame = {
-  id: symbol;
-  mode: "allow" | "suppress";
-};
-
-type AutoInstrumentationSuppressionState = {
-  frames: AutoInstrumentationSuppressionFrame[];
-};
+import iso, { type IsoAsyncLocalStorage } from "../isomorph";
 
 let autoInstrumentationSuppressionStore:
-  | IsoAsyncLocalStorage<AutoInstrumentationSuppressionState | undefined>
+  | IsoAsyncLocalStorage<boolean>
   | undefined;
 
 function suppressionStore() {
-  autoInstrumentationSuppressionStore ??= iso.newAsyncLocalStorage<
-    AutoInstrumentationSuppressionState | undefined
-  >();
+  autoInstrumentationSuppressionStore ??= iso.newAsyncLocalStorage<boolean>();
   return autoInstrumentationSuppressionStore;
 }
 
-function currentFrames(): AutoInstrumentationSuppressionFrame[] {
-  return suppressionStore().getStore()?.frames ?? [];
-}
-
 export function isAutoInstrumentationSuppressed(): boolean {
-  const frames = currentFrames();
-  return frames[frames.length - 1]?.mode === "suppress";
+  return suppressionStore().getStore() === true;
 }
 
 export function runWithAutoInstrumentationSuppressed<R>(callback: () => R): R {
-  const frame = {
-    id: Symbol("braintrust.auto-instrumentation-suppress"),
-    mode: "suppress" as const,
-  };
-  return suppressionStore().run(
-    { frames: [...currentFrames(), frame] },
-    callback,
-  );
+  return suppressionStore().run(true, callback);
 }
 
-export function bindAutoInstrumentationSuppressionToStart<T>(
-  tracingChannel: Pick<IsoTracingChannel<T>, "start">,
-): (() => void) | undefined {
-  const startChannel = tracingChannel.start;
-  if (!startChannel) {
-    return undefined;
-  }
-
-  const store = suppressionStore();
-  startChannel.bindStore(store, () => ({
-    frames: [
-      ...currentFrames(),
-      {
-        id: Symbol("braintrust.auto-instrumentation-suppress"),
-        mode: "suppress" as const,
-      },
-    ],
-  }));
-
-  return () => {
-    startChannel.unbindStore(store);
-  };
-}
-
-export function enterAutoInstrumentationAllowed(): () => void {
-  const frame = {
-    id: Symbol("braintrust.auto-instrumentation-allow"),
-    mode: "allow" as const,
-  };
-  // TODO(luca): Replace ALS.enterWith() with ALS.run()
-  // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above.
-  suppressionStore().enterWith({
-    frames: [...currentFrames(), frame],
-  });
-
-  return () => {
-    const frames = currentFrames().filter(
-      (candidate) => candidate.id !== frame.id,
-    );
-    // TODO(luca): Replace ALS.enterWith() with ALS.run()
-    // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above.
-    suppressionStore().enterWith(frames.length > 0 ? { frames } : undefined);
-  };
+export function runWithAutoInstrumentationAllowed<R>(callback: () => R): R {
+  return suppressionStore().run(undefined, callback);
 }
diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts
index fcc538deb..0cfd405de 100644
--- a/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts
+++ b/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts
@@ -38,6 +38,7 @@ const mockNewTracingChannel = iso.newTracingChannel as ReturnType<typeof vi.fn>;
 type MockTracingChannel = {
   handlers: any[];
   hasSubscribers: boolean;
+  intercept: ReturnType<typeof vi.fn>;
   subscribe: ReturnType<typeof vi.fn>;
   unsubscribe: ReturnType<typeof vi.fn>;
 };
@@ -66,6 +67,23 @@ describe("AISDKPlugin", () => {
       const channel: MockTracingChannel = {
         handlers: [],
         hasSubscribers: false,
+        intercept: vi.fn((interceptor: any) => {
+          const handlers = {
+            end: (event: any) =>
+              interceptor(
+                () => event.result,
+                event.self,
+                event.arguments ?? [],
+                {},
+              ),
+          };
+          channel.handlers.push(handlers);
+          return vi.fn(() => {
+            channel.handlers = channel.handlers.filter(
+              (candidate) => candidate !== handlers,
+            );
+          });
+        }),
         subscribe: vi.fn((handlers: any) => {
           channel.handlers.push(handlers);
           channel.hasSubscribers = true;
@@ -243,7 +261,7 @@ describe("AISDKPlugin", () => {
       const channel = mockChannels.get(
         "orchestrion:ai:createTelemetryDispatcher",
       );
-      expect(channel?.subscribe).toHaveBeenCalledTimes(1);
+      expect(channel?.intercept).toHaveBeenCalledTimes(1);
 
       channel?.handlers[0]?.end({
         arguments: [{ telemetry: {} }],
diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/plugins/ai-sdk-plugin.ts
index e6cbf92f4..d23ff888b 100644
--- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts
+++ b/js/src/instrumentation/plugins/ai-sdk-plugin.ts
@@ -193,7 +193,7 @@ export class AISDKPlugin extends BasePlugin {
     const denyOutputPaths =
       this.config.denyOutputPaths || DEFAULT_DENY_OUTPUT_PATHS;
 
-    this.unsubscribers.push(subscribeToAISDKV7TelemetryDispatcher());
+    this.unsubscribers.push(interceptAISDKV7TelemetryDispatcher());
     this.unsubscribers.push(subscribeToHarnessAgentCreateSession());
     this.unsubscribers.push(
       subscribeToHarnessContinuation(
@@ -839,31 +839,29 @@ function subscribeToHarnessContinuation(
   };
 }
 
-function subscribeToAISDKV7TelemetryDispatcher(): () => void {
-  const channel = aiSDKChannels.v7CreateTelemetryDispatcher.tracingChannel();
+function interceptAISDKV7TelemetryDispatcher(): () => void {
   const telemetry = braintrustAISDKTelemetry();
-  const handlers: IsoChannelHandlers<
-    ChannelMessage<typeof aiSDKChannels.v7CreateTelemetryDispatcher>
-  > = {
-    end: (event) => {
-      const telemetryOptions = event.arguments?.[0]?.telemetry;
-      if (telemetryOptions?.isEnabled === false) {
-        return;
+  return aiSDKChannels.v7CreateTelemetryDispatcher.intercept(
+    (target, thisArg, args) => {
+      const dispatcher = Reflect.apply(target, thisArg, args);
+      const telemetryOptions = args[0]?.telemetry;
+      if (telemetryOptions?.isEnabled !== false) {
+        try {
+          patchAISDKV7TelemetryDispatcher(
+            dispatcher,
+            telemetry,
+            telemetryOptions,
+          );
+        } catch (error) {
+          debugLogger.error(
+            "Error instrumenting AI SDK v7 telemetry dispatcher:",
+            error,
+          );
+        }
       }
-
-      patchAISDKV7TelemetryDispatcher(
-        event.result,
-        telemetry,
-        telemetryOptions,
-      );
+      return dispatcher;
     },
-  };
-
-  channel.subscribe(handlers);
-
-  return () => {
-    channel.unsubscribe(handlers);
-  };
+  );
 }
 
 function patchAISDKV7TelemetryDispatcher(
@@ -891,7 +889,6 @@ function patchAISDKV7TelemetryDispatcher(
   if (typeof telemetryOptions?.functionId === "string") {
     telemetryEventFields.functionId = telemetryOptions.functionId;
   }
-
   const eventWithOperationKey = (event: unknown): unknown => {
     if (!isObject(event)) {
       return event;
diff --git a/js/src/instrumentation/plugins/anthropic-sessions-plugin.test.ts b/js/src/instrumentation/plugins/anthropic-sessions-plugin.test.ts
index 1d4cb48ac..01cc32592 100644
--- a/js/src/instrumentation/plugins/anthropic-sessions-plugin.test.ts
+++ b/js/src/instrumentation/plugins/anthropic-sessions-plugin.test.ts
@@ -9,7 +9,6 @@ vi.mock("../../isomorph", () => ({
   default: {
     getEnv: vi.fn(),
     newAsyncLocalStorage: vi.fn(() => ({
-      enterWith: vi.fn(),
       getStore: vi.fn(() => undefined),
       run: vi.fn((_store: unknown, callback: () => unknown) => callback()),
     })),
diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-channels.ts b/js/src/instrumentation/plugins/claude-agent-sdk-channels.ts
index 959065035..c58e23144 100644
--- a/js/src/instrumentation/plugins/claude-agent-sdk-channels.ts
+++ b/js/src/instrumentation/plugins/claude-agent-sdk-channels.ts
@@ -11,7 +11,7 @@ export const claudeAgentSDKChannels = defineChannels(
     query: channel<
       [ClaudeAgentSDKQueryParams],
       AsyncIterable<ClaudeAgentSDKMessage>,
-      Record<string, never>,
+      Record<never, never>,
       ClaudeAgentSDKMessage
     >({
       channelName: "query",
diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts
index 3bf15777a..e0c2129c3 100644
--- a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts
+++ b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts
@@ -1,178 +1,37 @@
 import iso from "../../isomorph";
 
-type LocalToolParentResolver = (toolUseId: string) => Promise<string>;
+type ClaudeLocalToolParentResolver = (toolUseId: string) => Promise<string>;
 
-export type ClaudeAgentSDKLocalToolContext = {
-  resolveLocalToolParent?: LocalToolParentResolver;
-};
+const localToolContextStore =
+  iso.newAsyncLocalStorage<ClaudeLocalToolParentResolver>();
+const localToolParentResolversByToolUseId = new Map<
+  string,
+  ClaudeLocalToolParentResolver
+>();
 
-const LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED = Symbol.for(
-  "braintrust.claude_agent_sdk.local_tool_context_async_iterator_patched",
-);
-
-type AsyncLocalStorageLike<T> = {
-  enterWith: (store: T) => void;
-  getStore: () => T | undefined;
-  run: <R>(store: T, callback: () => R) => R;
-};
-
-function createLocalToolContextStore(): AsyncLocalStorageLike<ClaudeAgentSDKLocalToolContext> {
-  const maybeIsoWithAsyncLocalStorage = iso as {
-    newAsyncLocalStorage?: <T>() => AsyncLocalStorageLike<T>;
-  };
-
-  if (
-    typeof maybeIsoWithAsyncLocalStorage.newAsyncLocalStorage === "function"
-  ) {
-    return maybeIsoWithAsyncLocalStorage.newAsyncLocalStorage<ClaudeAgentSDKLocalToolContext>();
-  }
-
-  let currentStore: ClaudeAgentSDKLocalToolContext | undefined;
-  return {
-    enterWith(store) {
-      currentStore = store;
-    },
-    getStore() {
-      return currentStore;
-    },
-    run(store, callback) {
-      const previousStore = currentStore;
-      currentStore = store;
-      try {
-        return callback();
-      } finally {
-        currentStore = previousStore;
-      }
-    },
-  };
-}
-
-const localToolContextStore = createLocalToolContextStore();
-let fallbackLocalToolParentResolver: LocalToolParentResolver | undefined;
-
-export function createClaudeLocalToolContext(): ClaudeAgentSDKLocalToolContext {
-  return {};
-}
-
-function runWithClaudeLocalToolContext<R>(
+export function runWithClaudeLocalToolContext<R>(
   callback: () => R,
-  context?: ClaudeAgentSDKLocalToolContext,
+  resolver: ClaudeLocalToolParentResolver,
 ): R {
-  return localToolContextStore.run(
-    context ?? createClaudeLocalToolContext(),
-    callback,
-  );
+  return localToolContextStore.run(resolver, callback);
 }
 
-function ensureClaudeLocalToolContext():
-  | ClaudeAgentSDKLocalToolContext
-  | undefined {
-  const existing = localToolContextStore.getStore();
-  if (existing) {
-    return existing;
-  }
-
-  const created: ClaudeAgentSDKLocalToolContext = {};
-  // TODO(luca): Replace ALS.enterWith() with ALS.run()
-  // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above.
-  localToolContextStore.enterWith(created);
-  return created;
-}
-
-export function setClaudeLocalToolParentResolver(
-  resolver: LocalToolParentResolver,
+export function registerClaudeLocalToolParentResolver(
+  toolUseId: string,
+  resolver: ClaudeLocalToolParentResolver,
 ): void {
-  fallbackLocalToolParentResolver = resolver;
-  const context = ensureClaudeLocalToolContext();
-  if (!context) {
-    return;
-  }
-  context.resolveLocalToolParent = resolver;
-}
-
-export function getClaudeLocalToolParentResolver():
-  | LocalToolParentResolver
-  | undefined {
-  return (
-    localToolContextStore.getStore()?.resolveLocalToolParent ??
-    fallbackLocalToolParentResolver
-  );
-}
-
-function isAsyncIterable(value: unknown): value is AsyncIterable<unknown> {
-  return (
-    value !== null &&
-    typeof value === "object" &&
-    Symbol.asyncIterator in value &&
-    typeof value[Symbol.asyncIterator] === "function"
-  );
+  localToolParentResolversByToolUseId.set(toolUseId, resolver);
 }
 
-export function bindClaudeLocalToolContextToAsyncIterable<T>(
-  result: T,
-  localToolContext: ClaudeAgentSDKLocalToolContext,
-): T {
-  if (
-    !isAsyncIterable(result) ||
-    Object.isFrozen(result) ||
-    Object.isSealed(result)
-  ) {
-    return result;
+export function getClaudeLocalToolParentResolver(
+  toolUseId?: string,
+): ClaudeLocalToolParentResolver | undefined {
+  const currentResolver = localToolContextStore.getStore();
+  if (!toolUseId) {
+    return currentResolver;
   }
 
-  const stream = result as AsyncIterable<unknown> & {
-    [Symbol.asyncIterator]: (() => AsyncIterator<unknown>) & {
-      [LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED]?: boolean;
-    };
-  };
-  const originalAsyncIterator = stream[Symbol.asyncIterator];
-  if (originalAsyncIterator[LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED]) {
-    return result;
-  }
-
-  const patchedAsyncIterator = function (this: unknown) {
-    return runWithClaudeLocalToolContext(() => {
-      const iterator = Reflect.apply(originalAsyncIterator, this, []);
-      if (!iterator || typeof iterator !== "object") {
-        return iterator;
-      }
-
-      const patchMethod = (methodName: "next" | "return" | "throw") => {
-        const originalMethod = Reflect.get(iterator, methodName);
-        if (typeof originalMethod !== "function") {
-          return;
-        }
-
-        Reflect.set(iterator, methodName, (...args: unknown[]) =>
-          runWithClaudeLocalToolContext(
-            () =>
-              Reflect.apply(
-                originalMethod as (...methodArgs: unknown[]) => unknown,
-                iterator,
-                args,
-              ),
-            localToolContext,
-          ),
-        );
-      };
-
-      patchMethod("next");
-      patchMethod("return");
-      patchMethod("throw");
-      return iterator;
-    }, localToolContext);
-  };
-
-  Object.defineProperty(
-    patchedAsyncIterator,
-    LOCAL_TOOL_CONTEXT_ASYNC_ITERATOR_PATCHED,
-    {
-      configurable: false,
-      enumerable: false,
-      value: true,
-      writable: false,
-    },
-  );
-  Reflect.set(stream, Symbol.asyncIterator, patchedAsyncIterator);
-  return result;
+  const registeredResolver = localToolParentResolversByToolUseId.get(toolUseId);
+  localToolParentResolversByToolUseId.delete(toolUseId);
+  return currentResolver ?? registeredResolver;
 }
diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts
index 1a1c7db88..843309851 100644
--- a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts
+++ b/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts
@@ -53,7 +53,7 @@ export function wrapLocalClaudeToolHandler(
       ? `mcp__${metadata.serverName}__${metadata.toolName}`
       : metadata.toolName;
     const toolUseId = getToolUseIdFromExtra(handlerArgs[1]);
-    const localToolParentResolver = getClaudeLocalToolParentResolver();
+    const localToolParentResolver = getClaudeLocalToolParentResolver(toolUseId);
     const spanName = metadata.serverName
       ? `tool: ${metadata.serverName}/${metadata.toolName}`
       : `tool: ${metadata.toolName}`;
diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts
index f1d446bcd..1c35a4a27 100644
--- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts
+++ b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts
@@ -1,4 +1,5 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
+import { AsyncLocalStorage } from "node:async_hooks";
 
 // Mock iso's newTracingChannel - must be before any imports that use it
 const streamPatcherMock = vi.hoisted(() => ({
@@ -12,6 +13,7 @@ const streamPatcherMock = vi.hoisted(() => ({
 
 vi.mock("../../isomorph", () => ({
   default: {
+    newAsyncLocalStorage: <T>() => new AsyncLocalStorage<T>(),
     newTracingChannel: vi.fn(),
   },
 }));
@@ -123,11 +125,16 @@ describe("ClaudeAgentSDKPlugin", () => {
   let plugin: ClaudeAgentSDKPlugin;
   let mockChannel: any;
   let mockUnsubscribe: any;
+  let queryInterceptor: any;
 
   beforeEach(() => {
     streamPatcherMock.options = undefined;
     mockUnsubscribe = vi.fn();
     mockChannel = {
+      intercept: vi.fn((interceptor) => {
+        queryInterceptor = interceptor;
+        return mockUnsubscribe;
+      }),
       subscribe: vi.fn(),
       unsubscribe: mockUnsubscribe,
       hasSubscribers: false,
@@ -149,21 +156,15 @@ describe("ClaudeAgentSDKPlugin", () => {
       expect(mockNewTracingChannel).toHaveBeenCalledWith(
         "orchestrion:@anthropic-ai/claude-agent-sdk:query",
       );
-      expect(mockChannel.subscribe).toHaveBeenCalledTimes(1);
-      expect(mockChannel.subscribe).toHaveBeenCalledWith(
-        expect.objectContaining({
-          start: expect.any(Function),
-          end: expect.any(Function),
-          error: expect.any(Function),
-        }),
-      );
+      expect(mockChannel.intercept).toHaveBeenCalledTimes(1);
+      expect(mockChannel.intercept).toHaveBeenCalledWith(expect.any(Function));
     });
 
     it("should not subscribe twice if already enabled", () => {
       plugin.enable();
       plugin.enable();
 
-      expect(mockChannel.subscribe).toHaveBeenCalledTimes(1);
+      expect(mockChannel.intercept).toHaveBeenCalledTimes(1);
     });
 
     it("should store unsubscribe function", () => {
@@ -202,7 +203,34 @@ describe("ClaudeAgentSDKPlugin", () => {
 
     beforeEach(() => {
       plugin.enable();
-      handlers = mockChannel.subscribe.mock.calls[0][0];
+      handlers = {
+        start: (event: any) =>
+          queryInterceptor(
+            () => ({
+              async *[Symbol.asyncIterator]() {
+                // Keep the query span open so tests can drive stream callbacks.
+              },
+            }),
+            event.self,
+            event.arguments ?? [],
+            {},
+          ),
+        end: () => undefined,
+        error: (event: any) => {
+          try {
+            queryInterceptor(
+              () => {
+                throw event.error;
+              },
+              event.self,
+              event.arguments ?? [],
+              {},
+            );
+          } catch {
+            // The invocation interceptor preserves the target's exception.
+          }
+        },
+      };
     });
 
     describe("start handler", () => {
@@ -752,7 +780,7 @@ describe("ClaudeAgentSDKPlugin", () => {
       plugin.disable();
       plugin.enable();
 
-      expect(mockChannel.subscribe).toHaveBeenCalledTimes(2);
+      expect(mockChannel.intercept).toHaveBeenCalledTimes(2);
     });
 
     it("should properly clean up on multiple enable/disable cycles", () => {
diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts
index 96f52c9b6..86389ee9d 100644
--- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts
+++ b/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts
@@ -1,7 +1,6 @@
 import { BasePlugin } from "../core";
-import type { ChannelMessage } from "../core/channel-definitions";
 import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher";
-import type { IsoChannelHandlers } from "../../isomorph";
+import { debugLogger } from "../../debug-logger";
 import { startSpan as startBaseSpan } from "../../logger";
 import type { Span } from "../../logger";
 import {
@@ -23,10 +22,8 @@ import {
   wrapLocalMcpServerToolHandlers,
 } from "./claude-agent-sdk-local-tool-spans";
 import {
-  bindClaudeLocalToolContextToAsyncIterable,
-  createClaudeLocalToolContext,
-  setClaudeLocalToolParentResolver,
-  type ClaudeAgentSDKLocalToolContext,
+  registerClaudeLocalToolParentResolver,
+  runWithClaudeLocalToolContext,
 } from "./claude-agent-sdk-local-tool-context";
 import type {
   ClaudeAgentSDKHookCallback,
@@ -627,6 +624,7 @@ function createToolTracingHooks(
       (isLocalToolUse(input.tool_name, mcpServers) ||
         localToolHookNames.has(input.tool_name))
     ) {
+      registerClaudeLocalToolParentResolver(toolUseID, resolveParentSpan);
       return {};
     }
 
@@ -985,7 +983,7 @@ type QueryState = {
   latestRootLlmParentRef: { value: string | undefined };
   toolUseToParent: Map<string, string | null>;
   usageByMessageId: Map<string, ClaudeAgentSDKUsage>;
-  localToolContext: ClaudeAgentSDKLocalToolContext;
+  localToolParentResolver: ParentSpanResolver;
 };
 
 function setSubAgentPromptMessages(
@@ -1580,7 +1578,7 @@ async function finalizeQuerySpan(state: QueryState): Promise<void> {
 
 export class ClaudeAgentSDKPlugin extends BasePlugin {
   protected onEnable(): void {
-    this.subscribeToQuery();
+    this.interceptQuery();
   }
 
   protected onDisable(): void {
@@ -1590,272 +1588,276 @@ export class ClaudeAgentSDKPlugin extends BasePlugin {
     this.unsubscribers = [];
   }
 
-  private subscribeToQuery(): void {
-    const channel = claudeAgentSDKChannels.query.tracingChannel();
-    const spans = new WeakMap<object, QueryState>();
-
-    const handlers: IsoChannelHandlers<
-      ChannelMessage<typeof claudeAgentSDKChannels.query>
-    > = {
-      start: (event) => {
-        const params = (event.arguments[0] ?? {}) as ClaudeAgentSDKQueryParams;
-        const originalPrompt = params.prompt;
-        const options = params.options ?? {};
-        const promptIsAsyncIterable = isAsyncIterable(originalPrompt);
-        let promptStarted = false;
-        let capturedPromptMessages: ClaudeAgentSDKMessage[] | undefined;
-        let resolvePromptDone: (() => void) | undefined;
-        const promptDone = new Promise<void>((resolve) => {
-          resolvePromptDone = resolve;
-        });
+  private interceptQuery(): void {
+    const startQuery = (params: ClaudeAgentSDKQueryParams): QueryState => {
+      const originalPrompt = params.prompt;
+      const options = params.options ?? {};
+      const promptIsAsyncIterable = isAsyncIterable(originalPrompt);
+      let promptStarted = false;
+      let capturedPromptMessages: ClaudeAgentSDKMessage[] | undefined;
+      let resolvePromptDone: (() => void) | undefined;
+      const promptDone = new Promise<void>((resolve) => {
+        resolvePromptDone = resolve;
+      });
 
-        if (promptIsAsyncIterable) {
-          capturedPromptMessages = [];
-          const promptStream =
-            originalPrompt as AsyncIterable<ClaudeAgentSDKMessage>;
-          params.prompt = (async function* () {
-            promptStarted = true;
-            try {
-              for await (const message of promptStream) {
-                capturedPromptMessages!.push(message);
-                yield message;
-              }
-            } finally {
-              resolvePromptDone?.();
+      if (promptIsAsyncIterable) {
+        capturedPromptMessages = [];
+        const promptStream =
+          originalPrompt as AsyncIterable<ClaudeAgentSDKMessage>;
+        params.prompt = (async function* () {
+          promptStarted = true;
+          try {
+            for await (const message of promptStream) {
+              capturedPromptMessages!.push(message);
+              yield message;
             }
-          })();
-        }
+          } finally {
+            resolvePromptDone?.();
+          }
+        })();
+      }
 
-        const span = startBaseSpan(
-          withSpanInstrumentationName(
-            {
-              name: "Claude Agent",
-              spanAttributes: {
-                type: SpanTypeAttribute.TASK,
-              },
+      const span = startBaseSpan(
+        withSpanInstrumentationName(
+          {
+            name: "Claude Agent",
+            spanAttributes: {
+              type: SpanTypeAttribute.TASK,
             },
-            INSTRUMENTATION_NAMES.CLAUDE_AGENT_SDK,
-          ),
-        );
-        const startTime = getCurrentUnixTimestamp();
+          },
+          INSTRUMENTATION_NAMES.CLAUDE_AGENT_SDK,
+        ),
+      );
+      const startTime = getCurrentUnixTimestamp();
 
-        try {
-          span.log({
-            input:
-              typeof originalPrompt === "string"
-                ? originalPrompt
-                : promptIsAsyncIterable
-                  ? undefined
-                  : originalPrompt !== undefined
-                    ? String(originalPrompt)
-                    : undefined,
-            metadata: filterSerializableOptions(options),
-          });
-        } catch (error) {
-          // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage.
-          console.error("Error extracting input for Claude Agent SDK:", error);
-        }
+      try {
+        span.log({
+          input:
+            typeof originalPrompt === "string"
+              ? originalPrompt
+              : promptIsAsyncIterable
+                ? undefined
+                : originalPrompt !== undefined
+                  ? String(originalPrompt)
+                  : undefined,
+          metadata: filterSerializableOptions(options),
+        });
+      } catch (error) {
+        // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage.
+        console.error("Error extracting input for Claude Agent SDK:", error);
+      }
 
-        const activeToolSpans = new Map<string, Span>();
-        const activeLlmSpansByParentToolUse = new Map<string, Span>();
-        const conversationHistoryByParentKey = new Map<
-          string,
-          ClaudeConversationMessage[]
-        >();
-        const subAgentSpans = new Map<string, Span>();
-        const endedSubAgentSpans = new Set<string>();
-        const toolUseToParent = new Map<string, string | null>();
-        const latestLlmParentBySubAgentToolUse = new Map<string, string>();
-        const latestRootLlmParentRef = {
-          value: undefined as string | undefined,
-        };
-        const subAgentDetailsByToolUseId = new Map<string, SubAgentDetails>();
-        const taskIdToToolUseId = new Map<string, string>();
-        const promptMessagesByParentKey = new Map<
-          string,
-          ClaudeConversationMessage[]
-        >();
-        const promptSourcePriorityByParentKey = new Map<string, number>();
-        const localToolContext = createClaudeLocalToolContext();
-        const { hasLocalToolHandlers, localToolHookNames } =
-          prepareLocalToolHandlersInMcpServers(options.mcpServers);
-        const skipLocalToolHooks =
-          options[CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION] === true ||
-          hasLocalToolHandlers;
-        const resolveToolUseParentSpan: ParentSpanResolver = async (
-          toolUseID,
-          context,
-        ) => {
-          const trackedParentToolUseId = toolUseToParent.get(toolUseID);
-          const parentToolUseId =
-            trackedParentToolUseId ??
-            (context?.agentId
-              ? (taskIdToToolUseId.get(context.agentId) ?? null)
-              : null);
-          const parentKey = llmParentKey(parentToolUseId);
-          const activeLlmSpan = activeLlmSpansByParentToolUse.get(parentKey);
-          const latestLlmParent = parentToolUseId
-            ? latestLlmParentBySubAgentToolUse.get(parentToolUseId)
-            : latestRootLlmParentRef.value;
-
-          // Tool spans should be siblings of the driving LLM turn, but we still
-          // materialize that LLM span first so trace ordering reflects that the
-          // tool call was produced by the model.
-          if (!activeLlmSpan && !latestLlmParent) {
-            await ensureActiveLlmSpanForParentToolUse(
-              span,
-              activeLlmSpansByParentToolUse,
-              subAgentDetailsByToolUseId,
-              activeToolSpans,
-              subAgentSpans,
-              parentToolUseId,
-              getCurrentUnixTimestamp(),
-            );
-          }
+      const activeToolSpans = new Map<string, Span>();
+      const activeLlmSpansByParentToolUse = new Map<string, Span>();
+      const conversationHistoryByParentKey = new Map<
+        string,
+        ClaudeConversationMessage[]
+      >();
+      const subAgentSpans = new Map<string, Span>();
+      const endedSubAgentSpans = new Set<string>();
+      const toolUseToParent = new Map<string, string | null>();
+      const latestLlmParentBySubAgentToolUse = new Map<string, string>();
+      const latestRootLlmParentRef = {
+        value: undefined as string | undefined,
+      };
+      const subAgentDetailsByToolUseId = new Map<string, SubAgentDetails>();
+      const taskIdToToolUseId = new Map<string, string>();
+      const promptMessagesByParentKey = new Map<
+        string,
+        ClaudeConversationMessage[]
+      >();
+      const promptSourcePriorityByParentKey = new Map<string, number>();
+      const { hasLocalToolHandlers, localToolHookNames } =
+        prepareLocalToolHandlersInMcpServers(options.mcpServers);
+      const skipLocalToolHooks =
+        options[CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION] === true ||
+        hasLocalToolHandlers;
+      const resolveToolUseParentSpan: ParentSpanResolver = async (
+        toolUseID,
+        context,
+      ) => {
+        const trackedParentToolUseId = toolUseToParent.get(toolUseID);
+        const parentToolUseId =
+          trackedParentToolUseId ??
+          (context?.agentId
+            ? (taskIdToToolUseId.get(context.agentId) ?? null)
+            : null);
+        const parentKey = llmParentKey(parentToolUseId);
+        const activeLlmSpan = activeLlmSpansByParentToolUse.get(parentKey);
+        const latestLlmParent = parentToolUseId
+          ? latestLlmParentBySubAgentToolUse.get(parentToolUseId)
+          : latestRootLlmParentRef.value;
+
+        // Tool spans should be siblings of the driving LLM turn, but we still
+        // materialize that LLM span first so trace ordering reflects that the
+        // tool call was produced by the model.
+        if (!activeLlmSpan && !latestLlmParent) {
+          await ensureActiveLlmSpanForParentToolUse(
+            span,
+            activeLlmSpansByParentToolUse,
+            subAgentDetailsByToolUseId,
+            activeToolSpans,
+            subAgentSpans,
+            parentToolUseId,
+            getCurrentUnixTimestamp(),
+          );
+        }
 
-          if (parentToolUseId) {
-            const subAgentSpan = await ensureSubAgentSpan(
-              subAgentDetailsByToolUseId,
-              span,
-              activeToolSpans,
-              subAgentSpans,
-              parentToolUseId,
-            );
-            return subAgentSpan.export();
-          }
+        if (parentToolUseId) {
+          const subAgentSpan = await ensureSubAgentSpan(
+            subAgentDetailsByToolUseId,
+            span,
+            activeToolSpans,
+            subAgentSpans,
+            parentToolUseId,
+          );
+          return subAgentSpan.export();
+        }
 
-          return span.export();
-        };
+        return span.export();
+      };
+
+      const optionsWithHooks = injectTracingHooks(
+        options,
+        resolveToolUseParentSpan,
+        taskIdToToolUseId,
+        toolUseToParent,
+        activeToolSpans,
+        localToolHookNames,
+        skipLocalToolHooks,
+        subAgentDetailsByToolUseId,
+        subAgentSpans,
+        endedSubAgentSpans,
+      );
 
-        localToolContext.resolveLocalToolParent = resolveToolUseParentSpan;
-        setClaudeLocalToolParentResolver(resolveToolUseParentSpan);
-        const optionsWithHooks = injectTracingHooks(
-          options,
-          resolveToolUseParentSpan,
-          taskIdToToolUseId,
-          toolUseToParent,
-          activeToolSpans,
-          localToolHookNames,
-          skipLocalToolHooks,
-          subAgentDetailsByToolUseId,
-          subAgentSpans,
-          endedSubAgentSpans,
-        );
+      params.options = optionsWithHooks;
+
+      return {
+        activeLlmSpansByParentToolUse,
+        activePartialMessageIdByParentKey: new Map(),
+        activeToolSpans,
+        conversationHistoryByParentKey,
+        capturedPromptMessages,
+        currentMessageId: undefined,
+        currentMessageStartTime: startTime,
+        currentMessages: [],
+        endedSubAgentSpans,
+        finalOutputUsageMessageIds: new Set(),
+        finalResults: [],
+        options: optionsWithHooks,
+        originalPrompt,
+        processing: Promise.resolve(),
+        promptDone,
+        promptMessagesByParentKey,
+        promptStarted: () => promptStarted,
+        promptSourcePriorityByParentKey,
+        span,
+        subAgentDetailsByToolUseId,
+        subAgentSpans,
+        taskIdToToolUseId,
+        latestLlmParentBySubAgentToolUse,
+        latestRootLlmParentRef,
+        toolUseToParent,
+        usageByMessageId: new Map(),
+        localToolParentResolver: resolveToolUseParentSpan,
+      };
+    };
 
-        params.options = optionsWithHooks;
-        event.arguments[0] = params;
-
-        spans.set(event, {
-          activeLlmSpansByParentToolUse,
-          activePartialMessageIdByParentKey: new Map(),
-          activeToolSpans,
-          conversationHistoryByParentKey,
-          capturedPromptMessages,
-          currentMessageId: undefined,
-          currentMessageStartTime: startTime,
-          currentMessages: [],
-          endedSubAgentSpans,
-          finalOutputUsageMessageIds: new Set(),
-          finalResults: [],
-          options: optionsWithHooks,
-          originalPrompt,
-          processing: Promise.resolve(),
-          promptDone,
-          promptMessagesByParentKey,
-          promptStarted: () => promptStarted,
-          promptSourcePriorityByParentKey,
-          span,
-          subAgentDetailsByToolUseId,
-          subAgentSpans,
-          taskIdToToolUseId,
-          latestLlmParentBySubAgentToolUse,
-          latestRootLlmParentRef,
-          toolUseToParent,
-          usageByMessageId: new Map(),
-          localToolContext,
+    const finishQuery = (
+      state: QueryState,
+      result: AsyncIterable<ClaudeAgentSDKMessage>,
+    ): void => {
+      if (isAsyncIterable(result)) {
+        patchStreamIfNeeded(result, {
+          aroundNext: (callback) =>
+            runWithClaudeLocalToolContext(
+              callback,
+              state.localToolParentResolver,
+            ),
+          onChunk: (message: ClaudeAgentSDKMessage) => {
+            maybeTrackToolUseContext(state, message);
+            state.processing = state.processing
+              .then(() => handleStreamMessage(state, message))
+              .catch((error) => {
+                // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage.
+                console.error(
+                  "Error processing Claude Agent SDK stream chunk:",
+                  error,
+                );
+              });
+          },
+          onComplete: () =>
+            state.processing.then(() => finalizeQuerySpan(state)),
+          onError: (error: Error) =>
+            state.processing
+              .then(() => {
+                state.span.log({ error: error.message });
+              })
+              .then(() => finalizeQuerySpan(state)),
         });
-      },
-
-      end: (event) => {
-        const state = spans.get(event);
-        if (!state) {
-          return;
-        }
-
-        const eventResult = bindClaudeLocalToolContextToAsyncIterable(
-          event.result,
-          state.localToolContext,
-        );
-        if (eventResult === undefined) {
-          state.span.end();
-          spans.delete(event);
-          return;
-        }
 
-        if (isAsyncIterable(eventResult)) {
-          patchStreamIfNeeded(eventResult, {
-            onChunk: (message: ClaudeAgentSDKMessage) => {
-              maybeTrackToolUseContext(state, message);
-              state.processing = state.processing
-                .then(() => handleStreamMessage(state, message))
-                .catch((error) => {
-                  // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage.
-                  console.error(
-                    "Error processing Claude Agent SDK stream chunk:",
-                    error,
-                  );
-                });
-            },
-            onComplete: () =>
-              state.processing
-                .then(() => finalizeQuerySpan(state))
-                .finally(() => {
-                  spans.delete(event);
-                }),
-            onError: (error: Error) =>
-              state.processing
-                .then(() => {
-                  state.span.log({
-                    error: error.message,
-                  });
-                })
-                .then(() => finalizeQuerySpan(state))
-                .finally(() => {
-                  spans.delete(event);
-                }),
-          });
+        return;
+      }
 
-          return;
-        }
+      try {
+        state.span.log({ output: result });
+      } catch (error) {
+        // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage.
+        console.error("Error extracting output for Claude Agent SDK:", error);
+      } finally {
+        state.span.end();
+      }
+    };
 
+    this.unsubscribers.push(
+      claudeAgentSDKChannels.query.intercept((target, thisArg, args) => {
+        let state: QueryState | undefined;
         try {
-          state.span.log({ output: eventResult });
+          args[0] ??= {};
+          state = startQuery(args[0]);
         } catch (error) {
-          // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage.
-          console.error("Error extracting output for Claude Agent SDK:", error);
-        } finally {
-          state.span.end();
-          spans.delete(event);
+          debugLogger.error(
+            "Error starting Claude Agent SDK instrumentation:",
+            error,
+          );
         }
-      },
 
-      error: (event) => {
-        const state = spans.get(event);
-        if (!state || !event.error) {
-          return;
+        const invokeTarget = () => Reflect.apply(target, thisArg, args);
+        try {
+          const result = state
+            ? runWithClaudeLocalToolContext(
+                invokeTarget,
+                state.localToolParentResolver,
+              )
+            : invokeTarget();
+          if (state) {
+            try {
+              finishQuery(state, result);
+            } catch (error) {
+              debugLogger.error(
+                "Error finalizing Claude Agent SDK instrumentation:",
+                error,
+              );
+            }
+          }
+          return result;
+        } catch (error) {
+          if (state) {
+            try {
+              state.span.log({
+                error: error instanceof Error ? error.message : String(error),
+              });
+              state.span.end();
+            } catch (instrumentationError) {
+              debugLogger.error(
+                "Error handling Claude Agent SDK instrumentation failure:",
+                instrumentationError,
+              );
+            }
+          }
+          throw error;
         }
-
-        state.span.log({
-          error: event.error.message,
-        });
-        state.span.end();
-        spans.delete(event);
-      },
-    };
-
-    channel.subscribe(handlers);
-    this.unsubscribers.push(() => {
-      channel.unsubscribe(handlers);
-    });
+      }),
+    );
   }
 }
diff --git a/js/src/instrumentation/plugins/google-genai-plugin.test.ts b/js/src/instrumentation/plugins/google-genai-plugin.test.ts
index f14690e73..4985e0a65 100644
--- a/js/src/instrumentation/plugins/google-genai-plugin.test.ts
+++ b/js/src/instrumentation/plugins/google-genai-plugin.test.ts
@@ -6,9 +6,6 @@ vi.mock("../../isomorph", () => ({
     newAsyncLocalStorage: vi.fn(() => {
       let current: unknown;
       return {
-        enterWith: vi.fn((store: unknown) => {
-          current = store;
-        }),
         getStore: vi.fn(() => current),
         run: vi.fn((store: unknown, callback: () => unknown) => {
           const previous = current;
diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts b/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts
index 433361974..346901c78 100644
--- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts
+++ b/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts
@@ -99,7 +99,8 @@ describe("PiCodingAgentPlugin", () => {
       return makeStream(finalMessage);
     });
     const agent = makeAgent(originalStreamFn);
-    agent.state.tools = [bashTool()];
+    const tool = bashTool();
+    agent.state.tools = [tool];
     const session = makeSession(agent);
     const context = {
       systemPrompt: "system",
@@ -110,7 +111,7 @@ describe("PiCodingAgentPlugin", () => {
           timestamp: 1,
         },
       ],
-      tools: [bashTool()],
+      tools: [tool],
     };
 
     await interceptor(
@@ -127,6 +128,8 @@ describe("PiCodingAgentPlugin", () => {
           toolName: "bash",
           type: "tool_execution_start",
         });
+        await tool.execute?.("tool-1", { command: "printf pi_tool_ok" });
+        expect(isAutoInstrumentationSuppressed()).toBe(true);
         await this.agent.emit({
           isError: false,
           result: { stdout: "pi_tool_ok" },
@@ -459,6 +462,12 @@ function bashTool() {
   return {
     description: "Run a shell command.",
     name: "bash",
+    execute: vi.fn(async (..._args: unknown[]) => {
+      expect(isAutoInstrumentationSuppressed()).toBe(false);
+      await Promise.resolve();
+      expect(isAutoInstrumentationSuppressed()).toBe(false);
+      return { stdout: "pi_tool_ok" };
+    }),
     parameters: {
       type: "object",
       properties: {
diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts
index da6c50b30..83381b4ac 100644
--- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts
+++ b/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts
@@ -12,7 +12,7 @@ import { getCurrentUnixTimestamp } from "../../util";
 import { SpanTypeAttribute, isObject } from "../../../util/index";
 import { processInputAttachments } from "../../wrappers/attachment-utils";
 import {
-  enterAutoInstrumentationAllowed,
+  runWithAutoInstrumentationAllowed,
   runWithAutoInstrumentationSuppressed,
 } from "../auto-instrumentation-suppression";
 import { piCodingAgentChannels } from "./pi-coding-agent-channels";
@@ -62,7 +62,6 @@ type PiLlmSpanState = {
 };
 
 type PiToolSpanState = {
-  restoreAutoInstrumentation?: () => void;
   span: Span;
 };
 
@@ -73,9 +72,10 @@ type PiAgentPatchState = {
 
 const piAgentPatchStates = new WeakMap<PiAgent, PiAgentPatchState>();
 const piAgentEventSubscriptions = new WeakSet<PiAgent>();
-let piPromptContextStore:
-  | IsoAsyncLocalStorage<PiPromptState | undefined>
-  | undefined;
+const PI_TOOL_EXECUTE_WRAPPED = Symbol.for(
+  "braintrust.pi_coding_agent.tool_execute_wrapped",
+);
+let piPromptContextStore: IsoAsyncLocalStorage<PiPromptState> | undefined;
 
 export class PiCodingAgentPlugin extends BasePlugin {
   private readonly activePromptStates = new Set<PiPromptState>();
@@ -167,6 +167,7 @@ function startPiPromptRun(
     return undefined;
   }
   installPiAgentInstrumentation(agent);
+  wrapPiToolExecutors(agent.state?.tools);
 
   const metadata = {
     ...extractSessionMetadata(session),
@@ -233,10 +234,8 @@ function isPiAgent(value: unknown): value is PiAgent {
   );
 }
 
-function promptContextStore(): IsoAsyncLocalStorage<PiPromptState | undefined> {
-  piPromptContextStore ??= iso.newAsyncLocalStorage<
-    PiPromptState | undefined
-  >();
+function promptContextStore(): IsoAsyncLocalStorage<PiPromptState> {
+  piPromptContextStore ??= iso.newAsyncLocalStorage<PiPromptState>();
   return piPromptContextStore;
 }
 
@@ -299,6 +298,7 @@ function makeInstrumentedStreamFn(
       return invokeOriginal();
     }
 
+    wrapPiToolExecutors(context.tools);
     const llmState = await startPiLlmSpan(state, model, context, options);
     try {
       const stream = await runWithAutoInstrumentationSuppressed(invokeOriginal);
@@ -310,6 +310,41 @@ function makeInstrumentedStreamFn(
   };
 }
 
+function wrapPiToolExecutors(tools: PiTool[] | undefined): void {
+  if (!tools) {
+    return;
+  }
+
+  for (const tool of tools) {
+    try {
+      const execute = tool.execute;
+      if (
+        typeof execute !== "function" ||
+        (execute as typeof execute & { [PI_TOOL_EXECUTE_WRAPPED]?: boolean })[
+          PI_TOOL_EXECUTE_WRAPPED
+        ]
+      ) {
+        continue;
+      }
+
+      const wrappedExecute = function (this: unknown, ...args: unknown[]) {
+        return runWithAutoInstrumentationAllowed(() =>
+          Reflect.apply(execute, this, args),
+        );
+      };
+      Object.defineProperty(wrappedExecute, PI_TOOL_EXECUTE_WRAPPED, {
+        configurable: false,
+        enumerable: false,
+        value: true,
+        writable: false,
+      });
+      tool.execute = wrappedExecute;
+    } catch (error) {
+      logInstrumentationError("Pi Coding Agent tool wrapping", error);
+    }
+  }
+}
+
 async function startPiLlmSpan(
   state: PiPromptState,
   model: PiModel,
@@ -528,35 +563,26 @@ async function startPiToolSpan(
     return;
   }
 
-  const restoreAutoInstrumentation = enterAutoInstrumentationAllowed();
   const metadata = {
     "gen_ai.tool.call.id": event.toolCallId,
     "gen_ai.tool.name": event.toolName,
     "pi_coding_agent.tool.name": event.toolName,
   };
-  try {
-    const span = startBaseSpan(
-      withSpanInstrumentationName(
-        {
-          event: {
-            input: event.args,
-            metadata,
-          },
-          name: event.toolName || "tool",
-          parent: await state.span.export(),
-          spanAttributes: { type: SpanTypeAttribute.TOOL },
+  const span = startBaseSpan(
+    withSpanInstrumentationName(
+      {
+        event: {
+          input: event.args,
+          metadata,
         },
-        INSTRUMENTATION_NAMES.PI_CODING_AGENT,
-      ),
-    );
-    state.activeToolSpans.set(event.toolCallId, {
-      restoreAutoInstrumentation,
-      span,
-    });
-  } catch (error) {
-    restoreAutoInstrumentation();
-    throw error;
-  }
+        name: event.toolName || "tool",
+        parent: await state.span.export(),
+        spanAttributes: { type: SpanTypeAttribute.TOOL },
+      },
+      INSTRUMENTATION_NAMES.PI_CODING_AGENT,
+    ),
+  );
+  state.activeToolSpans.set(event.toolCallId, { span });
 }
 
 function finishPiToolSpan(
@@ -582,11 +608,7 @@ function finishPiToolSpan(
       output: event.result,
     });
   } finally {
-    try {
-      toolState.span.end();
-    } finally {
-      toolState.restoreAutoInstrumentation?.();
-    }
+    toolState.span.end();
   }
 }
 
@@ -671,14 +693,10 @@ function finishPiLlmSpan(
 
 function finishOpenToolSpans(state: PiPromptState, error?: unknown): void {
   for (const [, toolState] of state.activeToolSpans) {
-    try {
-      safeLog(toolState.span, {
-        error: error ? toLoggedError(error) : "Pi tool did not complete",
-      });
-      toolState.span.end();
-    } finally {
-      toolState.restoreAutoInstrumentation?.();
-    }
+    safeLog(toolState.span, {
+      error: error ? toLoggedError(error) : "Pi tool did not complete",
+    });
+    toolState.span.end();
   }
   state.activeToolSpans.clear();
 }
diff --git a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts
index 54340d1a3..a2b3f0527 100644
--- a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts
+++ b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts
@@ -1,35 +1,26 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 
-const {
-  mockBindStore,
-  mockWithCurrent,
-  mockNewAsyncLocalStorage,
-  mockStartSpan,
-  mockUnbindStore,
-} = vi.hoisted(() => ({
-  mockBindStore: vi.fn(),
-  mockWithCurrent: vi.fn(),
-  mockNewAsyncLocalStorage: vi.fn(() => {
-    let current: unknown;
-    return {
-      enterWith: vi.fn((store: unknown) => {
-        current = store;
-      }),
-      getStore: vi.fn(() => current),
-      run: vi.fn((store: unknown, callback: () => unknown) => {
-        const previous = current;
-        current = store;
-        try {
-          return callback();
-        } finally {
-          current = previous;
-        }
-      }),
-    };
+const { mockWithCurrent, mockNewAsyncLocalStorage, mockStartSpan } = vi.hoisted(
+  () => ({
+    mockWithCurrent: vi.fn(),
+    mockNewAsyncLocalStorage: vi.fn(() => {
+      let current: unknown;
+      return {
+        getStore: vi.fn(() => current),
+        run: vi.fn((store: unknown, callback: () => unknown) => {
+          const previous = current;
+          current = store;
+          try {
+            return callback();
+          } finally {
+            current = previous;
+          }
+        }),
+      };
+    }),
+    mockStartSpan: vi.fn(),
   }),
-  mockStartSpan: vi.fn(),
-  mockUnbindStore: vi.fn(),
-}));
+);
 
 vi.mock("../../isomorph", () => ({
   default: {
@@ -75,13 +66,48 @@ describe("StrandsAgentSDKPlugin", () => {
     handlersByName = new Map();
     spans = [];
     mockNewTracingChannel.mockImplementation((name: string) => ({
-      start: {
-        bindStore: mockBindStore,
-        unbindStore: mockUnbindStore,
-      },
-      subscribe: vi.fn((handlers) => handlersByName.set(name, handlers)),
+      intercept: vi.fn((interceptor) => {
+        const handlers = {
+          end: (event: any) =>
+            interceptor(
+              () =>
+                typeof event.invoke === "function"
+                  ? event.invoke()
+                  : event.result,
+              event.self,
+              event.arguments ?? [],
+              {
+                ...(event.agent ? { agent: event.agent } : {}),
+                ...(event.orchestrator
+                  ? { orchestrator: event.orchestrator }
+                  : {}),
+              },
+            ),
+          error: (event: any) => {
+            try {
+              interceptor(
+                () => {
+                  throw event.error;
+                },
+                event.self,
+                event.arguments ?? [],
+                {
+                  ...(event.agent ? { agent: event.agent } : {}),
+                  ...(event.orchestrator
+                    ? { orchestrator: event.orchestrator }
+                    : {}),
+                },
+              );
+            } catch {
+              // The real interceptor preserves the target error.
+            }
+          },
+          start: vi.fn(),
+        };
+        handlersByName.set(name, handlers);
+        return vi.fn();
+      }),
       traceSync: vi.fn((fn) => fn()),
-      unsubscribe: vi.fn(),
     }));
     currentSpan = undefined;
     mockWithCurrent.mockImplementation((span: any, callback: () => unknown) => {
@@ -122,7 +148,7 @@ describe("StrandsAgentSDKPlugin", () => {
     vi.clearAllMocks();
   });
 
-  it("subscribes to Strands stream channels and binds suppression", () => {
+  it("intercepts Strands stream channels", () => {
     const plugin = new StrandsAgentSDKPlugin();
     plugin.enable();
 
@@ -135,11 +161,7 @@ describe("StrandsAgentSDKPlugin", () => {
     expect(
       handlersByName.has("orchestrion:@strands-agents/sdk:Swarm.stream"),
     ).toBe(true);
-    expect(mockBindStore).toHaveBeenCalledTimes(3);
-
     plugin.disable();
-
-    expect(mockUnbindStore).toHaveBeenCalledTimes(3);
   });
 
   it("records agent model and tool spans from stream events", async () => {
@@ -226,6 +248,10 @@ describe("StrandsAgentSDKPlugin", () => {
     );
     const event = {
       arguments: ["hello", undefined],
+      invoke: () => {
+        suppressionStates.push(isAutoInstrumentationSuppressed());
+        return stream;
+      },
       moduleVersion: "1.6.0",
       result: stream,
       self: agent,
@@ -239,7 +265,15 @@ describe("StrandsAgentSDKPlugin", () => {
     }
 
     expect(chunks).toHaveLength(6);
-    expect(suppressionStates).toEqual([true, true, true, true, true, true]);
+    expect(suppressionStates).toEqual([
+      true,
+      true,
+      true,
+      true,
+      true,
+      true,
+      true,
+    ]);
     const rootSpan = spans.find((span) => span.args.name === "Agent: helper");
     const modelSpan = spans.find(
       (span) => span.args.name === "Strands model: gpt-4o-mini",
diff --git a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts
index 255578a13..13a434cff 100644
--- a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts
+++ b/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts
@@ -1,7 +1,5 @@
 import { BasePlugin, toLoggedError } from "../core";
-import type { ChannelMessage } from "../core/channel-definitions";
 import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher";
-import type { IsoChannelHandlers } from "../../isomorph";
 import { debugLogger } from "../../debug-logger";
 import {
   Attachment,
@@ -18,10 +16,7 @@ import { LRUCache } from "../../lru-cache";
 import { getCurrentUnixTimestamp } from "../../util";
 import { SpanTypeAttribute, isObject } from "../../../util/index";
 import { convertDataToBlob } from "../../wrappers/attachment-utils";
-import {
-  bindAutoInstrumentationSuppressionToStart,
-  runWithAutoInstrumentationSuppressed,
-} from "../auto-instrumentation-suppression";
+import { runWithAutoInstrumentationSuppressed } from "../auto-instrumentation-suppression";
 import { strandsAgentSDKChannels } from "./strands-agent-sdk-channels";
 import type {
   StrandsAfterModelCallEvent,
@@ -105,12 +100,12 @@ export class StrandsAgentSDKPlugin extends BasePlugin {
   private readonly activeChildParents: ActiveChildParents = new WeakMap();
 
   protected onEnable(): void {
-    this.subscribeToAgentStream();
-    this.subscribeToMultiAgentStream(
+    this.interceptAgentStream();
+    this.interceptMultiAgentStream(
       strandsAgentSDKChannels.graphStream,
       "Graph.stream",
     );
-    this.subscribeToMultiAgentStream(
+    this.interceptMultiAgentStream(
       strandsAgentSDKChannels.swarmStream,
       "Swarm.stream",
     );
@@ -123,145 +118,124 @@ export class StrandsAgentSDKPlugin extends BasePlugin {
     this.unsubscribers = [];
   }
 
-  private subscribeToAgentStream(): void {
-    const channel = strandsAgentSDKChannels.agentStream.tracingChannel();
-    const states = new WeakMap<object, AgentStreamState>();
-    const unbindAutoInstrumentationSuppression =
-      bindAutoInstrumentationSuppressionToStart(channel);
-
-    const handlers: IsoChannelHandlers<
-      ChannelMessage<typeof strandsAgentSDKChannels.agentStream>
-    > = {
-      start: (event) => {
-        const state = startAgentStream(event, this.activeChildParents);
-        if (state) {
-          states.set(event, state);
-        }
-      },
-      end: (event) => {
-        const state = states.get(event);
-        if (!state) {
-          return;
-        }
-
-        const result = event.result;
-        if (isAsyncIterable(result)) {
-          patchStreamIfNeeded<StrandsAgentStreamEvent>(result, {
-            aroundNext: (callback) =>
-              runWithAutoInstrumentationSuppressed(callback),
-            onChunk: (chunk) => handleAgentStreamEvent(state, chunk),
-            onComplete: () => {
-              finalizeAgentStream(state);
-              states.delete(event);
-            },
-            onError: (error) => {
-              finalizeAgentStream(state, error);
-              states.delete(event);
-            },
-          });
-          return;
-        }
-
-        finalizeAgentStream(state, undefined, result);
-        states.delete(event);
-      },
-      error: (event) => {
-        const state = states.get(event);
-        if (!state || !event.error) {
-          return;
-        }
-        finalizeAgentStream(state, event.error);
-        states.delete(event);
-      },
-    };
-
-    channel.subscribe(handlers);
-    this.unsubscribers.push(() => {
-      unbindAutoInstrumentationSuppression?.();
-      channel.unsubscribe(handlers);
-    });
+  private interceptAgentStream(): void {
+    this.unsubscribers.push(
+      strandsAgentSDKChannels.agentStream.intercept(
+        (target, thisArg, args, additional) =>
+          instrumentStrandsStreamInvocation<
+            AgentStreamState,
+            StrandsAgentStreamEvent,
+            ReturnType<typeof target>
+          >({
+            finalize: finalizeAgentStream,
+            handleChunk: handleAgentStreamEvent,
+            invoke: () => Reflect.apply(target, thisArg, args),
+            name: "Strands Agent SDK",
+            start: () =>
+              startAgentStream(
+                args[0],
+                extractAgent(additional.agent, thisArg),
+                this.activeChildParents,
+              ),
+          }),
+      ),
+    );
   }
 
-  private subscribeToMultiAgentStream(
+  private interceptMultiAgentStream(
     channel: MultiAgentStreamChannel,
     operation: MultiAgentStreamState["operation"],
   ): void {
-    const tracingChannel = channel.tracingChannel();
-    const states = new WeakMap<object, MultiAgentStreamState>();
-    const unbindAutoInstrumentationSuppression =
-      bindAutoInstrumentationSuppressionToStart(tracingChannel);
-
-    const handlers: IsoChannelHandlers<ChannelMessage<typeof channel>> = {
-      start: (event) => {
-        const state = startMultiAgentStream(
-          event,
-          operation,
-          this.activeChildParents,
-        );
-        if (state) {
-          states.set(event, state);
-        }
-      },
-      end: (event) => {
-        const state = states.get(event);
-        if (!state) {
-          return;
-        }
+    this.unsubscribers.push(
+      channel.intercept((target, thisArg, args, additional) =>
+        instrumentStrandsStreamInvocation<
+          MultiAgentStreamState,
+          StrandsMultiAgentStreamEvent,
+          ReturnType<typeof target>
+        >({
+          finalize: (state, error, output) =>
+            finalizeMultiAgentStream(
+              state,
+              this.activeChildParents,
+              error,
+              output,
+            ),
+          handleChunk: (state, chunk) =>
+            handleMultiAgentStreamEvent(state, chunk, this.activeChildParents),
+          invoke: () => Reflect.apply(target, thisArg, args),
+          name: "Strands multi-agent",
+          start: () =>
+            startMultiAgentStream(
+              args[0],
+              extractOrchestrator(additional.orchestrator, thisArg),
+              operation,
+              this.activeChildParents,
+            ),
+        }),
+      ),
+    );
+  }
+}
 
-        const result = event.result;
-        if (isAsyncIterable(result)) {
-          patchStreamIfNeeded<StrandsMultiAgentStreamEvent>(result, {
-            aroundNext: (callback) =>
-              runWithAutoInstrumentationSuppressed(callback),
-            onChunk: (chunk) =>
-              handleMultiAgentStreamEvent(
-                state,
-                chunk,
-                this.activeChildParents,
-              ),
-            onComplete: () => {
-              finalizeMultiAgentStream(state, this.activeChildParents);
-              states.delete(event);
-            },
-            onError: (error) => {
-              finalizeMultiAgentStream(state, this.activeChildParents, error);
-              states.delete(event);
-            },
-          });
-          return;
-        }
+function instrumentStrandsStreamInvocation<TState, TChunk, TResult>(options: {
+  finalize: (state: TState, error?: unknown, output?: unknown) => void;
+  handleChunk: (state: TState, chunk: TChunk) => void;
+  invoke: () => TResult;
+  name: string;
+  start: () => TState;
+}): TResult {
+  let state: TState | undefined;
+  try {
+    state = options.start();
+  } catch (error) {
+    debugLogger.error(`Error starting ${options.name} instrumentation:`, error);
+  }
 
-        finalizeMultiAgentStream(
-          state,
-          this.activeChildParents,
-          undefined,
-          result,
+  let result: TResult;
+  try {
+    result = runWithAutoInstrumentationSuppressed(options.invoke);
+  } catch (error) {
+    if (state) {
+      try {
+        options.finalize(state, error);
+      } catch (instrumentationError) {
+        debugLogger.error(
+          `Error handling ${options.name} instrumentation failure:`,
+          instrumentationError,
         );
-        states.delete(event);
-      },
-      error: (event) => {
-        const state = states.get(event);
-        if (!state || !event.error) {
-          return;
-        }
-        finalizeMultiAgentStream(state, this.activeChildParents, event.error);
-        states.delete(event);
-      },
-    };
-
-    tracingChannel.subscribe(handlers);
-    this.unsubscribers.push(() => {
-      unbindAutoInstrumentationSuppression?.();
-      tracingChannel.unsubscribe(handlers);
-    });
+      }
+    }
+    throw error;
+  }
+
+  if (state) {
+    try {
+      if (isAsyncIterable(result)) {
+        patchStreamIfNeeded<TChunk>(result, {
+          aroundNext: (callback) =>
+            runWithAutoInstrumentationSuppressed(callback),
+          onChunk: (chunk) => options.handleChunk(state, chunk),
+          onComplete: () => options.finalize(state),
+          onError: (error) => options.finalize(state, error),
+        });
+      } else {
+        options.finalize(state, undefined, result);
+      }
+    } catch (error) {
+      debugLogger.error(
+        `Error finalizing ${options.name} instrumentation:`,
+        error,
+      );
+    }
   }
+  return result;
 }
 
 function startAgentStream(
-  event: ChannelMessage<typeof strandsAgentSDKChannels.agentStream>,
+  input: unknown,
+  agent: StrandsAgent | undefined,
   activeChildParents: ActiveChildParents,
-): AgentStreamState | undefined {
-  const agent = extractAgent(event);
+): AgentStreamState {
   const model = agent?.model;
   const metadata = {
     ...extractAgentMetadata(agent),
@@ -273,17 +247,14 @@ function startAgentStream(
     ? getOnlyChildParent(activeChildParents, agent)
     : undefined;
   const attachmentCache = createStrandsAttachmentCache();
-  const input = processStrandsInputAttachments(
-    event.arguments[0],
-    attachmentCache,
-  );
+  const processedInput = processStrandsInputAttachments(input, attachmentCache);
   const span = parentSpan
     ? withCurrent(parentSpan, () =>
         startBaseSpan(
           withSpanInstrumentationName(
             {
               event: {
-                input,
+                input: processedInput,
                 metadata,
               },
               name: formatAgentSpanName(agent),
@@ -297,7 +268,7 @@ function startAgentStream(
         withSpanInstrumentationName(
           {
             event: {
-              input,
+              input: processedInput,
               metadata,
             },
             name: formatAgentSpanName(agent),
@@ -318,11 +289,11 @@ function startAgentStream(
 }
 
 function startMultiAgentStream(
-  event: ChannelMessage<MultiAgentStreamChannel>,
+  input: unknown,
+  orchestrator: StrandsMultiAgent | undefined,
   operation: MultiAgentStreamState["operation"],
   activeChildParents: ActiveChildParents,
 ): MultiAgentStreamState {
-  const orchestrator = extractOrchestrator(event);
   const metadata = {
     "strands.operation": operation,
     provider: "strands",
@@ -331,14 +302,14 @@ function startMultiAgentStream(
   const parentSpan = orchestrator
     ? getOnlyChildParent(activeChildParents, orchestrator)
     : undefined;
-  const input = processStrandsInputAttachments(event.arguments[0]);
+  const processedInput = processStrandsInputAttachments(input);
   const span = parentSpan
     ? withCurrent(parentSpan, () =>
         startBaseSpan(
           withSpanInstrumentationName(
             {
               event: {
-                input,
+                input: processedInput,
                 metadata,
               },
               name:
@@ -355,7 +326,7 @@ function startMultiAgentStream(
         withSpanInstrumentationName(
           {
             event: {
-              input,
+              input: processedInput,
               metadata,
             },
             name:
@@ -847,19 +818,18 @@ function finalizeMultiAgentStream(
   state.span.end();
 }
 
-function extractAgent(
-  event: ChannelMessage<typeof strandsAgentSDKChannels.agentStream>,
-): StrandsAgent | undefined {
-  const candidate = event.agent ?? event.self;
+function extractAgent(agent: unknown, self: unknown): StrandsAgent | undefined {
+  const candidate = agent ?? self;
   return isObject(candidate) && typeof candidate.stream === "function"
     ? (candidate as StrandsAgent)
     : undefined;
 }
 
 function extractOrchestrator(
-  event: ChannelMessage<MultiAgentStreamChannel>,
+  orchestrator: unknown,
+  self: unknown,
 ): StrandsMultiAgent | undefined {
-  const candidate = event.orchestrator ?? event.self;
+  const candidate = orchestrator ?? self;
   return isObject(candidate) && typeof candidate.stream === "function"
     ? (candidate as StrandsMultiAgent)
     : undefined;
diff --git a/js/src/instrumentation/registry.test.ts b/js/src/instrumentation/registry.test.ts
index fdd49a5c2..9d01c94ef 100644
--- a/js/src/instrumentation/registry.test.ts
+++ b/js/src/instrumentation/registry.test.ts
@@ -5,7 +5,6 @@ vi.mock("../isomorph", () => ({
   default: {
     newTracingChannel: vi.fn(),
     newAsyncLocalStorage: vi.fn(() => ({
-      enterWith: vi.fn(),
       getStore: vi.fn(() => undefined),
       run: vi.fn((_store: unknown, callback: () => unknown) => callback()),
     })),
diff --git a/js/src/isomorph.ts b/js/src/isomorph.ts
index d23ddcc7e..a6adf3b48 100644
--- a/js/src/isomorph.ts
+++ b/js/src/isomorph.ts
@@ -21,7 +21,6 @@ export type IsoAsyncLocalStorage<T> = GlobalHookAsyncLocalStorage<T>;
 class DefaultAsyncLocalStorage<T> implements IsoAsyncLocalStorage<T> {
   constructor() {}
 
-  enterWith(_: T): void {}
   run<R>(_: T | undefined, callback: () => R): R {
     return callback();
   }
diff --git a/js/src/vendor-sdk-types/pi-coding-agent.ts b/js/src/vendor-sdk-types/pi-coding-agent.ts
index 13df2a1b9..0c2810da5 100644
--- a/js/src/vendor-sdk-types/pi-coding-agent.ts
+++ b/js/src/vendor-sdk-types/pi-coding-agent.ts
@@ -139,6 +139,7 @@ export interface PiContext {
 export interface PiTool {
   name: string;
   description?: string;
+  execute?: (this: unknown, ...args: unknown[]) => unknown;
   parameters?: unknown;
   [key: string]: unknown;
 }
diff --git a/js/src/wrappers/ai-sdk/telemetry.ts b/js/src/wrappers/ai-sdk/telemetry.ts
index 1fb2c9812..5bb41f885 100644
--- a/js/src/wrappers/ai-sdk/telemetry.ts
+++ b/js/src/wrappers/ai-sdk/telemetry.ts
@@ -28,7 +28,6 @@ import type {
   AISDKResult,
   AISDKRerankResult,
 } from "../../vendor-sdk-types/ai-sdk";
-import iso from "../../isomorph";
 import type {
   AISDKV7LanguageModelCallStartEvent,
   AISDKV7OperationEvent,
@@ -74,9 +73,6 @@ type EmbedSpanState = CallSpanState & {
 export function braintrustAISDKTelemetry(): any {
   const operations = new Map<string, OperationState>();
   const operationKeysByCallId = new Map<string, string[]>();
-  const workflowOperationKeyStore = iso.newAsyncLocalStorage<
-    string | undefined
-  >();
   const modelSpans = new Map<string, Span[]>();
   const objectSpans = new Map<string, Span>();
   const embedSpans = new Map<string, EmbedSpanState>();
@@ -135,13 +131,6 @@ export function braintrustAISDKTelemetry(): any {
     }
 
     operations.delete(operationKey);
-    if (workflowOperationKeyStore.getStore() === operationKey) {
-      // TODO(luca): Replace ALS.enterWith() with ALS.run() once direct
-      // telemetry can wrap the full WorkflowAgent callback lifecycle.
-      // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above.
-      workflowOperationKeyStore.enterWith(undefined);
-    }
-
     const keys = operationKeysByCallId.get(state.callId);
     if (!keys) {
       return;
@@ -208,16 +197,12 @@ export function braintrustAISDKTelemetry(): any {
       }
     }
 
-    const workflowOperationKey = workflowOperationKeyStore.getStore();
-    if (workflowOperationKey && keys.includes(workflowOperationKey)) {
-      return workflowOperationKey;
-    }
-
-    if (callId === "workflow-agent") {
-      return undefined;
-    }
-
-    return mode === "finish" ? keys[0] : keys[keys.length - 1];
+    // Dispatcher instrumentation stamps an explicit key on every callback.
+    // Direct registerTelemetry() usage has no callback boundary to carry one,
+    // so overlapping operations with the same callId are best-effort.
+    return callId === "workflow-agent" || mode === "active"
+      ? keys[keys.length - 1]
+      : keys[0];
   };
 
   const operationKeyFromEvent = (
@@ -237,20 +222,19 @@ export function braintrustAISDKTelemetry(): any {
           return operationKey;
         }
 
-        const workflowOperationKey = workflowOperationKeyStore.getStore();
-        if (workflowOperationKey && operations.has(workflowOperationKey)) {
-          return workflowOperationKey;
+        // Some direct WorkflowAgent telemetry callbacks use a child callId
+        // instead of the operation's shared `workflow-agent` callId. Without
+        // a dispatcher key, route these to the newest active workflow as a
+        // deterministic best-effort fallback.
+        const workflowAgentKeys = operationKeysByCallId.get("workflow-agent");
+        if (workflowAgentKeys?.length) {
+          return workflowAgentKeys[workflowAgentKeys.length - 1];
         }
 
         return callId === "workflow-agent" ? undefined : callId;
       }
     }
 
-    const workflowOperationKey = workflowOperationKeyStore.getStore();
-    if (workflowOperationKey && operations.has(workflowOperationKey)) {
-      return workflowOperationKey;
-    }
-
     const wrapperSpan = currentWorkflowAgentWrapperSpan();
     if (wrapperSpan?.spanId) {
       for (const [operationKey, state] of operations) {
@@ -266,8 +250,8 @@ export function braintrustAISDKTelemetry(): any {
     // WorkflowAgent uses this callId on the operation, but omits it from
     // tool start/end callbacks in @ai-sdk/workflow@1.0.x.
     const workflowAgentKeys = operationKeysByCallId.get("workflow-agent");
-    if (workflowAgentKeys?.length === 1) {
-      return workflowAgentKeys[0];
+    if (workflowAgentKeys?.length) {
+      return workflowAgentKeys[workflowAgentKeys.length - 1];
     }
 
     if (operations.size === 1) {
@@ -521,15 +505,6 @@ export function braintrustAISDKTelemetry(): any {
           return;
         }
 
-        if (workflowAgent) {
-          // Direct registerTelemetry() calls do not receive the hidden
-          // dispatcher operation key used by auto-instrumentation.
-          // TODO(luca): Replace ALS.enterWith() with ALS.run() once direct
-          // telemetry can wrap the full WorkflowAgent callback lifecycle.
-          // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above.
-          workflowOperationKeyStore.enterWith(operationKey);
-        }
-
         let metadata = metadataFromEvent(event);
         const logPayload: {
           input?: unknown;
diff --git a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts
index db24f0bcf..ddc6d678e 100644
--- a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts
+++ b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts
@@ -54,11 +54,11 @@ function wrapClaudeAgentQuery(
         thisArg === proxy || thisArg === undefined
           ? (defaultThis ?? thisArg)
           : thisArg;
-      return claudeAgentSDKChannels.query.traceSync(
-        () => Reflect.apply(target, invocationTarget, [wrappedParams]),
-        // The channel carries no extra context fields, but the generated
-        // StartOf<> type for Record<string, never> is overly strict here.
-        { arguments: [wrappedParams] } as never,
+      return claudeAgentSDKChannels.query.invoke(
+        target,
+        invocationTarget,
+        [wrappedParams],
+        {},
       );
     },
   });
diff --git a/js/src/wrappers/strands-agent-sdk.test.ts b/js/src/wrappers/strands-agent-sdk.test.ts
index f811800d4..66d2ad46c 100644
--- a/js/src/wrappers/strands-agent-sdk.test.ts
+++ b/js/src/wrappers/strands-agent-sdk.test.ts
@@ -1,15 +1,20 @@
 import { afterEach, describe, expect, it, vi } from "vitest";
 
-const { traceSync } = vi.hoisted(() => ({
-  traceSync: vi.fn((fn: () => unknown, _event?: unknown) => fn()),
+const { invoke } = vi.hoisted(() => ({
+  invoke: vi.fn(
+    (
+      target: Function,
+      thisArg: unknown,
+      args: unknown[],
+      _additional?: unknown,
+    ) => Reflect.apply(target, thisArg, args),
+  ),
 }));
 
 vi.mock("../isomorph", () => ({
   default: {
     newTracingChannel: vi.fn(() => ({
-      subscribe: vi.fn(),
-      traceSync,
-      unsubscribe: vi.fn(),
+      invoke,
     })),
   },
 }));
@@ -61,14 +66,12 @@ describe("wrapStrandsAgentSDK", () => {
     expect(result).toMatchObject({
       lastMessage: { role: "assistant", content: [{ text: "world" }] },
     });
-    expect(traceSync).toHaveBeenCalledTimes(2);
-    expect(traceSync.mock.calls[0][1]).toMatchObject({
-      arguments: ["hello", undefined],
-      self: expect.objectContaining({ name: "assistant" }),
-    });
-    expect(traceSync.mock.calls[1][1]).toMatchObject({
-      arguments: ["world", undefined],
+    expect(invoke).toHaveBeenCalledTimes(2);
+    expect(invoke.mock.calls[0][2]).toEqual(["hello", undefined]);
+    expect(invoke.mock.calls[0][3]).toMatchObject({
+      agent: expect.objectContaining({ name: "assistant" }),
     });
+    expect(invoke.mock.calls[1][2]).toEqual(["world", undefined]);
   });
 
   it("wraps Graph and Swarm stream/invoke", async () => {
@@ -101,12 +104,14 @@ describe("wrapStrandsAgentSDK", () => {
     await expect(new wrapped.Swarm().invoke("swarm")).resolves.toMatchObject({
       status: "COMPLETED",
     });
-    expect(traceSync).toHaveBeenCalledTimes(2);
-    expect(traceSync.mock.calls[0][1]).toMatchObject({
-      arguments: ["graph", undefined],
+    expect(invoke).toHaveBeenCalledTimes(2);
+    expect(invoke.mock.calls[0][2]).toEqual(["graph", undefined]);
+    expect(invoke.mock.calls[0][3]).toMatchObject({
+      orchestrator: expect.anything(),
     });
-    expect(traceSync.mock.calls[1][1]).toMatchObject({
-      arguments: ["swarm", undefined],
+    expect(invoke.mock.calls[1][2]).toEqual(["swarm", undefined]);
+    expect(invoke.mock.calls[1][3]).toMatchObject({
+      orchestrator: expect.anything(),
     });
   });
 
@@ -126,7 +131,7 @@ describe("wrapStrandsAgentSDK", () => {
     ) as any;
 
     await new wrapped.Agent().invoke("hello");
-    expect(traceSync).toHaveBeenCalledTimes(1);
+    expect(invoke).toHaveBeenCalledTimes(1);
   });
 
   it("preserves private-field-safe method binding", async () => {
diff --git a/js/src/wrappers/strands-agent-sdk.ts b/js/src/wrappers/strands-agent-sdk.ts
index 6c6173034..5cf339789 100644
--- a/js/src/wrappers/strands-agent-sdk.ts
+++ b/js/src/wrappers/strands-agent-sdk.ts
@@ -21,8 +21,8 @@ const WRAPPED_INSTANCE = Symbol.for(
 );
 
 /**
- * Wraps the Strands Agent SDK with Braintrust tracing. The wrapper emits
- * diagnostics-channel events; the Strands plugin owns span lifecycle.
+ * Wraps the Strands Agent SDK with Braintrust tracing. The wrapper invokes
+ * typed instrumentation channels; the Strands plugin owns span lifecycle.
  */
 export function wrapStrandsAgentSDK<T>(sdk: T): T {
   if (!sdk || typeof sdk !== "object") {
@@ -157,13 +157,11 @@ function wrapAgentInstance(agent: StrandsAgent): StrandsAgent {
             StrandsInvokeArgs,
             StrandsInvokeOptions | undefined,
           ];
-          return strandsAgentSDKChannels.agentStream.traceSync(
-            () => Reflect.apply(value, target, callArgs),
-            {
-              agent: proxy,
-              arguments: callArgs,
-              self: proxy,
-            } as never,
+          return strandsAgentSDKChannels.agentStream.invoke(
+            value as StrandsAgent["stream"],
+            target,
+            callArgs,
+            { agent: proxy },
           );
         };
       }
@@ -219,13 +217,13 @@ function wrapMultiAgentInstance(
             kind === "graph"
               ? strandsAgentSDKChannels.graphStream
               : strandsAgentSDKChannels.swarmStream;
-          return channel.traceSync(
-            () => Reflect.apply(value, target, callArgs),
+          return channel.invoke(
+            value as StrandsMultiAgent["stream"],
+            target,
+            callArgs,
             {
-              arguments: callArgs,
               orchestrator: proxy,
-              self: proxy,
-            } as never,
+            },
           );
         };
       }
diff --git a/js/src/wrappers/vitest/context-manager.ts b/js/src/wrappers/vitest/context-manager.ts
index 03cb41799..da25ee17c 100644
--- a/js/src/wrappers/vitest/context-manager.ts
+++ b/js/src/wrappers/vitest/context-manager.ts
@@ -54,12 +54,6 @@ class VitestContextManager {
     return this.contextStorage.getStore();
   }
 
-  setContext(context: VitestExperimentContext): void {
-    // TODO(luca): Replace ALS.enterWith() with ALS.run()
-    // eslint-disable-next-line no-restricted-syntax -- Existing ALS.enterWith() usage tracked by the TODO above.
-    this.contextStorage.enterWith(context);
-  }
-
   runInContext<R>(context: VitestExperimentContext, callback: () => R): R {
     return this.contextStorage.run(context, callback);
   }
diff --git a/js/src/wrappers/vitest/wrapper.ts b/js/src/wrappers/vitest/wrapper.ts
index 4be2adf5b..b3c7344fa 100644
--- a/js/src/wrappers/vitest/wrapper.ts
+++ b/js/src/wrappers/vitest/wrapper.ts
@@ -87,8 +87,8 @@ export function wrapTest<VitestContext = unknown>(
 
       // Capture context at registration time (during wrapDescribe factory execution)
       // as a fallback. Vitest's async test runner creates new async contexts for
-      // each test, so AsyncLocalStorage.enterWith() set in the describe factory
-      // doesn't propagate to test execution. The captured context is used when
+      // each test, so the describe factory's AsyncLocalStorage context doesn't
+      // propagate to test execution. The captured context is used when
       // getExperimentContext() returns null at runtime.
       const registrationContext = getExperimentContext();
 
@@ -259,9 +259,7 @@ export function wrapDescribe(
           config.onProgress({ type: "suite_start", suiteName });
         }
 
-        contextManager.setContext(lazyContext);
-
-        factory();
+        contextManager.runInContext(lazyContext, factory);
 
         if (afterAll) {
           afterAll(async () => {