ref: Remove AsyncLocalStorage.enterWith() usage

.changeset/remove-enter-with.md

@@ -0,0 +1,5 @@
+---
+"braintrust": patch
+---
+
+ref: Remove `AsyncLocalStorage.enterWith()` usage

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;
       },

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;
 }

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);
     });
   });

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);
 }

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: {} }],

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;

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()),
     })),

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",