diff --git a/js/src/edge-light/config.ts b/js/src/edge-light/config.ts index 69287ef63..c32c40635 100644 --- a/js/src/edge-light/config.ts +++ b/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 = () => new runtimeAsyncLocalStorage(); } + iso.newTracingChannel = ( + nameOrChannels: string | IsoTracingChannelCollection, + ) => + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + tracingChannel(nameOrChannels as string | object) as IsoTracingChannel; + 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; } diff --git a/js/src/edge-runtime-bootstrap.test.ts b/js/src/edge-runtime-bootstrap.test.ts index 6dc61cafa..b3e81a2ac 100644 --- a/js/src/edge-runtime-bootstrap.test.ts +++ b/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"); + }); }, ); diff --git a/js/src/workerd/config.ts b/js/src/workerd/config.ts index 3bdba2a9d..cc4f40b7a 100644 --- a/js/src/workerd/config.ts +++ b/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 = () => new runtimeAsyncLocalStorage(); } + iso.newTracingChannel = ( + nameOrChannels: string | IsoTracingChannelCollection, + ) => + // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + tracingChannel(nameOrChannels as string | object) as IsoTracingChannel; + 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; }