fix: Add tracing polyfill to workerd and edge-light runtimes

js/src/edge-light/config.ts

@@ -1,6 +1,13 @@
 import iso from "../isomorph";
+import type {
+  IsoTracingChannel,
+  IsoTracingChannelCollection,
+} from "../isomorph";
 import { _internalSetInitialState } from "../logger";
 import { resolveRuntimeAsyncLocalStorage } from "../runtime-async-local-storage";
+import { tracingChannel } from "dc-browser";
+import { patchTracingChannel } from "../auto-instrumentations/patch-tracing-channel";
+import { registry } from "../instrumentation/registry";
 
 let edgeLightConfigured = false;
 
@@ -20,6 +27,13 @@ export function configureEdgeLight(): void {
     iso.newAsyncLocalStorage = <T>() => new runtimeAsyncLocalStorage<T>();
   }
 
+  iso.newTracingChannel = <M = unknown>(
+    nameOrChannels: string | IsoTracingChannelCollection<M>,
+  ) =>
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    tracingChannel(nameOrChannels as string | object) as IsoTracingChannel<M>;
+  patchTracingChannel(tracingChannel);
+
   iso.getEnv = (name: string) => {
     if (typeof process === "undefined" || typeof process.env === "undefined") {
       return undefined;
@@ -39,5 +53,6 @@ export function configureEdgeLight(): void {
   };
 
   _internalSetInitialState();
+  registry.enable();
   edgeLightConfigured = true;
 }

js/src/edge-runtime-bootstrap.test.ts

@@ -108,5 +108,44 @@ describe.each([
       expect(result.childRootSpanId).toBe(root.spanId);
       expect(await backgroundLogger.drain()).toHaveLength(3);
     });
+
+    test("wrapAISDK logs spans in edge runtimes", async () => {
+      const braintrust = await import(entrypoint);
+
+      braintrust._exportsForTestingOnly.setInitialTestState();
+      await braintrust._exportsForTestingOnly.simulateLoginForTests();
+
+      const backgroundLogger =
+        braintrust._exportsForTestingOnly.useTestBackgroundLogger();
+      braintrust.initLogger({
+        projectId: "test-project-id",
+        projectName,
+      });
+
+      await braintrust.traced(async () => "ok", { name: "root" });
+
+      const wrapped = braintrust.wrapAISDK({
+        generateText: async () => ({
+          finishReason: "stop",
+          text: "ok",
+          usage: {
+            completionTokens: 1,
+            promptTokens: 1,
+            totalTokens: 2,
+          },
+        }),
+      });
+
+      await wrapped.generateText({
+        model: { modelId: "fake-model", provider: "fake-provider" },
+        prompt: "hello",
+      });
+
+      const spans = await backgroundLogger.drain();
+      const spanNames = spans.map((span) => span.span_attributes?.name);
+
+      expect(spanNames).toContain("root");
+      expect(spanNames).toContain("generateText");
+    });
   },
 );

js/src/workerd/config.ts

@@ -1,6 +1,13 @@
 import iso from "../isomorph";
+import type {
+  IsoTracingChannel,
+  IsoTracingChannelCollection,
+} from "../isomorph";
 import { _internalSetInitialState } from "../logger";
 import { resolveRuntimeAsyncLocalStorage } from "../runtime-async-local-storage";
+import { tracingChannel } from "dc-browser";
+import { patchTracingChannel } from "../auto-instrumentations/patch-tracing-channel";
+import { registry } from "../instrumentation/registry";
 
 let workerdConfigured = false;
 
@@ -20,6 +27,13 @@ export function configureWorkerd(): void {
     iso.newAsyncLocalStorage = <T>() => new runtimeAsyncLocalStorage<T>();
   }
 
+  iso.newTracingChannel = <M = unknown>(
+    nameOrChannels: string | IsoTracingChannelCollection<M>,
+  ) =>
+    // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
+    tracingChannel(nameOrChannels as string | object) as IsoTracingChannel<M>;
+  patchTracingChannel(tracingChannel);
+
   iso.getEnv = (name: string) => {
     if (typeof process === "undefined" || typeof process.env === "undefined") {
       return undefined;
@@ -39,5 +53,6 @@ export function configureWorkerd(): void {
   };
 
   _internalSetInitialState();
+  registry.enable();
   workerdConfigured = true;
 }