From 61fafeabde35788d00c9db59870f45ad1a76c3eb Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:22:17 +0000 Subject: [PATCH 1/2] wip ref: Replace diagnostic channels with proprietary global hooks --- .../template/app/api/test/route.ts | 33 +- .../template/scenario.ts | 4 +- .../scenario.impl.mjs | 2 +- js/NOTICE | 6 + js/licenses/node-diagnostics-channel/LICENSE | 19 + js/package.json | 1 - js/src/auto-instrumentations/README.md | 566 +------ js/src/auto-instrumentations/bundler/next.ts | 41 +- .../auto-instrumentations/bundler/plugin.ts | 27 +- .../bundler/webpack-loader.ts | 5 +- js/src/auto-instrumentations/hook.mts | 12 +- js/src/auto-instrumentations/index.ts | 2 +- .../auto-instrumentations/loader/cjs-patch.ts | 2 +- .../loader/special-case-patches.ts | 2 +- .../orchestrion-js/index.ts | 3 +- .../orchestrion-js/matcher.ts | 5 +- .../orchestrion-js/transformer.ts | 8 +- .../orchestrion-js/transforms.ts | 196 +-- .../orchestrion-js/types.ts | 2 +- .../patch-tracing-channel.test.ts | 255 --- .../patch-tracing-channel.ts | 146 -- js/src/browser/config.ts | 9 - js/src/edge-light/config.ts | 13 - js/src/global-instrumentation-hooks.test.ts | 232 +++ js/src/global-instrumentation-hooks.ts | 474 ++++++ js/src/instrumentation/README.md | 1471 ++--------------- js/src/instrumentation/braintrust-plugin.ts | 2 +- js/src/instrumentation/core/channel.ts | 2 +- js/src/instrumentation/core/plugin.ts | 2 +- js/src/instrumentation/core/stream-patcher.ts | 2 +- js/src/instrumentation/core/types.ts | 6 +- js/src/instrumentation/index.ts | 4 +- .../plugins/genkit-plugin.test.ts | 16 +- .../plugins/langchain-plugin.test.ts | 7 - js/src/instrumentation/registry.test.ts | 4 +- js/src/instrumentation/registry.ts | 4 +- js/src/isomorph.ts | 220 +-- js/src/logger.ts | 4 +- .../node/apply-auto-instrumentation-entry.ts | 4 - js/src/node/config.ts | 8 - js/src/workerd/config.ts | 13 - .../error-handling.test.ts | 4 +- .../event-content.test.ts | 6 +- .../fixtures/global-hook-listener.cjs | 176 ++ .../fixtures/listener-cjs.cjs | 6 +- .../fixtures/listener-esm.mjs | 20 +- .../orchestrion-js/arguments_mutation/test.js | 7 +- .../orchestrion-js/common/preamble.js | 153 +- .../orchestrion-js/polyfill_cjs/mod.js | 9 - .../orchestrion-js/polyfill_cjs/polyfill.js | 8 - .../orchestrion-js/polyfill_cjs/test.js | 19 - .../orchestrion-js/polyfill_mjs/mod.mjs | 9 - .../orchestrion-js/polyfill_mjs/polyfill.js | 9 - .../orchestrion-js/polyfill_mjs/test.mjs | 17 - .../test-api-promise-preservation.mjs | 45 - .../function-behavior.test.ts | 2 - .../auto-instrumentations/loader-hook.test.ts | 28 +- .../multiple-instrumentations.test.ts | 2 - .../auto-instrumentations/next-config.test.ts | 14 +- .../orchestrion-js-upstream.test.ts | 17 +- .../runtime-execution.test.ts | 8 - .../auto-instrumentations/test-helpers.ts | 30 +- .../transformation.test.ts | 165 +- js/tsup.config.ts | 2 +- pnpm-lock.yaml | 8 - 65 files changed, 1496 insertions(+), 3102 deletions(-) create mode 100644 js/licenses/node-diagnostics-channel/LICENSE delete mode 100644 js/src/auto-instrumentations/patch-tracing-channel.test.ts delete mode 100644 js/src/auto-instrumentations/patch-tracing-channel.ts create mode 100644 js/src/global-instrumentation-hooks.test.ts create mode 100644 js/src/global-instrumentation-hooks.ts create mode 100644 js/tests/auto-instrumentations/fixtures/global-hook-listener.cjs delete mode 100644 js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/mod.js delete mode 100644 js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/polyfill.js delete mode 100644 js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/test.js delete mode 100644 js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/mod.mjs delete mode 100644 js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/polyfill.js delete mode 100644 js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/test.mjs delete mode 100644 js/tests/auto-instrumentations/fixtures/test-api-promise-preservation.mjs diff --git a/e2e/scenarios/nextjs-auto-instrumentation/template/app/api/test/route.ts b/e2e/scenarios/nextjs-auto-instrumentation/template/app/api/test/route.ts index b3eb05183..282e699dd 100644 --- a/e2e/scenarios/nextjs-auto-instrumentation/template/app/api/test/route.ts +++ b/e2e/scenarios/nextjs-auto-instrumentation/template/app/api/test/route.ts @@ -1,12 +1,20 @@ -import "braintrust"; // Triggers configureNode(), which applies patchTracingChannel +import "braintrust"; // Registers Braintrust's global instrumentation hooks. import OpenAI from "openai"; -import { tracingChannel } from "node:diagnostics_channel"; import { createServer } from "node:http"; import type { AddressInfo } from "node:net"; // Must be dynamic — this route starts an in-process HTTP server per request. export const dynamic = "force-dynamic"; +type InstrumentationHook = { + subscribe(handlers: InstrumentationHookHandlers): void; + unsubscribe(handlers: InstrumentationHookHandlers): boolean; +}; + +type InstrumentationHookHandlers = { + start(): void; +}; + export async function GET() { // Spin up a minimal mock OpenAI server so no real API calls are made. const mockServer = createServer((_req, res) => { @@ -33,19 +41,20 @@ export async function GET() { ); const { port } = mockServer.address() as AddressInfo; - const channel = tracingChannel("orchestrion:openai:chat.completions.create"); - let channelFired = false; + const hooks = ( + globalThis as typeof globalThis & { + __braintrust_instrumentation_hooks?: Map; + } + ).__braintrust_instrumentation_hooks; + const hook = hooks?.get("orchestrion:openai:chat.completions.create"); + let hookFired = false; const subscriber = { start: () => { - channelFired = true; + hookFired = true; }, - end: () => {}, - asyncStart: () => {}, - asyncEnd: () => {}, - error: () => {}, }; - channel.subscribe(subscriber); + hook?.subscribe(subscriber); try { const client = new OpenAI({ @@ -58,9 +67,9 @@ export async function GET() { messages: [{ role: "user", content: "hi" }], }); } finally { - channel.unsubscribe(subscriber); + hook?.unsubscribe(subscriber); mockServer.close(); } - return Response.json({ instrumented: channelFired }); + return Response.json({ instrumented: hookFired }); } diff --git a/e2e/scenarios/nextjs-auto-instrumentation/template/scenario.ts b/e2e/scenarios/nextjs-auto-instrumentation/template/scenario.ts index 259a2bf29..c7287267e 100644 --- a/e2e/scenarios/nextjs-auto-instrumentation/template/scenario.ts +++ b/e2e/scenarios/nextjs-auto-instrumentation/template/scenario.ts @@ -98,12 +98,12 @@ async function main() { if (!data.instrumented) { throw new Error( - `OpenAI tracing channel did not fire; Next.js ${bundler} instrumentation is not working`, + `OpenAI global hook did not fire; Next.js ${bundler} instrumentation is not working`, ); } console.log( - `OpenAI tracing channel fired at runtime; Next.js ${bundler} instrumentation is active`, + `OpenAI global hook fired at runtime; Next.js ${bundler} instrumentation is active`, ); const edgeResponse = await httpGet(`http://localhost:${PORT}/api/edge`); diff --git a/e2e/scenarios/openai-agents-instrumentation/scenario.impl.mjs b/e2e/scenarios/openai-agents-instrumentation/scenario.impl.mjs index 37106f25c..1edcb6ebb 100644 --- a/e2e/scenarios/openai-agents-instrumentation/scenario.impl.mjs +++ b/e2e/scenarios/openai-agents-instrumentation/scenario.impl.mjs @@ -37,7 +37,7 @@ export async function runOpenAIAgentsAutoInstrumentationScenario() { // Remove the built-in OpenAI trace exporter so it doesn't POST to // api.openai.com/v1/traces/ingest with the cassette placeholder key. - // Braintrust captures trace events via diagnostics_channel independently. + // Braintrust captures trace events via global instrumentation hooks independently. setTraceProcessors([]); await runTracedScenario({ diff --git a/js/NOTICE b/js/NOTICE index 4a357d5a7..71d1b1af2 100644 --- a/js/NOTICE +++ b/js/NOTICE @@ -18,3 +18,9 @@ licenses/import-in-the-middle/NOTICE. This package also includes a modified fork of require-in-the-middle 8.0.1. require-in-the-middle is licensed under the MIT License. See licenses/require-in-the-middle/LICENSE. + +The global instrumentation hook runtime is adapted from the Node.js +diagnostics_channel TracingChannel implementation. +Copyright Node.js contributors. All rights reserved. +It is licensed under the MIT License. +See licenses/node-diagnostics-channel/LICENSE. diff --git a/js/licenses/node-diagnostics-channel/LICENSE b/js/licenses/node-diagnostics-channel/LICENSE new file mode 100644 index 000000000..d8d7f9437 --- /dev/null +++ b/js/licenses/node-diagnostics-channel/LICENSE @@ -0,0 +1,19 @@ +Copyright Node.js contributors. All rights reserved. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to +deal in the Software without restriction, including without limitation the +rights to use, copy, modify, merge, publish, distribute, sublicense, and/or +sell copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS +IN THE SOFTWARE. diff --git a/js/package.json b/js/package.json index 1b8d78807..577d92ebe 100644 --- a/js/package.json +++ b/js/package.json @@ -224,7 +224,6 @@ "cli-progress": "^3.12.0", "cli-table3": "^0.6.5", "cors": "^2.8.5", - "dc-browser": "^1.0.4", "dotenv": "^16.4.5", "esbuild": "0.28.1", "esquery": "^1.7.0", diff --git a/js/src/auto-instrumentations/README.md b/js/src/auto-instrumentations/README.md index f2a319210..8091ebcfb 100644 --- a/js/src/auto-instrumentations/README.md +++ b/js/src/auto-instrumentations/README.md @@ -1,525 +1,127 @@ -# Auto-Instrumentation Internals +# Braintrust Auto-Instrumentation -This document describes the internal architecture of the Braintrust auto-instrumentation system. +Braintrust auto-instrumentation uses the vendored Orchestrion-JS transformer to +wrap selected AI SDK functions at load time or bundle time. Transformed code +invokes tracing-compatible hooks stored in a shared global registry; it does not +import or publish Node.js `diagnostics_channel` events. -## Architecture Overview +## Instrumentation Configs -The auto-instrumentation system uses an internal fork of [orchestrion-js](https://github.com/nodejs/orchestrion-js) to perform AST transformations at load-time or build-time, injecting `diagnostics_channel` instrumentation into AI SDK methods. +Each config identifies a package file and function: -## How It Works - -### 1. Instrumentation Configs - -Each SDK integration defines **instrumentation configs** that specify which functions to instrument: - -```typescript -{ - channelName: 'chat.completions.create', +```ts +const config = { + channelName: "chat.completions.create", module: { - name: 'openai', - versionRange: '>=4.0.0', - filePath: 'resources/chat/completions.mjs' + name: "openai", + versionRange: ">=4.0.0 <7.0.0", + filePath: "resources/chat/completions.mjs", }, functionQuery: { - className: 'Completions', - methodName: 'create', - kind: 'Async' - } -} -``` - -**Note:** The `channelName` should NOT include the `orchestrion:` prefix. The code-transformer automatically prepends `orchestrion:` + `module.name` + `:` to these names. - -**Important:** The code-transformer uses the **exact module name** from the config. For example: - -- Module `openai` → Channel `orchestrion:openai:chat.completions.create` -- Module `@anthropic-ai/sdk` → Channel `orchestrion:@anthropic-ai/sdk:messages.create` -- Module `@google/genai` → Channel `orchestrion:@google/genai:models.generateContent` - -Plugin subscriptions must match these **exact** channel names for instrumentation to work. - -### 2. AST Transformation - -At build-time or load-time, orchestrion-js transforms the source code by wrapping instrumented functions with `tracingChannel` calls: - -**Before:** - -```typescript -class Completions { - async create(params) { - // ... actual implementation - } -} -``` - -**After:** - -```typescript -import { tracingChannel } from 'dc-polyfill'; - -class Completions { - async create(params) { - const channel = tracingChannel('orchestrion:openai:chat.completions.create'); - return channel.tracePromise( - () => /* original implementation */, - { arguments: [params] } - ); - } -} -``` - -### 3. Event Subscription - -The `BraintrustPlugin` (in the main `braintrust` package at `js/src/instrumentation/braintrust-plugin.ts`) subscribes to these channels and converts the events into Braintrust spans: - -```typescript -channel.subscribe({ - start: (event) => { - // Create span and store it on the event for access in other handlers - const span = startSpan({ name: 'Chat Completion', type: 'llm' }); - span.log({ input: event.arguments[0].messages }); - event.span = span; - }, - asyncStart: (event) => { - // Called when the async callback/promise is reached - // Retrieve span from event and log the result - const span = event.span; - span.log({ output: event.result.choices, metrics: { tokens: ... } }); - span.end(); - }, - error: (event) => { - // Retrieve span from event and log the error - const span = event.span; - span.log({ error: event.error }); - span.end(); - } -}); -``` - -## Load-Time vs Build-Time Instrumentation - -### Load-Time (Node.js Hook) - -The `hook.mjs` file uses Node.js's module loader hooks to intercept module loading and apply transformations on-the-fly. This approach: - -- Works with both ESM and CJS modules -- Requires no build step -- Adds minimal overhead at startup -- Is the recommended approach for Node.js applications - -### Build-Time (Bundler Plugins) - -The bundler plugins (Vite, Webpack, esbuild, Rollup) apply transformations during the build process. This approach: - -- Required for browser environments (no Node.js hooks available) -- Zero runtime overhead (transformations done at build time) -- Works with bundled Node.js applications - -## Plugin Architecture - -The plugin system follows the OpenTelemetry pattern: - -- **Core Infrastructure**: `js/src/instrumentation/core/` in the main `braintrust` package - - `BasePlugin` - Abstract base class for plugins - - Channel utilities (`createChannelName`, `parseChannelName`, etc.) - - Event type definitions - -- **Braintrust Implementation**: `js/src/instrumentation/braintrust-plugin.ts` in the main `braintrust` package - - `BraintrustPlugin` - Converts diagnostics_channel events into Braintrust spans - - Automatically enabled when the Braintrust SDK is loaded - - Supports configuration to disable specific integrations - -- **Plugin Registry**: `js/src/instrumentation/registry.ts` in the main `braintrust` package - - `PluginRegistry` - Manages plugin lifecycle - - Automatically enables `BraintrustPlugin` on SDK initialization - - Reads configuration from `configureInstrumentation()` API and `BRAINTRUST_DISABLE_INSTRUMENTATION` environment variable - - `configureInstrumentation()` - Public API for disabling specific integrations - -- **Integration Configs**: `src/configs/` in this package - - SDK-specific instrumentation configurations - - Consumed by orchestrion-js for transformation - -### Auto-Enable Mechanism - -The instrumentation is automatically enabled when the Braintrust SDK is loaded: - -1. **Node.js**: When `configureNode()` is called (in `src/node.ts`), it calls `registry.enable()` -2. **Browser**: When `configureBrowser()` is called (in `src/browser-config.ts`), it calls `registry.enable()` -3. The registry creates and enables `BraintrustPlugin` with the appropriate configuration -4. Individual integrations (OpenAI, Anthropic, etc.) can be disabled via: - - Programmatic API: `configureInstrumentation({ integrations: { openai: false } })` - - Environment variable: `BRAINTRUST_DISABLE_INSTRUMENTATION=openai,anthropic` - -### Configuration Priority - -Configuration is merged in the following order (later overrides earlier): - -1. Default configuration (all integrations enabled) -2. Programmatic configuration via `configureInstrumentation()` -3. Environment variables - -**Important**: `configureInstrumentation()` must be called before the SDK is loaded (before importing any AI SDKs) to take effect. After the registry is enabled, configuration changes will log a warning and be ignored. - -## Adding Support for a New SDK - -1. **Create a config file** in `src/configs/` (e.g., `anthropic.ts`): - -```typescript -import type { InstrumentationConfig } from "braintrust/vite"; - -export const anthropicConfigs: InstrumentationConfig[] = [ - { - channelName: "messages.create", - module: { - name: "@anthropic-ai/sdk", - versionRange: ">=0.20.0", - filePath: "resources/messages.mjs", - }, - functionQuery: { - className: "Messages", - methodName: "create", - kind: "Async", - }, + className: "Completions", + methodName: "create", + kind: "Async", }, -]; -``` - -2. **Export the configs** from `src/index.ts`: - -```typescript -export { anthropicConfigs } from "./configs/anthropic"; -``` - -3. **Add channel handlers** in `BraintrustPlugin` (in the main `braintrust` package): - -```typescript -// In js/src/instrumentation/braintrust-plugin.ts -protected onEnable() { - // ... existing OpenAI handlers ... - - // Add Anthropic handler - const anthropicChannel = tracingChannel('orchestrion:anthropic:messages.create'); - anthropicChannel.subscribe({ - asyncStart: (event) => { - // Create span and log input - }, - asyncEnd: (event) => { - // Log output and end span - }, - error: (event) => { - // Log error and end span - } - }); -} -``` - -## Environment Support: Node.js vs Browser - -### Isomorphic Pattern - -The instrumentation system uses an **isomorphic pattern** to work seamlessly across environments: - -- **Node.js**: Uses native `node:diagnostics_channel` for optimal performance -- **Browser**: Uses `dc-browser` polyfill for compatibility -- **Plugin Code**: Uses `iso.newTracingChannel()` which abstracts the environment difference - -This ensures that: - -1. Both plugin subscription code AND transformed SDK code use the same channel implementation -2. Events properly propagate between emitters and subscribers -3. No channel registry mismatches occur - -### Node.js Setup - -**For runtime applications using the loader hook:** - -```bash -node --import @braintrust/auto-instrumentations/hook.mjs app.js -``` - -The hook automatically uses `node:diagnostics_channel`. - -**For bundled applications:** - -When using bundler plugins (Vite, Webpack, etc.) in Node.js: - -```javascript -// vite.config.js -import { braintrustVitePlugin } from "braintrust/vite"; - -export default { - plugins: [braintrustVitePlugin()], }; ``` -Braintrust-prefixed bundler plugins use Node.js's native `node:diagnostics_channel` module by default. - -### Browser Setup - -**For browser/edge builds:** - -```javascript -// vite.config.js -import { braintrustVitePlugin } from "braintrust/vite"; +`channelName` omits the prefix. Orchestrion constructs the stable identifier: -export default { - plugins: [braintrustVitePlugin({ useDiagnosticChannelCompatShim: true })], -}; -``` - -Setting `useDiagnosticChannelCompatShim: true` ensures the code-transformer injects `dc-browser` imports for environments where Node.js's `diagnostics_channel` module is unavailable. -The deprecated plugin exports are still available and preserve the old `browser` option, which defaults to `true` when omitted; the new Braintrust-prefixed option defaults to `false`. - -**Important:** The `useDiagnosticChannelCompatShim` option must match your target environment: - -- Mismatch (e.g., enabling the compatibility shim but running in Node.js) causes channel registry conflicts -- Plugin code uses the iso pattern and adapts automatically -- Only the transformed SDK code is affected by the compatibility shim option - -### Next.js Setup - -Wrap your Next.js config with `wrapNextjsConfigWithBraintrust` to let Braintrust choose the -webpack plugin or Turbopack loader based on the active Next.js build: - -```javascript -// next.config.mjs -import { wrapNextjsConfigWithBraintrust } from "braintrust/next"; - -const nextConfig = {}; - -export default wrapNextjsConfigWithBraintrust(nextConfig); -``` - -For webpack builds, Braintrust uses Next.js' webpack build context to pass -`browser: true` for client and edge bundles and `browser: false` for Node.js -server bundles. For Turbopack builds, Braintrust inserts webpack loader -rules with runtime conditions so client and edge bundles use `browser: true` -and Node.js server bundles use `browser: false`. - -## Advanced: Custom Plugins - -Third parties can create custom plugins by extending `BasePlugin`: - -```typescript -import { BasePlugin } from "braintrust/instrumentation"; -import { tracingChannel } from "dc-polyfill"; - -class MyCustomPlugin extends BasePlugin { - protected onEnable() { - const channel = tracingChannel( - "orchestrion:openai:chat.completions.create", - ); - - channel.subscribe({ - asyncEnd: (event) => { - // Custom handling - e.g., log to console, send to external service, etc. - console.log("OpenAI call completed:", event.result); - }, - }); - } - - protected onDisable() { - // Cleanup subscriptions - } -} -``` - -Note: Plugins must be constructed and enabled within the core Braintrust library. Due to ESM loading semantics, users cannot manually enable plugins from application code. - -## Testing - -The test suite covers: - -- **Config validation** (`src/configs/*.test.ts`) - Verify config structure -- **AST transformation** (`tests/auto-instrumentations/transformation.test.ts`) - Verify orchestrion-js transforms code correctly -- **Runtime execution** (`tests/auto-instrumentations/runtime-execution.test.ts`) - Verify transformed code executes correctly -- **Integration tests** (`tests/auto-instrumentations/integration.test.ts`) - Verify channel name alignment and event flow -- **Error handling** (`tests/auto-instrumentations/error-handling.test.ts`) - Verify errors propagate correctly -- **Plugin tests** (`src/instrumentation/plugins/*.test.ts`) - Unit tests for each plugin - -## Troubleshooting - -### No Traces Appearing - -If auto-instrumentation isn't creating traces, check the following: - -#### 1. Channel Name Alignment - -**Problem:** Plugin subscriptions don't match the channel names emitted by code-transformer. - -**Symptoms:** SDK calls execute normally but no spans appear in Braintrust. - -**Solution:** Verify channel names match exactly: - -```javascript -// Config (in src/configs/anthropic.ts) -module: { - name: "@anthropic-ai/sdk"; -} -channelName: "messages.create"; - -// Results in emitted channel: -("orchestrion:@anthropic-ai/sdk:messages.create"); // Full scoped package name - -// Plugin MUST subscribe to the exact same name: -iso.newTracingChannel("orchestrion:@anthropic-ai/sdk:messages.create"); +```text +orchestrion:: ``` -**Common mistakes:** - -- Using shortened names (e.g., `anthropic` instead of `@anthropic-ai/sdk`) -- Missing the `@` scope prefix -- Wrong package name entirely - -Run integration tests to verify alignment: - -```bash -pnpm test -- tests/auto-instrumentations/integration.test.ts -``` - -#### 2. Browser/Node.js Mismatch - -**Problem:** Bundler diagnostics channel compatibility setting doesn't match runtime environment. - -**Symptoms:** +The corresponding typed channel definition and plugin subscription must use the +same identifier. -- "Cannot find module 'node:diagnostics_channel'" errors in browser -- Events not propagating in Node.js +## Generated Runtime Contract -**Solution:** +For every configured channel, transformed modules lazily look up: -- For Node.js apps: Use a Braintrust-prefixed bundler plugin with default options -- For browser apps: Set `useDiagnosticChannelCompatShim: true` -- For Node.js runtime apps: Use the loader hook instead of bundling - -#### 3. APIPromise Compatibility (Anthropic SDK) - -**Problem:** Anthropic's `APIPromise` has incompatible constructor with `Promise`. - -**Symptoms:** - -- Errors like "APIPromise constructor called with wrong arguments" -- Traces failing specifically with Anthropic SDK - -**Solution:** The auto-instrumentation hook automatically patches `APIPromise` to fix this. Ensure you're using the hook: - -```bash -node --import @braintrust/auto-instrumentations/hook.mjs app.js -``` - -Or when using bundlers, the patch is applied automatically when plugins are enabled. - -#### 4. Import Order Issues - -**Problem:** SDK imported before Braintrust is configured. - -**Solution:** Import and configure Braintrust before importing AI SDKs: - -```javascript -// CORRECT -import { configureNode } from "braintrust"; -configureNode(); - -import Anthropic from "@anthropic-ai/sdk"; // Now instrumented - -// INCORRECT -import Anthropic from "@anthropic-ai/sdk"; // Not instrumented yet - -import { configureNode } from "braintrust"; -configureNode(); // Too late! -``` - -### Debugging Channel Events - -To debug channel event flow, enable debug logging: - -```bash -DEBUG=braintrust:* node --import @braintrust/auto-instrumentations/hook.mjs app.js -``` - -Or add manual logging in your code: - -```javascript -import iso from "braintrust/isomorph"; - -const channel = iso.newTracingChannel( - "orchestrion:@anthropic-ai/sdk:messages.create", +```js +globalThis.__braintrust_instrumentation_hooks?.get( + "orchestrion:openai:chat.completions.create", ); - -channel.subscribe({ - start: (event) => { - console.log("Channel event received:", event); - }, -}); ``` -### Known Issues +The lookup is retried until a hook exists, then cached. This has two important +properties: -#### TypeScript Errors with Bundler Plugins +- Loading an instrumented provider before Braintrust is safe; calls run normally. +- Registering Braintrust later enables tracing without retransformation. -If TypeScript has resolution issues with direct internal imports, use the public Braintrust bundler entrypoints: +When a hook has no subscribers, the original function is called directly. +Otherwise the generated wrapper invokes `tracePromise`, `traceSync`, or +`traceCallback` with a context containing the original `arguments`, `self` when +available, and `moduleVersion`. -```javascript -import { braintrustVitePlugin } from "braintrust/vite"; -``` - -#### ESM vs CJS Mixing +The hook lifecycle mirrors tracing channels: -The loader hook works best with pure ESM projects. For CJS projects: +1. `start` before the target call +2. `end` after its synchronous portion +3. `asyncStart` and `asyncEnd` when an asynchronous result settles +4. `error` for synchronous throws, promise rejections, or callback errors -- Use the bundler plugin approach -- Or ensure all instrumented SDKs are ESM +The same context object is passed through every phase. Subscribers may mutate +arguments or returned streams before user code continues. -## Migration Guide +## Global Registry -### From dc-browser to Isomorphic Pattern +The SDK installs `globalThis.__braintrust_instrumentation_hooks` with a +non-enumerable, non-writable property descriptor. Its value is a mutable +`Map` shared by all Braintrust SDK copies in the realm. -If you have custom plugins that import `dc-browser` directly: +The implementation lives in `src/global-instrumentation-hooks.ts`. It supports: -**Before:** +- all five lifecycle phases +- multiple subscribers and complete unsubscription +- `bindStore` / `unbindStore` for async-context propagation +- sync, promise, and callback tracing operators +- preservation of Promise subclasses, thenables, and non-Promise return values -```javascript -import { tracingChannel } from "dc-browser"; - -const channel = tracingChannel("my-channel"); -``` +Manual wrappers use the same registry through typed channel definitions, so +manual and auto-instrumented paths share lifecycle and span behavior. -**After:** +## Loaders and Bundlers -```javascript -import iso from "braintrust/isomorph"; +The unified Node hook instruments ESM and CJS: -const channel = iso.newTracingChannel("my-channel"); +```bash +node --import braintrust/hook.mjs app.mjs ``` -This ensures your plugin works in both Node.js (using `node:diagnostics_channel`) and browser (using `dc-browser`) environments. +Bundler integrations are available for esbuild, Vite, Rollup, Webpack, Next.js, +and Turbopack. Generated provider code is runtime-independent and contains no +Node built-in or browser-shim import. -### Channel Name Updates +`useDiagnosticChannelCompatShim` remains accepted by the public bundler options +for source compatibility. It no longer changes the hook transport; it only +retains the legacy browser-target hint used to skip Node-specific special-case +patches. -If you created custom instrumentations before the channel name fixes: +## Adding an Instrumentation -**Before:** +1. Add the narrowest supported package/version/file/function config under + `configs/`. +2. Define a typed channel with the same package and operation identifier. +3. Add or update a plugin that subscribes through the shared channel helpers. +4. Keep manual wrappers on that same typed channel. +5. Add transformation/runtime coverage and a provider e2e scenario when the + user-visible trace contract changes. -```javascript -// Plugin incorrectly used shortened names -channel.subscribe("orchestrion:anthropic:messages.create"); +Instrumentation must preserve target behavior, including receivers, argument +mutation, errors, async context, streams, and custom Promise APIs. -// Config used full name -module: { - name: "@anthropic-ai/sdk"; -} -// ❌ Mismatch - events won't be received! -``` +## Testing -**After:** +Relevant suites live in: -```javascript -// Plugin uses EXACT module name from config -channel.subscribe("orchestrion:@anthropic-ai/sdk:messages.create"); +- `src/global-instrumentation-hooks.test.ts` +- `tests/auto-instrumentations/orchestrion-js-upstream.test.ts` +- `tests/auto-instrumentations/transformation.test.ts` +- `tests/auto-instrumentations/runtime-execution.test.ts` +- `tests/auto-instrumentations/loader-hook.test.ts` -// Config uses full name -module: { - name: "@anthropic-ai/sdk"; -} -// ✅ Match - events will be received -``` +The transformation suites assert that output contains the global registry lookup +and contains neither `diagnostics_channel` nor `dc-browser`. Provider e2e tests +must run in cassette replay mode after instrumentation changes. diff --git a/js/src/auto-instrumentations/bundler/next.ts b/js/src/auto-instrumentations/bundler/next.ts index 53112eb0d..774a4ab98 100644 --- a/js/src/auto-instrumentations/bundler/next.ts +++ b/js/src/auto-instrumentations/bundler/next.ts @@ -1,5 +1,5 @@ import { createRequire } from "node:module"; -import { isAbsolute, join, relative } from "node:path"; +import { join } from "node:path"; import { webpackPlugin } from "./webpack"; type MaybePromise = T | Promise; @@ -159,53 +159,14 @@ function wrapTurbopackConfig( turbopackConfig: TurbopackConfig | undefined, ): TurbopackConfig { const config = { ...(turbopackConfig ?? {}) }; - const resolveAlias = - config.resolveAlias && - typeof config.resolveAlias === "object" && - !Array.isArray(config.resolveAlias) - ? config.resolveAlias - : {}; const rules = config.rules && typeof config.rules === "object" && !Array.isArray(config.rules) ? config.rules : {}; - let dcBrowserPath: string | undefined; - - try { - dcBrowserPath = createRequire( - requireFromProject.resolve("braintrust/package.json"), - ).resolve("dc-browser"); - } catch { - try { - dcBrowserPath = requireFromProject.resolve("dc-browser"); - } catch { - dcBrowserPath = undefined; - } - } - - // Absolute Turbopack aliases are interpreted as server-relative imports, so - // make the resolved dependency path relative to the app config directory. - if (dcBrowserPath && isAbsolute(dcBrowserPath)) { - const relativeDcBrowserPath = relative( - process.cwd(), - dcBrowserPath, - ).replace(/\\/g, "/"); - dcBrowserPath = relativeDcBrowserPath.startsWith(".") - ? relativeDcBrowserPath - : `./${relativeDcBrowserPath}`; - } - return { ...config, - // Turbopack resolves modules emitted by our loader from the app graph. The - // browser diagnostics-channel shim is Braintrust's dependency, so alias it - // to the copy installed with Braintrust while still letting user aliases win. - resolveAlias: - dcBrowserPath && resolveAlias["dc-browser"] === undefined - ? { ...resolveAlias, "dc-browser": dcBrowserPath } - : config.resolveAlias, rules: addBraintrustTurbopackRule(rules), }; } diff --git a/js/src/auto-instrumentations/bundler/plugin.ts b/js/src/auto-instrumentations/bundler/plugin.ts index dda1b7490..2f7295f75 100644 --- a/js/src/auto-instrumentations/bundler/plugin.ts +++ b/js/src/auto-instrumentations/bundler/plugin.ts @@ -21,9 +21,8 @@ export interface LegacyBundlerPluginOptions { /** * Whether to bundle for browser environments. * - * When true, uses 'dc-browser' for browser-compatible diagnostics_channel polyfill. - * When false, uses Node.js built-in 'diagnostics_channel' and 'async_hooks'. - * Defaults to true (assumes browser build). + * This remains as a legacy target hint for special-case patches. Global + * instrumentation hooks are runtime-independent. */ browser?: boolean; } @@ -40,13 +39,11 @@ export interface BundlerPluginOptions { instrumentations?: InstrumentationConfig[]; /** - * Use the `diagnostics_channel` compatibility shim in patched code instead - * of Node.js's built-in `diagnostics_channel` module. + * Legacy browser target hint. * - * Enable this for browser, edge, or worker bundles where Node's - * `diagnostics_channel` module is unavailable. Leave it disabled for Node.js - * bundles so transformed SDK code publishes on the native `diagnostics_channel` - * registry. + * Global instrumentation hooks no longer require a diagnostics-channel shim. + * The option is retained for source compatibility and only controls whether + * Node-specific special-case patches are skipped. * * @default false */ @@ -76,11 +73,8 @@ export const unplugin = createUnplugin( additionalInstrumentations: options.instrumentations, }); - // Default to browser build, use polyfill unless explicitly disabled - const dcModule = options.browser === false ? undefined : "dc-browser"; - // Create the code transformer instrumentor - const instrumentationMatcher = create(allInstrumentations, dcModule); + const instrumentationMatcher = create(allInstrumentations); return { name: "code-transformer", @@ -164,13 +158,8 @@ export const unplugin = createUnplugin( // Transform the code const moduleType = isModule ? "esm" : "cjs"; const result = transformer.transform(code, moduleType); - const transformedCode = result.code.replace( - /const \{tracingChannel: ([A-Za-z_$][\w$]*)\} = ([A-Za-z_$][\w$]*);/g, - "const $1 = $2.tracingChannel;", - ); - return { - code: transformedCode, + code: result.code, map: result.map, }; } catch (error) { diff --git a/js/src/auto-instrumentations/bundler/webpack-loader.ts b/js/src/auto-instrumentations/bundler/webpack-loader.ts index c218c21a6..63d06a1c1 100644 --- a/js/src/auto-instrumentations/bundler/webpack-loader.ts +++ b/js/src/auto-instrumentations/bundler/webpack-loader.ts @@ -58,8 +58,7 @@ function getMatcher(options: LegacyBundlerPluginOptions): Matcher { const allInstrumentations = getDefaultInstrumentationConfigs({ additionalInstrumentations: options.instrumentations, }); - const dcModule = options.browser ? "dc-browser" : undefined; - const configHash = JSON.stringify({ allInstrumentations, dcModule }); + const configHash = JSON.stringify({ allInstrumentations }); if (matcherCache.has(configHash)) { return matcherCache.get(configHash)!; @@ -71,7 +70,7 @@ function getMatcher(options: LegacyBundlerPluginOptions): Matcher { } } - const matcher = create(allInstrumentations, dcModule ?? null); + const matcher = create(allInstrumentations); matcherCache.set(configHash, matcher); return matcher; } diff --git a/js/src/auto-instrumentations/hook.mts b/js/src/auto-instrumentations/hook.mts index 10626cebd..fd6172dc4 100644 --- a/js/src/auto-instrumentations/hook.mts +++ b/js/src/auto-instrumentations/hook.mts @@ -5,7 +5,7 @@ * node --import @braintrust/auto-instrumentations/hook.mjs app.js * * This hook performs AST transformation at load-time for BOTH ESM and CJS modules, - * injecting TracingChannel calls into AI SDK functions. + * injecting global instrumentation hook calls into AI SDK functions. * * Many modern apps use a mix of ESM and CJS modules, so this single hook * handles both: @@ -22,22 +22,12 @@ import { BraintrustObservabilityExporter } from "../wrappers/mastra.js"; import { installMastraExporterFactory } from "./loader/mastra-observability-patch.js"; import { getDefaultAutoInstrumentationConfigs } from "./configs/all.js"; import { ModulePatch } from "./loader/cjs-patch.js"; -import { patchTracingChannel } from "./patch-tracing-channel.js"; const state = ((globalThis as any)[ Symbol.for("braintrust.applyAutoInstrumentation") ] ??= {}) as { applied?: boolean }; const alreadyApplied = state.applied; -// Patch diagnostics_channel.tracePromise to handle APIPromise correctly. -// MUST be done here (before any SDK code runs) to fix Anthropic APIPromise incompatibility. -// Construct the module path dynamically to prevent build from stripping "node:" prefix. -if (!alreadyApplied) { - const dcPath = ["node", "diagnostics_channel"].join(":"); - const dc: any = await import(/* @vite-ignore */ dcPath as any); - patchTracingChannel(dc.tracingChannel); -} - if (!alreadyApplied) { const allConfigs = getDefaultAutoInstrumentationConfigs(); diff --git a/js/src/auto-instrumentations/index.ts b/js/src/auto-instrumentations/index.ts index ea0533ac6..1c16b2d46 100644 --- a/js/src/auto-instrumentations/index.ts +++ b/js/src/auto-instrumentations/index.ts @@ -1,7 +1,7 @@ /** * @braintrust/auto-instrumentations * - * Auto-instrumentation for AI SDKs using orchestrion-js and diagnostics_channel. + * Auto-instrumentation for AI SDKs using orchestrion-js and global hooks. * * This package provides: * - Instrumentation configs for orchestrion-js diff --git a/js/src/auto-instrumentations/loader/cjs-patch.ts b/js/src/auto-instrumentations/loader/cjs-patch.ts index 670a6a5b1..df5c448ae 100644 --- a/js/src/auto-instrumentations/loader/cjs-patch.ts +++ b/js/src/auto-instrumentations/loader/cjs-patch.ts @@ -30,7 +30,7 @@ export class ModulePatch { /** * Patches the Node.js module class method that is responsible for compiling code. * If a module is found that has an instrumentator, it will transform the code before compiling it - * with tracing channel methods. + * with global hook calls. */ patch() { const self = this; diff --git a/js/src/auto-instrumentations/loader/special-case-patches.ts b/js/src/auto-instrumentations/loader/special-case-patches.ts index 4b984c158..d73c53010 100644 --- a/js/src/auto-instrumentations/loader/special-case-patches.ts +++ b/js/src/auto-instrumentations/loader/special-case-patches.ts @@ -4,7 +4,7 @@ * ⚠️ ANTI-PATTERN — DO NOT EXTEND CASUALLY. * * Every entry in this file represents a target SDK that doesn't expose a - * stable extension point we can hook through diagnostics_channel + the + * stable extension point we can hook through global instrumentation hooks + the * internal Orchestrion matcher. New integrations should * **prefer the standard channel-handler / `BasePlugin` pattern** used by * every other integration in `js/src/instrumentation/plugins/*-plugin.ts`. diff --git a/js/src/auto-instrumentations/orchestrion-js/index.ts b/js/src/auto-instrumentations/orchestrion-js/index.ts index 0192c0d9c..245ea7cf8 100644 --- a/js/src/auto-instrumentations/orchestrion-js/index.ts +++ b/js/src/auto-instrumentations/orchestrion-js/index.ts @@ -11,9 +11,8 @@ import type { InstrumentationConfig } from "./types"; */ export function create( configs: InstrumentationConfig[], - dcModule?: string | null, ): InstrumentationMatcher { - return new InstrumentationMatcher(configs, dcModule); + return new InstrumentationMatcher(configs); } export type { InstrumentationConfig, ModuleType } from "./types"; diff --git a/js/src/auto-instrumentations/orchestrion-js/matcher.ts b/js/src/auto-instrumentations/orchestrion-js/matcher.ts index 66a61f94c..073317085 100644 --- a/js/src/auto-instrumentations/orchestrion-js/matcher.ts +++ b/js/src/auto-instrumentations/orchestrion-js/matcher.ts @@ -13,12 +13,10 @@ import type { InstrumentationConfig } from "./types"; */ export class InstrumentationMatcher { private configs: InstrumentationConfig[] = []; - private dcModule: string; private transformers: Record = {}; - constructor(configs: InstrumentationConfig[], dcModule?: string | null) { + constructor(configs: InstrumentationConfig[]) { this.configs = configs; - this.dcModule = dcModule || "diagnostics_channel"; } /** @@ -54,7 +52,6 @@ export class InstrumentationMatcher { version, filePath, configs, - this.dcModule, ); return this.transformers[id]; diff --git a/js/src/auto-instrumentations/orchestrion-js/transformer.ts b/js/src/auto-instrumentations/orchestrion-js/transformer.ts index ecd0d66f3..8d1385c21 100644 --- a/js/src/auto-instrumentations/orchestrion-js/transformer.ts +++ b/js/src/auto-instrumentations/orchestrion-js/transformer.ts @@ -22,7 +22,7 @@ type TraceOperator = "traceCallback" | "tracePromise" | "traceSync"; /** * Applies instrumentation configs to JavaScript source by parsing it into an - * AST, locating target functions, injecting diagnostics_channel tracing + * AST, locating target functions, injecting global instrumentation hooks * wrappers, and regenerating the source. */ export class Transformer { @@ -30,24 +30,21 @@ export class Transformer { private version: string; private filePath: string; private configs: InstrumentationConfig[] = []; - private dcModule: string; constructor( moduleName: string, version: string, filePath: string, configs: InstrumentationConfig[], - dcModule: string, ) { this.moduleName = moduleName; this.version = version; this.filePath = filePath; this.configs = configs; - this.dcModule = dcModule; } /** - * Instruments `code` by injecting diagnostics_channel tracing around the + * Instruments `code` by injecting global tracing hooks around the * target functions defined by this transformer's configs. */ transform(code: string | Buffer, moduleType: ModuleType): TransformOutput { @@ -91,7 +88,6 @@ export class Transformer { const query = astQuery || this.fromFunctionQuery(resolvedFunctionQuery); const state: TransformState = { ...config, - dcModule: this.dcModule, moduleType, moduleVersion: this.version, functionQuery: resolvedFunctionQuery, diff --git a/js/src/auto-instrumentations/orchestrion-js/transforms.ts b/js/src/auto-instrumentations/orchestrion-js/transforms.ts index 92a1adfaa..32787bd81 100644 --- a/js/src/auto-instrumentations/orchestrion-js/transforms.ts +++ b/js/src/auto-instrumentations/orchestrion-js/transforms.ts @@ -5,6 +5,7 @@ import esquery from "esquery"; import { parse } from "meriyah"; +import { GLOBAL_INSTRUMENTATION_HOOKS_KEY } from "../../global-instrumentation-hooks"; import type { FunctionQuery, InstrumentationConfig, ModuleType } from "./types"; type AnyNode = any; @@ -17,7 +18,6 @@ type TransformFn = ( type TraceOperator = "traceCallback" | "tracePromise" | "traceSync"; export interface TransformState extends InstrumentationConfig { - dcModule: string; moduleType: ModuleType; moduleVersion: string; functionQuery: FunctionQuery; @@ -25,71 +25,44 @@ export interface TransformState extends InstrumentationConfig { functionIndex?: number; } -const tracingChannelPredicate = (node: AnyNode): boolean => - node.declarations?.[0]?.id?.properties?.[0]?.value?.name === - "tr_ch_apm_tracingChannel"; - const CHANNEL_REGEX = /[^\w]/g; function formatChannelVariable(channelName: string): string { return `tr_ch_apm$${channelName.replace(CHANNEL_REGEX, "_")}`; } -export const transforms: Record = { - tracingChannelImport({ dcModule, moduleType }, node) { - if (node.body.some(tracingChannelPredicate)) { - return; - } - - const options = { module: moduleType === "esm" }; - const index = node.body.findIndex( - (child: AnyNode) => child.directive === "use strict", - ); - const dc = - moduleType === "esm" - ? `import tr_ch_apm_dc from "${dcModule}"` - : `const tr_ch_apm_dc = ${"require"}("${dcModule}")`; - const tracingChannel = - "const { tracingChannel: tr_ch_apm_tracingChannel } = tr_ch_apm_dc"; - const hasSubscribers = `const tr_ch_apm_hasSubscribers = ch => ch.start.hasSubscribers - || ch.end.hasSubscribers - || ch.asyncStart.hasSubscribers - || ch.asyncEnd.hasSubscribers - || ch.error.hasSubscribers`; - - node.body.splice( - index + 1, - 0, - parse(dc, options as any).body[0], - parse(tracingChannel, options as any).body[0], - parse(hasSubscribers, options as any).body[0], - ); - }, +function formatChannelGetter(channelName: string): string { + return `tr_ch_apm$get_${channelName.replace(CHANNEL_REGEX, "_")}`; +} - tracingChannelDeclaration(state, node) { +export const transforms: Record = { + tracingHookDeclaration(state, node) { const { channelName, module: { name }, } = state; const channelVariable = formatChannelVariable(channelName); + const channelGetter = formatChannelGetter(channelName); if ( node.body.some( - (child: AnyNode) => - child.declarations?.[0]?.id?.name === channelVariable, + (child: AnyNode) => child.declarations?.[0]?.id?.name === channelGetter, ) ) { return; } - transforms.tracingChannelImport(state, node, null, []); - - const index = node.body.findIndex(tracingChannelPredicate); + const index = node.body.findIndex( + (child: AnyNode) => child.directive === "use strict", + ); const code = ` - const ${channelVariable} = tr_ch_apm_tracingChannel("orchestrion:${name}:${channelName}") + let ${channelVariable}; + const ${channelGetter} = () => ${channelVariable} ??= globalThis[${JSON.stringify( + GLOBAL_INSTRUMENTATION_HOOKS_KEY, + )}]?.get?.("orchestrion:${name}:${channelName}"); `; - node.body.splice(index + 1, 0, parse(code).body[0]); + node.body.splice(index + 1, 0, ...parse(code).body); }, traceCallback: traceAny, @@ -117,7 +90,7 @@ function traceFunction( node: AnyNode, program: AnyNode, ): void { - transforms.tracingChannelDeclaration(state, program, null, []); + transforms.tracingHookDeclaration(state, program, null, []); const { functionQuery } = state; const methodName = @@ -172,7 +145,7 @@ function traceInstanceMethod( let ctor = classBody.body.find(({ kind }: AnyNode) => kind === "constructor"); - transforms.tracingChannelDeclaration(state, program, null, []); + transforms.tracingHookDeclaration(state, program, null, []); if (!ctor) { ctor = ( @@ -307,141 +280,46 @@ function wrapCallback(state: TransformState): AnyNode { channelName, functionQuery: { callbackIndex = -1 }, } = state; - const channelVariable = formatChannelVariable(channelName); + const channelGetter = formatChannelGetter(channelName); return parse(` function wrapper () { - const __apm$cb = Array.prototype.at.call(arguments, ${callbackIndex}); - - if (!${channelVariable}.start.hasSubscribers) return __apm$traced(); - - function __apm$wrappedCb(err, res) { - if (err) { - __apm$ctx.error = err; - ${channelVariable}.error.publish(__apm$ctx); - } else { - __apm$ctx.result = res; - } - - ${channelVariable}.asyncStart.runStores(__apm$ctx, () => { - try { - if (__apm$cb) { - return __apm$cb.apply(this, arguments); - } - } finally { - ${channelVariable}.asyncEnd.publish(__apm$ctx); - } - }); - } - - if (typeof __apm$cb !== 'function') { - return __apm$traced(); - } - Array.prototype.splice.call(arguments, ${callbackIndex}, 1, __apm$wrappedCb); - - return ${channelVariable}.start.runStores(__apm$ctx, () => { - try { - return __apm$traced(); - } catch (err) { - __apm$ctx.error = err; - ${channelVariable}.error.publish(__apm$ctx); - throw err; - } finally { - __apm$ctx.self ??= this; - ${channelVariable}.end.publish(__apm$ctx); - } - }); + const __apm$hook = ${channelGetter}(); + if (!__apm$hook?.hasSubscribers) return __apm$traced(); + __apm$ctx.self ??= this; + return __apm$hook.traceCallback( + __apm$traced, + ${callbackIndex}, + __apm$ctx + ); } `); } function wrapPromise(state: TransformState): AnyNode { const { channelName } = state; - const channelVariable = formatChannelVariable(channelName); + const channelGetter = formatChannelGetter(channelName); return parse(` function wrapper () { - if (!tr_ch_apm_hasSubscribers(${channelVariable})) return __apm$traced(); - - return ${channelVariable}.start.runStores(__apm$ctx, () => { - try { - let promise = __apm$traced(); - if (typeof promise?.then !== 'function') { - __apm$ctx.result = promise; - return promise; - } - // Mirror Node.js core diagnostics_channel behaviour: for native Promise - // instances, chain normally (safe since there is no subclass API to - // preserve). For Promise subclasses and other thenables, side-chain the - // callbacks for event publishing and return the original so that any - // subclass-specific methods (e.g. APIPromise.withResponse()) remain - // accessible to the caller. - if (promise instanceof Promise && promise.constructor === Promise) { - return promise.then( - result => { - __apm$ctx.result = result; - ${channelVariable}.asyncStart.publish(__apm$ctx); - ${channelVariable}.asyncEnd.publish(__apm$ctx); - return result; - }, - err => { - __apm$ctx.error = err; - ${channelVariable}.error.publish(__apm$ctx); - ${channelVariable}.asyncStart.publish(__apm$ctx); - ${channelVariable}.asyncEnd.publish(__apm$ctx); - throw err; - } - ); - } - promise.then( - result => { - __apm$ctx.result = result; - ${channelVariable}.asyncStart.publish(__apm$ctx); - ${channelVariable}.asyncEnd.publish(__apm$ctx); - }, - err => { - __apm$ctx.error = err; - ${channelVariable}.error.publish(__apm$ctx); - ${channelVariable}.asyncStart.publish(__apm$ctx); - ${channelVariable}.asyncEnd.publish(__apm$ctx); - } - ); - return promise; - } catch (err) { - __apm$ctx.error = err; - ${channelVariable}.error.publish(__apm$ctx); - throw err; - } finally { - __apm$ctx.self ??= this; - ${channelVariable}.end.publish(__apm$ctx); - } - }); + const __apm$hook = ${channelGetter}(); + if (!__apm$hook?.hasSubscribers) return __apm$traced(); + __apm$ctx.self ??= this; + return __apm$hook.tracePromise(__apm$traced, __apm$ctx); } `); } function wrapSync(state: TransformState): AnyNode { const { channelName } = state; - const channelVariable = formatChannelVariable(channelName); + const channelGetter = formatChannelGetter(channelName); return parse(` function wrapper () { - if (!tr_ch_apm_hasSubscribers(${channelVariable})) return __apm$traced(); - - return ${channelVariable}.start.runStores(__apm$ctx, () => { - try { - const result = __apm$traced(); - __apm$ctx.result = result; - return result; - } catch (err) { - __apm$ctx.error = err; - ${channelVariable}.error.publish(__apm$ctx); - throw err; - } finally { - __apm$ctx.self ??= this; - ${channelVariable}.end.publish(__apm$ctx); - } - }); + const __apm$hook = ${channelGetter}(); + if (!__apm$hook?.hasSubscribers) return __apm$traced(); + __apm$ctx.self ??= this; + return __apm$hook.traceSync(__apm$traced, __apm$ctx); } `); } diff --git a/js/src/auto-instrumentations/orchestrion-js/types.ts b/js/src/auto-instrumentations/orchestrion-js/types.ts index 67237c18e..603519cc7 100644 --- a/js/src/auto-instrumentations/orchestrion-js/types.ts +++ b/js/src/auto-instrumentations/orchestrion-js/types.ts @@ -59,7 +59,7 @@ export type FunctionQuery = */ export interface InstrumentationConfig { /** - * The name of the diagnostics channel to publish to. + * The name of the global instrumentation hook to invoke. */ channelName: string; diff --git a/js/src/auto-instrumentations/patch-tracing-channel.test.ts b/js/src/auto-instrumentations/patch-tracing-channel.test.ts deleted file mode 100644 index e346717a6..000000000 --- a/js/src/auto-instrumentations/patch-tracing-channel.test.ts +++ /dev/null @@ -1,255 +0,0 @@ -/** - * Regression tests for patchTracingChannel. - * - * The bug: Node.js diagnostics_channel's tracePromise calls PromisePrototypeThen - * (the original, unoverridden Promise.prototype.then) on the return value. This uses - * the Symbol.species protocol to create the result promise — it calls - * new SubClass(resolve, reject) — which breaks for Promise subclasses that don't - * accept a standard (resolve, reject) executor as their first argument (e.g., - * Anthropic's APIPromise, which takes a responsePromise). - * - * The fix: call result.then(resolve, reject) directly, which invokes whatever .then() - * override the subclass provides (typically a safe version that delegates to an inner - * native Promise, avoiding the species protocol entirely). - */ - -import { describe, it, expect } from "vitest"; -import { patchTracingChannel } from "./patch-tracing-channel"; - -/** - * Creates a fresh TracingChannel-like class with the ORIGINAL broken tracePromise - * behavior, simulating Node.js diagnostics_channel before the patch is applied. - * - * The key difference from the patched version: it calls - * `Promise.prototype.then.call(result, ...)` (equivalent to Node.js's internal - * PromisePrototypeThen), which triggers Symbol.species on Promise subclasses. - */ -function makeUnpatchedTracingChannel() { - class FakeChannel { - readonly hasSubscribers = true; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - publish(_msg: any): void {} - // eslint-disable-next-line @typescript-eslint/no-explicit-any - runStores(_msg: any, fn: () => any): any { - return fn(); - } - } - - class FakeTracingChannel { - start = new FakeChannel(); - end = new FakeChannel(); - asyncStart = new FakeChannel(); - asyncEnd = new FakeChannel(); - error = new FakeChannel(); - - tracePromise( - // eslint-disable-next-line @typescript-eslint/no-explicit-any - fn: (...args: any[]) => any, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - context: Record = {}, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - thisArg: any, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - ...args: any[] - ): Promise { - const { start, end, asyncStart, asyncEnd, error } = this; - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - function reject(err: any) { - context.error = err; - error.publish(context); - asyncStart.publish(context); - asyncEnd.publish(context); - return Promise.reject(err); - } - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - function resolve(result: any) { - context.result = result; - asyncStart.publish(context); - asyncEnd.publish(context); - return result; - } - - start.publish(context); - - try { - const result = Reflect.apply(fn, thisArg, args); - - // BROKEN: calls original Promise.prototype.then directly, equivalent to - // Node.js's internal PromisePrototypeThen(result, resolve, reject). - // This triggers Symbol.species on Promise subclasses, calling - // new SubClass(resolve, reject) — which breaks for non-standard constructors. - // eslint-disable-next-line @typescript-eslint/no-explicit-any - return (Promise.prototype.then as any).call( - result as Promise, - resolve, - reject, - ); - } catch (err) { - context.error = err; - error.publish(context); - throw err; - } finally { - end.publish(context); - } - } - } - - return FakeTracingChannel; -} - -/** - * Minimal reproduction of Anthropic's APIPromise pattern: - * - Non-standard constructor: takes a responsePromise, not a (resolve, reject) executor - * - Overrides .then() to delegate to an inner native Promise (avoiding species protocol) - * - Sets Symbol.species to a class that throws, reproducing the exact error - * - * When PromisePrototypeThen(apiPromise, resolve, reject) is called (the original .then), - * it uses Symbol.species to create the result promise: new FailingSpecies(resolve, reject) - * which throws "Promise resolve or reject function is not callable". - * - * When apiPromise.then(resolve, reject) is called (the patched behavior), it invokes - * the safe .then() override which delegates to the inner native promise, bypassing - * Symbol.species entirely. - */ -class MockAPIPromise extends Promise { - readonly #innerPromise: Promise; - - constructor(responsePromise: Promise) { - let res!: (value: T | PromiseLike) => void; - super((resolve) => { - res = resolve; - }); - this.#innerPromise = Promise.resolve(responsePromise); - this.#innerPromise.then(res); - } - - // Safe override: delegates to the inner native Promise instead of `this`. - // Calling .then() on a native Promise uses native Promise as the species, - // so it never tries to call new MockAPIPromise(resolve, reject). - override then( - onfulfilled?: ((value: T) => TResult1 | PromiseLike) | null, - onrejected?: ((reason: unknown) => TResult2 | PromiseLike) | null, - ): Promise { - return this.#innerPromise.then(onfulfilled, onrejected); - } - - async withResponse() { - return { - data: await this.#innerPromise, - response: { ok: true }, - }; - } - - // Symbol.species returns a class that throws, reproducing the exact error seen - // when Node.js's PromisePrototypeThen tries to construct the result promise. - static override get [Symbol.species](): PromiseConstructor { - return class { - constructor() { - throw new TypeError( - "Promise resolve or reject function is not callable", - ); - } - } as unknown as PromiseConstructor; - } -} - -describe("patchTracingChannel", () => { - it("unpatched tracePromise throws when given a Promise subclass with incompatible Symbol.species", () => { - // This reproduces the error reported with Anthropic's APIPromise + Node.js bundler path. - const FakeTCClass = makeUnpatchedTracingChannel(); - const channel = new FakeTCClass(); - const apiPromise = new MockAPIPromise(Promise.resolve("hello")); - - // The unpatched tracePromise calls Promise.prototype.then.call(result, ...) which - // uses Symbol.species to create the result promise, throwing synchronously. - expect(() => channel.tracePromise(() => apiPromise, {}, null)).toThrow( - "Promise resolve or reject function is not callable", - ); - }); - - it("patched tracePromise correctly handles Promise subclasses with non-standard Symbol.species", async () => { - const FakeTCClass = makeUnpatchedTracingChannel(); - const channel = new FakeTCClass(); - - // Apply the patch — replaces FakeTCClass.prototype.tracePromise with the fixed version - patchTracingChannel(() => channel); - - const context: Record = {}; - const apiPromise = new MockAPIPromise(Promise.resolve("hello")); - - // The patched tracePromise calls result.then(resolve, reject) directly, which - // invokes our safe .then() override that avoids the species protocol. - const result = await channel.tracePromise(() => apiPromise, context, null); - - expect(result).toBe("hello"); - expect(context.result).toBe("hello"); - }); - - it("patched tracePromise preserves helper methods on promise subclasses", async () => { - const FakeTCClass = makeUnpatchedTracingChannel(); - const channel = new FakeTCClass(); - patchTracingChannel(() => channel); - - const apiPromise = new MockAPIPromise(Promise.resolve("hello")); - const traced = channel.tracePromise(() => apiPromise, {}, null); - const withResponse = await (traced as any).withResponse(); - - expect(traced).toBe(apiPromise); - expect(withResponse.data).toBe("hello"); - expect(withResponse.response.ok).toBe(true); - }); - - it("patched tracePromise correctly handles plain async functions", async () => { - const FakeTCClass = makeUnpatchedTracingChannel(); - const channel = new FakeTCClass(); - patchTracingChannel(() => channel); - - const context: Record = {}; - const result = await channel.tracePromise( - async () => "world", - context, - null, - ); - - expect(result).toBe("world"); - expect(context.result).toBe("world"); - }); - - it("patched tracePromise propagates rejections and sets context.error", async () => { - const FakeTCClass = makeUnpatchedTracingChannel(); - const channel = new FakeTCClass(); - patchTracingChannel(() => channel); - - const context: Record = {}; - const err = new Error("api error"); - - await expect( - channel.tracePromise( - async () => { - throw err; - }, - context, - null, - ), - ).rejects.toBe(err); - - expect(context.error).toBe(err); - }); - - it("is idempotent — applying the patch twice produces correct behavior", async () => { - const FakeTCClass = makeUnpatchedTracingChannel(); - const channel = new FakeTCClass(); - - patchTracingChannel(() => channel); - patchTracingChannel(() => channel); - - const context: Record = {}; - const apiPromise = new MockAPIPromise(Promise.resolve(99)); - const result = await channel.tracePromise(() => apiPromise, context, null); - - expect(result).toBe(99); - expect(context.result).toBe(99); - }); -}); diff --git a/js/src/auto-instrumentations/patch-tracing-channel.ts b/js/src/auto-instrumentations/patch-tracing-channel.ts deleted file mode 100644 index 349e05e14..000000000 --- a/js/src/auto-instrumentations/patch-tracing-channel.ts +++ /dev/null @@ -1,146 +0,0 @@ -/** - * Patches TracingChannel.prototype to handle APIPromise and other Promise subclasses - * that change the constructor signature (violating the species contract). - * - * node:diagnostics_channel's tracePromise wraps the result with Promise.resolve(), - * which calls the subclass constructor with the wrong signature for classes like - * Anthropic's APIPromise. This patch uses duck-typing (.then check) instead. - * - * This is applied both in the loader hook (hook.mts) for the --import path, - * and in configureNode/configureBrowser for the bundler plugin path. - */ - -// eslint-disable-next-line @typescript-eslint/no-explicit-any -export function patchTracingChannel( - tracingChannelFn: (name: string) => any, -): void { - const dummyChannel = tracingChannelFn("__braintrust_probe__"); - const TracingChannel = dummyChannel?.constructor; - - if (!TracingChannel?.prototype) { - return; - } - - if ( - !Object.getOwnPropertyDescriptor(TracingChannel.prototype, "hasSubscribers") - ) { - Object.defineProperty(TracingChannel.prototype, "hasSubscribers", { - configurable: true, - enumerable: false, - get(this: { - start?: { hasSubscribers?: boolean }; - end?: { hasSubscribers?: boolean }; - asyncStart?: { hasSubscribers?: boolean }; - asyncEnd?: { hasSubscribers?: boolean }; - error?: { hasSubscribers?: boolean }; - }) { - return Boolean( - this.start?.hasSubscribers || - this.end?.hasSubscribers || - this.asyncStart?.hasSubscribers || - this.asyncEnd?.hasSubscribers || - this.error?.hasSubscribers, - ); - }, - }); - } - - if (TracingChannel.prototype.tracePromise) { - TracingChannel.prototype.tracePromise = function ( - // eslint-disable-next-line @typescript-eslint/no-explicit-any - fn: any, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - context: any = {}, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - thisArg: any, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - ...args: any[] - ) { - const start = this.start; - const end = this.end; - const asyncStart = this.asyncStart; - const asyncEnd = this.asyncEnd; - const error = this.error; - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - function publishRejected(err: any) { - context.error = err; - error?.publish(context); - asyncStart?.publish(context); - asyncEnd?.publish(context); - } - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - function publishResolved(result: any) { - context.result = result; - asyncStart?.publish(context); - asyncEnd?.publish(context); - } - - // Use runStores (not just publish) so fn runs inside the ALS context - // established by bindStore — required for span context to propagate across awaits. - // PATCHED: inside the callback, use duck-type thenable check instead of - // PromisePrototypeThen, which triggers Symbol.species and breaks Promise subclasses - // like Anthropic's APIPromise that have non-standard constructors. - return start.runStores(context, () => { - try { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const result: any = Reflect.apply(fn, thisArg, args); - end?.publish(context); - - if ( - result && - (typeof result === "object" || typeof result === "function") && - typeof result.then === "function" - ) { - if (result.constructor === Promise) { - return result.then( - (res) => { - publishResolved(res); - return res; - }, - (err) => { - publishRejected(err); - return Promise.reject(err); - }, - ); - } - - // Preserve the original promise-like object so SDK helper methods - // like Anthropic APIPromise.withResponse() remain available. - void result.then( - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (resolved: any) => { - try { - publishResolved(resolved); - } catch { - // Preserve wrapped promise semantics even if instrumentation fails. - } - }, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - (err: any) => { - try { - publishRejected(err); - } catch { - // Preserve wrapped promise semantics even if instrumentation fails. - } - }, - ); - - return result; - } - - context.result = result; - asyncStart?.publish(context); - asyncEnd?.publish(context); - return result; - } catch (err) { - context.error = err; - error?.publish(context); - end?.publish(context); - throw err; - } - }); - }; - } -} diff --git a/js/src/browser/config.ts b/js/src/browser/config.ts index 5aa398669..c2d649bdb 100644 --- a/js/src/browser/config.ts +++ b/js/src/browser/config.ts @@ -1,8 +1,6 @@ -import { tracingChannel } from "dc-browser"; import iso from "../isomorph"; import { _internalSetInitialState } from "../logger"; import { registry } from "../instrumentation/registry"; -import { patchTracingChannel } from "../auto-instrumentations/patch-tracing-channel"; // This is copied from next.js. It seems they define AsyncLocalStorage in the edge // environment, even though it's not defined in the browser. @@ -34,13 +32,6 @@ export function configureBrowser(): void { // Ignore } - iso.newTracingChannel = <_M = any>(nameOrChannels: string | object) => - tracingChannel(nameOrChannels as any) as any; - - // Patch TracingChannel.prototype.tracePromise to handle APIPromise and other - // Promise subclasses (mirrors the fix in hook.mts for the --import loader path). - patchTracingChannel(tracingChannel); - iso.getEnv = (name: string) => { if (typeof process === "undefined" || typeof process.env === "undefined") { return undefined; diff --git a/js/src/edge-light/config.ts b/js/src/edge-light/config.ts index fe84d53a8..20302d2cd 100644 --- a/js/src/edge-light/config.ts +++ b/js/src/edge-light/config.ts @@ -1,12 +1,6 @@ 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; @@ -27,13 +21,6 @@ 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; diff --git a/js/src/global-instrumentation-hooks.test.ts b/js/src/global-instrumentation-hooks.test.ts new file mode 100644 index 000000000..b1628d98c --- /dev/null +++ b/js/src/global-instrumentation-hooks.test.ts @@ -0,0 +1,232 @@ +import { AsyncLocalStorage } from "node:async_hooks"; +import { randomUUID } from "node:crypto"; +import { tracingChannel } from "node:diagnostics_channel"; +import { describe, expect, it } from "vitest"; +import { + GLOBAL_INSTRUMENTATION_HOOKS_KEY, + getGlobalTracingChannel, + newGlobalTracingChannel, +} from "./global-instrumentation-hooks"; + +function uniqueChannelName(label: string): string { + return `test:${label}:${randomUUID()}`; +} + +describe("global instrumentation hooks", () => { + it("installs a non-enumerable, immutable global registry", () => { + const name = uniqueChannelName("descriptor"); + const channel = newGlobalTracingChannel(name); + const descriptor = Object.getOwnPropertyDescriptor( + globalThis, + GLOBAL_INSTRUMENTATION_HOOKS_KEY, + ); + + expect(descriptor).toMatchObject({ + configurable: false, + enumerable: false, + writable: false, + }); + expect(descriptor?.value).toBeInstanceOf(Map); + expect(getGlobalTracingChannel(name)).toBe(channel); + expect(Object.keys(globalThis)).not.toContain( + GLOBAL_INSTRUMENTATION_HOOKS_KEY, + ); + }); + + it("shares subscriptions and supports complete unsubscription", () => { + const name = uniqueChannelName("subscriptions"); + const first = newGlobalTracingChannel>(name); + const second = newGlobalTracingChannel>(name); + const events: string[] = []; + const handlers = { + start: () => events.push("start"), + end: () => events.push("end"), + }; + + expect(first).toBe(second); + first.subscribe(handlers); + expect(second.hasSubscribers).toBe(true); + expect(second.traceSync(() => 42, {})).toBe(42); + expect(events).toEqual(["start", "end"]); + expect(second.unsubscribe(handlers)).toBe(true); + expect(first.hasSubscribers).toBe(false); + expect(second.unsubscribe(handlers)).toBe(false); + }); + + it("does not publish legacy diagnostics-channel events", () => { + const name = uniqueChannelName("hard-cutover"); + const diagnostics = tracingChannel(name); + const diagnosticsEvents: unknown[] = []; + const diagnosticsHandler = (message: unknown) => + diagnosticsEvents.push(message); + diagnostics.start.subscribe(diagnosticsHandler); + + const hook = newGlobalTracingChannel>(name); + const hookEvents: unknown[] = []; + hook.subscribe({ start: (message) => hookEvents.push(message) }); + hook.traceSync(() => "result", {}); + + expect(hookEvents).toHaveLength(1); + expect(diagnosticsEvents).toHaveLength(0); + diagnostics.start.unsubscribe(diagnosticsHandler); + }); + + it("mirrors sync lifecycle ordering and rethrows errors", () => { + const channel = newGlobalTracingChannel>( + uniqueChannelName("sync"), + ); + const lifecycle: string[] = []; + const contexts: Record[] = []; + channel.subscribe({ + start: (context) => { + lifecycle.push("start"); + contexts.push(context); + }, + end: (context) => { + lifecycle.push("end"); + contexts.push(context); + }, + error: (context) => { + lifecycle.push("error"); + contexts.push(context); + }, + }); + + const successContext: Record = {}; + expect(channel.traceSync(() => "result", successContext)).toBe("result"); + expect(successContext.result).toBe("result"); + expect(lifecycle).toEqual(["start", "end"]); + expect(contexts).toEqual([successContext, successContext]); + + lifecycle.length = 0; + contexts.length = 0; + const error = new Error("boom"); + const errorContext: Record = {}; + expect(() => + channel.traceSync(() => { + throw error; + }, errorContext), + ).toThrow(error); + expect(errorContext.error).toBe(error); + expect(lifecycle).toEqual(["start", "error", "end"]); + expect(contexts).toEqual([errorContext, errorContext, errorContext]); + }); + + it("mirrors promise lifecycle ordering for resolution and rejection", async () => { + const channel = newGlobalTracingChannel>( + uniqueChannelName("promise"), + ); + const lifecycle: string[] = []; + channel.subscribe({ + start: () => lifecycle.push("start"), + end: () => lifecycle.push("end"), + asyncStart: () => lifecycle.push("asyncStart"), + asyncEnd: () => lifecycle.push("asyncEnd"), + error: () => lifecycle.push("error"), + }); + + const successContext: Record = {}; + await expect( + channel.tracePromise(async () => "result", successContext), + ).resolves.toBe("result"); + expect(successContext.result).toBe("result"); + expect(lifecycle).toEqual(["start", "end", "asyncStart", "asyncEnd"]); + + lifecycle.length = 0; + const error = new Error("rejected"); + const errorContext: Record = {}; + await expect( + channel.tracePromise(async () => { + throw error; + }, errorContext), + ).rejects.toBe(error); + expect(errorContext.error).toBe(error); + expect(lifecycle).toEqual([ + "start", + "end", + "error", + "asyncStart", + "asyncEnd", + ]); + }); + + it("preserves promise subclasses and non-Promise return values", async () => { + class HelperPromise extends Promise { + withResponse(): Promise<{ data: T }> { + return this.then((data) => ({ data })); + } + } + + const channel = newGlobalTracingChannel>( + uniqueChannelName("promise-subclass"), + ); + const lifecycle: string[] = []; + channel.subscribe({ + end: () => lifecycle.push("end"), + asyncStart: () => lifecycle.push("asyncStart"), + asyncEnd: () => lifecycle.push("asyncEnd"), + }); + const original = new HelperPromise((resolve) => resolve("ok")); + const traced = channel.tracePromise(() => original, {}); + + expect(traced).toBe(original); + await expect(traced.withResponse()).resolves.toEqual({ data: "ok" }); + lifecycle.length = 0; + + const nonPromise = channel.tracePromise( + (() => 42) as unknown as () => PromiseLike, + {}, + ); + expect(nonPromise).toBe(42); + expect(lifecycle).toEqual(["end", "asyncStart", "asyncEnd"]); + }); + + it("wraps callbacks without changing arguments or receiver semantics", async () => { + const channel = newGlobalTracingChannel>( + uniqueChannelName("callback"), + ); + const lifecycle: string[] = []; + channel.subscribe({ + start: () => lifecycle.push("start"), + end: () => lifecycle.push("end"), + asyncStart: () => lifecycle.push("asyncStart"), + asyncEnd: () => lifecycle.push("asyncEnd"), + }); + + const receiver = { label: "receiver" }; + const result = await new Promise((resolve) => { + channel.traceCallback( + function (this: typeof receiver, value: string, callback: Function) { + expect(this).toBe(receiver); + callback.call(this, null, value); + return "immediate"; + }, + 1, + {}, + receiver, + "done", + function (this: typeof receiver, error: unknown, value: string) { + expect(this).toBe(receiver); + expect(error).toBeNull(); + resolve(value); + }, + ); + }); + + expect(result).toBe("done"); + expect(lifecycle).toEqual(["start", "asyncStart", "asyncEnd", "end"]); + }); + + it("runs traced functions inside bound stores", () => { + const channel = newGlobalTracingChannel>( + uniqueChannelName("stores"), + ); + const storage = new AsyncLocalStorage(); + channel.start.bindStore(storage, () => "bound"); + + expect(channel.hasSubscribers).toBe(true); + expect(channel.traceSync(() => storage.getStore(), {})).toBe("bound"); + expect(channel.start.unbindStore(storage)).toBe(true); + expect(channel.hasSubscribers).toBe(false); + }); +}); diff --git a/js/src/global-instrumentation-hooks.ts b/js/src/global-instrumentation-hooks.ts new file mode 100644 index 000000000..510230e73 --- /dev/null +++ b/js/src/global-instrumentation-hooks.ts @@ -0,0 +1,474 @@ +/* + * Adapted from Node.js diagnostics_channel's TracingChannel implementation. + * Copyright Node.js contributors. Licensed under the MIT License. + * See licenses/node-diagnostics-channel/LICENSE. + */ + +export const GLOBAL_INSTRUMENTATION_HOOKS_KEY = + "__braintrust_instrumentation_hooks"; + +export interface GlobalHookAsyncLocalStorage { + enterWith(store: T): void; + run(store: T | undefined, callback: () => R): R; + getStore(): T | undefined; +} + +export type GlobalHookMessageFunction< + M = any, + N extends string | symbol = string, +> = (message: M, name: N) => void; + +export type GlobalHookTransformFunction = (message: M) => S; + +export interface GlobalHookChannel< + M = any, + N extends string | symbol = string, +> { + readonly name: N; + readonly hasSubscribers: boolean; + subscribe(subscription: GlobalHookMessageFunction): void; + unsubscribe(subscription: GlobalHookMessageFunction): boolean; + bindStore( + store: GlobalHookAsyncLocalStorage, + transform?: GlobalHookTransformFunction, + ): void; + unbindStore(store: GlobalHookAsyncLocalStorage): boolean; + publish(message: M): void; + runStores any>( + message: M, + fn: F, + thisArg?: ThisParameterType, + ...args: Parameters + ): ReturnType; +} + +export interface GlobalTracingChannelCollection { + readonly start?: GlobalHookChannel; + readonly end?: GlobalHookChannel; + readonly asyncStart?: GlobalHookChannel; + readonly asyncEnd?: GlobalHookChannel; + readonly error?: GlobalHookChannel; +} + +export interface GlobalHookHandlers { + start?: (context: M, name: string) => void; + end?: (context: M, name: string) => void; + asyncStart?: (context: M, name: string) => void; + asyncEnd?: (context: M, name: string) => void; + error?: (context: M, name: string) => void; +} + +export interface GlobalTracingChannel< + M = any, +> extends GlobalTracingChannelCollection { + readonly start: GlobalHookChannel; + readonly end: GlobalHookChannel; + readonly asyncStart: GlobalHookChannel; + readonly asyncEnd: GlobalHookChannel; + readonly error: GlobalHookChannel; + readonly hasSubscribers: boolean; + subscribe(handlers: GlobalHookHandlers): void; + unsubscribe(handlers: GlobalHookHandlers): boolean; + traceSync any>( + fn: F, + message?: M, + thisArg?: ThisParameterType, + ...args: Parameters + ): ReturnType; + tracePromise PromiseLike>( + fn: F, + message?: M, + thisArg?: ThisParameterType, + ...args: Parameters + ): ReturnType; + traceCallback any>( + fn: F, + position?: number, + message?: M, + thisArg?: ThisParameterType, + ...args: Parameters + ): ReturnType; +} + +type StoreEntry = [ + GlobalHookAsyncLocalStorage, + GlobalHookTransformFunction | undefined, +]; + +function reportError(error: unknown): void { + queueMicrotask(() => { + throw error; + }); +} + +function defaultTransform(message: M): M { + return message; +} + +function wrapStoreRun( + store: GlobalHookAsyncLocalStorage, + message: M, + next: () => unknown, + transform: GlobalHookTransformFunction = defaultTransform, +): () => unknown { + return () => { + let context: unknown; + try { + context = transform(message); + } catch (error) { + reportError(error); + return next(); + } + return store.run(context, next); + }; +} + +class HookChannel< + M, + N extends string | symbol = string, +> implements GlobalHookChannel { + private subscribers: GlobalHookMessageFunction[] = []; + private stores = new Map< + GlobalHookAsyncLocalStorage, + GlobalHookTransformFunction | undefined + >(); + + constructor(readonly name: N) {} + + get hasSubscribers(): boolean { + return this.subscribers.length > 0 || this.stores.size > 0; + } + + subscribe(subscription: GlobalHookMessageFunction): void { + if (typeof subscription !== "function") { + throw new TypeError("subscription must be a function"); + } + this.subscribers = [...this.subscribers, subscription]; + } + + unsubscribe(subscription: GlobalHookMessageFunction): boolean { + const index = this.subscribers.indexOf(subscription); + if (index === -1) { + return false; + } + this.subscribers = [ + ...this.subscribers.slice(0, index), + ...this.subscribers.slice(index + 1), + ]; + return true; + } + + bindStore( + store: GlobalHookAsyncLocalStorage, + transform?: GlobalHookTransformFunction, + ): void { + if (!store || typeof store.run !== "function") { + throw new TypeError("store must have a run method"); + } + this.stores.set( + store as GlobalHookAsyncLocalStorage, + transform as GlobalHookTransformFunction | undefined, + ); + } + + unbindStore(store: GlobalHookAsyncLocalStorage): boolean { + return this.stores.delete(store as GlobalHookAsyncLocalStorage); + } + + publish(message: M): void { + const subscribers = this.subscribers; + for (const subscriber of subscribers) { + try { + subscriber(message, this.name); + } catch (error) { + reportError(error); + } + } + } + + runStores any>( + message: M, + fn: F, + thisArg?: ThisParameterType, + ...args: Parameters + ): ReturnType { + let run = () => { + this.publish(message); + return Reflect.apply(fn, thisArg, args); + }; + for (const [store, transform] of this.stores.entries() as Iterable< + StoreEntry + >) { + run = wrapStoreRun(store, message, run, transform); + } + return run() as ReturnType; + } +} + +const traceEvents = [ + "start", + "end", + "asyncStart", + "asyncEnd", + "error", +] as const; + +class TracingHook implements GlobalTracingChannel { + readonly start: GlobalHookChannel; + readonly end: GlobalHookChannel; + readonly asyncStart: GlobalHookChannel; + readonly asyncEnd: GlobalHookChannel; + readonly error: GlobalHookChannel; + + constructor(nameOrChannels: string | GlobalTracingChannelCollection) { + if (typeof nameOrChannels === "string") { + this.start = new HookChannel(`tracing:${nameOrChannels}:start`); + this.end = new HookChannel(`tracing:${nameOrChannels}:end`); + this.asyncStart = new HookChannel(`tracing:${nameOrChannels}:asyncStart`); + this.asyncEnd = new HookChannel(`tracing:${nameOrChannels}:asyncEnd`); + this.error = new HookChannel(`tracing:${nameOrChannels}:error`); + return; + } + + this.start = nameOrChannels.start ?? new HookChannel("tracing:start"); + this.end = nameOrChannels.end ?? new HookChannel("tracing:end"); + this.asyncStart = + nameOrChannels.asyncStart ?? new HookChannel("tracing:asyncStart"); + this.asyncEnd = + nameOrChannels.asyncEnd ?? new HookChannel("tracing:asyncEnd"); + this.error = nameOrChannels.error ?? new HookChannel("tracing:error"); + } + + get hasSubscribers(): boolean { + return ( + this.start.hasSubscribers || + this.end.hasSubscribers || + this.asyncStart.hasSubscribers || + this.asyncEnd.hasSubscribers || + this.error.hasSubscribers + ); + } + + subscribe(handlers: GlobalHookHandlers): void { + for (const eventName of traceEvents) { + const handler = handlers[eventName]; + if (handler) { + this[eventName].subscribe(handler); + } + } + } + + unsubscribe(handlers: GlobalHookHandlers): boolean { + let done = true; + for (const eventName of traceEvents) { + const handler = handlers[eventName]; + if (handler && !this[eventName].unsubscribe(handler)) { + done = false; + } + } + return done; + } + + traceSync any>( + fn: F, + message: M = {} as M, + thisArg?: ThisParameterType, + ...args: Parameters + ): ReturnType { + if (!this.hasSubscribers) { + return Reflect.apply(fn, thisArg, args) as ReturnType; + } + + const context = message as Record; + return this.start.runStores(message, () => { + try { + const result = Reflect.apply(fn, thisArg, args); + context.result = result; + return result; + } catch (error) { + context.error = error; + this.error.publish(message); + throw error; + } finally { + this.end.publish(message); + } + }); + } + + tracePromise PromiseLike>( + fn: F, + message: M = {} as M, + thisArg?: ThisParameterType, + ...args: Parameters + ): ReturnType { + if (!this.hasSubscribers) { + return Reflect.apply(fn, thisArg, args) as ReturnType; + } + + const context = message as Record; + return this.start.runStores(message, () => { + let ended = false; + try { + const result = Reflect.apply(fn, thisArg, args); + this.end.publish(message); + ended = true; + + if ( + !result || + (typeof result !== "object" && typeof result !== "function") || + typeof result.then !== "function" + ) { + context.result = result; + this.asyncStart.publish(message); + this.asyncEnd.publish(message); + return result; + } + + const resolve = (resolved: unknown) => { + context.result = resolved; + this.asyncStart.publish(message); + this.asyncEnd.publish(message); + return resolved; + }; + const reject = (error: unknown) => { + context.error = error; + this.error.publish(message); + this.asyncStart.publish(message); + this.asyncEnd.publish(message); + throw error; + }; + + if (result instanceof Promise && result.constructor === Promise) { + return result.then(resolve, reject); + } + + void result.then(resolve, (error: unknown) => { + try { + reject(error); + } catch { + // The original promise-like object is returned below. Keep the + // instrumentation side-chain from changing its rejection behavior. + } + }); + return result; + } catch (error) { + context.error = error; + this.error.publish(message); + if (!ended) { + this.end.publish(message); + } + throw error; + } + }) as ReturnType; + } + + traceCallback any>( + fn: F, + position = -1, + message: M = {} as M, + thisArg?: ThisParameterType, + ...args: Parameters + ): ReturnType { + if (!this.hasSubscribers) { + return Reflect.apply(fn, thisArg, args); + } + + const context = message as Record; + const callArgs = + args.length > 0 + ? args + : ((context.arguments as ArrayLike | undefined) ?? args); + const callback = Array.prototype.at.call(callArgs, position); + if (typeof callback !== "function") { + return Reflect.apply(fn, thisArg, args); + } + + const { asyncStart, asyncEnd, error: errorChannel } = this; + function wrappedCallback(this: unknown, error: unknown, result: unknown) { + if (error) { + context.error = error; + errorChannel.publish(message); + } else { + context.result = result; + } + + return asyncStart.runStores(message, () => { + try { + return Reflect.apply(callback, this, arguments); + } finally { + asyncEnd.publish(message); + } + }); + } + + Array.prototype.splice.call(callArgs, position, 1, wrappedCallback); + return this.start.runStores(message, () => { + try { + return Reflect.apply(fn, thisArg, args); + } catch (error) { + context.error = error; + this.error.publish(message); + throw error; + } finally { + this.end.publish(message); + } + }); + } +} + +type HookRegistry = Map>; + +const fallbackRegistry: HookRegistry = new Map(); + +function getHookRegistry(): HookRegistry { + const target = globalThis as Record; + const existing = target[GLOBAL_INSTRUMENTATION_HOOKS_KEY]; + if (existing instanceof Map) { + return existing as HookRegistry; + } + if (existing !== undefined) { + return fallbackRegistry; + } + + const registry: HookRegistry = new Map(); + try { + Object.defineProperty(globalThis, GLOBAL_INSTRUMENTATION_HOOKS_KEY, { + configurable: false, + enumerable: false, + value: registry, + writable: false, + }); + return registry; + } catch { + return fallbackRegistry; + } +} + +export function newGlobalTracingChannel( + nameOrChannels: string | GlobalTracingChannelCollection, +): GlobalTracingChannel { + if (typeof nameOrChannels !== "string") { + return new TracingHook(nameOrChannels); + } + + const registry = getHookRegistry(); + const existing = registry.get(nameOrChannels); + if (existing) { + return existing as GlobalTracingChannel; + } + + const hook = new TracingHook(nameOrChannels); + registry.set(nameOrChannels, hook); + return hook; +} + +export function getGlobalTracingChannel( + name: string, +): GlobalTracingChannel | undefined { + const registry = (globalThis as Record)[ + GLOBAL_INSTRUMENTATION_HOOKS_KEY + ]; + return registry instanceof Map + ? (registry.get(name) as GlobalTracingChannel | undefined) + : undefined; +} diff --git a/js/src/instrumentation/README.md b/js/src/instrumentation/README.md index 8ed64dfa6..9fe703090 100644 --- a/js/src/instrumentation/README.md +++ b/js/src/instrumentation/README.md @@ -1,1407 +1,160 @@ -# Writing Plugins for Orchestrion Instrumentation +# Writing Braintrust Instrumentation Plugins -This guide explains how to write plugins that consume diagnostics channel events from orchestrion instrumentations. +Braintrust instrumentation plugins consume tracing-compatible events from the +internal global hook registry and convert them into spans. Auto-instrumented +provider code and manual wrappers use the same typed channels, so extraction, +stream handling, and span behavior stay aligned. -## Table of Contents +## Architecture -- [Overview](#overview) -- [Understanding Orchestrion](#understanding-orchestrion) -- [Diagnostics Channel API](#diagnostics-channel-api) -- [Plugin Architecture](#plugin-architecture) -- [Event Lifecycle](#event-lifecycle) -- [Stream Handling](#stream-handling) -- [Complete Examples](#complete-examples) -- [Best Practices](#best-practices) -- [Testing](#testing) +An instrumentation has four parts: ---- +1. An Orchestrion config identifies the provider function for automatic + transformation. +2. A typed channel defines its arguments, result, extra event fields, and stable + `orchestrion::` identifier. +3. A plugin subscribes to that channel and maps events into Braintrust spans. +4. A manual wrapper invokes the same typed channel when transformation is not + available. -## Overview +The global hook transport is internal. Plugins should use `defineChannels`, +`traceAsyncChannel`, `traceStreamingChannel`, `traceSyncStreamChannel`, or +`BasePlugin` helpers rather than reading the global property directly. -### What is Orchestrion? +## Lifecycle -[Orchestrion](https://github.com/nodejs/orchestrion-js) is a code transformation library that automatically instruments JavaScript/TypeScript code at bundle-time or load-time. It injects calls to Node.js's [diagnostics_channel API](https://nodejs.org/api/diagnostics_channel.html) to enable observability without manual code changes that may be enabled or disabled at any time. +The event lifecycle is compatible with Node tracing channels: -**Key concepts:** +- `start`: before the synchronous portion of the target function +- `end`: after the synchronous portion completes +- `asyncStart`: when an asynchronous result begins settling +- `asyncEnd`: after that result settles and before user continuation +- `error`: when the target throws, rejects, or reports a callback error -- **FunctionQuery** - Declarative pattern for targeting functions to instrument -- **Transformation** - SWC-based AST modification at compile/load time -- **Diagnostics Channels** - Standard Node.js API for publishing instrumentation events +Every phase receives the same mutable context object: -### How It Works - -1. **Configuration** - You specify which functions to instrument via JSON config -2. **Transformation** - Orchestrion modifies code to publish to diagnostics channels -3. **Subscription** - Your plugin subscribes to those channels and handles events - -**Example flow:** - -``` -User calls: client.chat.completions.create(params) - ↓ -Orchestrion publishes: start event - ↓ -Your plugin subscriber receives: { arguments: [params] } - ↓ -Synchronous portion executes (function body runs, returns promise) - ↓ -Orchestrion publishes: end event with the promise - ↓ -User awaits the promise - ↓ -Promise begins to settle - ↓ -Orchestrion publishes: asyncStart event - ↓ -Promise settles with the result - ↓ -Orchestrion publishes: asyncEnd event - ↓ -Your plugin subscriber receives: { result: response } - ↓ -User's await completes, user code continues with result -``` - ---- - -## Understanding Orchestrion - -### Instrumentation Configuration - -Orchestrion uses JSON configs to specify what to instrument: - -```json -{ - "channelName": "orchestrion:openai:chat.completions.create", - "module": { - "name": "openai", - "versionRange": ">=4.0.0 <6.0.0", - "filePath": "resources/chat/completions.mjs" - }, - "functionQuery": { - "ClassMethod": { - "className": "Completions", - "methodName": "create", - "kind": "Async" - } - } -} -``` - -### What Orchestrion Generates - -For the above config, orchestrion transforms the target function to: - -```javascript -// Before transformation -class Completions { - async create(params) { - // ... implementation - } -} - -// After orchestrion transformation -class Completions { - async create(params) { - const __apm$original_args = arguments; - const __apm$traced = async () => { - const __apm$wrapped = async (params) => { - // ... original implementation - }; - return __apm$wrapped.apply(null, __apm$original_args); - }; - - if (!tr_ch_apm$create.hasSubscribers) { - return __apm$traced(); - } - - return tr_ch_apm$create.tracePromise(__apm$traced, { - arguments, - self: this, - moduleVersion: "5.0.0", - }); - } -} -``` - -The `tracePromise` function publishes events to the diagnostics channel. - ---- - -## Diagnostics Channel API - -### Channel Naming Convention - -Orchestrion uses the format: `orchestrion::` - -Examples: - -- `orchestrion:openai:chat.completions.create` -- `orchestrion:anthropic:messages.create` -- `orchestrion:vercel-ai:generateText` - -### Event Types - -For async functions, orchestrion uses `tracingChannel.tracePromise()` which publishes events at specific points: - -#### 1. Start Event - -Published when the function is called, before the synchronous portion executes. - -```typescript -interface StartEvent { - // Function arguments (from 'arguments' keyword) - arguments: ArrayLike; - - // The 'this' context - self?: unknown; - - // Module version (from package.json) - moduleVersion?: string; -} -``` - -#### 2. End Event - -Published after the synchronous portion completes (when the promise is returned). - -```typescript -interface EndEvent { - // Same context object from start +```ts +interface InstrumentationContext { arguments: ArrayLike; self?: unknown; moduleVersion?: string; + result?: unknown; + error?: unknown; } ``` -#### 3. AsyncStart Event - -Published when the promise begins to settle (async continuation starts). - -```typescript -interface AsyncStartEvent { - // Same context object - arguments: ArrayLike; - self?: unknown; - moduleVersion?: string; - - // The resolved value (BY REFERENCE - can be mutated!) - result: unknown; -} -``` - -#### 4. AsyncEnd Event - -Published when the promise finishes settling, **before control returns to user code**. - -```typescript -interface AsyncEndEvent { - // Same context object - arguments: ArrayLike; - self?: unknown; - moduleVersion?: string; - - // The resolved value (BY REFERENCE - can be mutated!) - result: unknown; -} -``` - -**Important:** The `result` is passed by reference. If it's an object (including streams), you can mutate it **before** the user's code continues. - -#### 5. Error Event - -Published if the function throws an error or the promise rejects. - -```typescript -interface ErrorEvent { - // Same context object - arguments: ArrayLike; - self?: unknown; - moduleVersion?: string; - - // The error that was thrown or rejection reason - error: Error; -} -``` - -### Subscribing to Channels - -Use the `tracingChannel` API from `dc-browser` (or Node.js `diagnostics_channel`): - -```typescript -import { tracingChannel } from "dc-browser"; - -const channel = tracingChannel("orchestrion:openai:chat.completions.create"); - -channel.subscribe({ - start: (event) => { - console.log("Function called with:", event.arguments); - }, - - end: (event) => { - console.log("Synchronous portion complete, promise returned"); - }, - - asyncStart: (event) => { - console.log("Promise beginning to settle"); - }, - - asyncEnd: (event) => { - console.log("Promise settled with:", event.result); - }, - - error: (event) => { - console.log("Function threw or promise rejected:", event.error); - }, -}); -``` - ---- - -## Plugin Architecture - -### BasePlugin Class - -Extend `BasePlugin` to create a plugin with lifecycle management: - -```typescript -import { BasePlugin } from "./core"; -import { tracingChannel } from "dc-browser"; - -export class MyPlugin extends BasePlugin { - private unsubscribers: Array<() => void> = []; - - protected onEnable(): void { - // Called when plugin is enabled - // Subscribe to channels here - this.subscribeToOpenAI(); - } - - protected onDisable(): void { - // Called when plugin is disabled - // Clean up subscriptions - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - } - - private subscribeToOpenAI(): void { - const channel = tracingChannel( - "orchestrion:openai:chat.completions.create", - ); - - const handlers = { - asyncStart: (event) => { - /* ... */ - }, - asyncEnd: (event) => { - /* ... */ - }, - error: (event) => { - /* ... */ - }, - }; - - channel.subscribe(handlers); - - // Store unsubscribe function for cleanup - this.unsubscribers.push(() => { - channel.unsubscribe(handlers); - }); - } -} -``` - -### Plugin Lifecycle - -```typescript -const plugin = new MyPlugin(); - -// Enable the plugin (calls onEnable) -plugin.enable(); - -// ... plugin is active, handling events ... - -// Disable the plugin (calls onDisable) -plugin.disable(); -``` - ---- - -## Event Lifecycle - -### Correlating Events - -Each invocation gets a unique context object shared across all events. Your may store your own data on the event object, though it's recommended to use a symbol property so as to not interfere with other consumers. You may also use a `WeakMap` to correlate start/end/asyncStart/asyncEnd/error events, if you prefer to avoid visibility to other consumers: - -```typescript -import { tracingChannel } from "dc-browser"; +The generated wrapper creates `arguments`, `self`, and `moduleVersion`. +Tracing operators add `result` or `error`. -const channel = tracingChannel("orchestrion:openai:chat.completions.create"); -const spans = new WeakMap(); +## Defining Typed Channels -channel.subscribe({ - start: (event) => { - // Create span and store it keyed by context object - const span = startSpan({ name: "Chat Completion" }); - spans.set(event, span); +Define the smallest types needed by instrumentation: - // or: - // event.span = span; +```ts +const providerChannels = defineChannels( + "provider-package", + { + create: channel< + [CreateParams], + CreateResult, + { providerRequestId?: string } + >({ + channelName: "messages.create", + kind: "async", + }), }, - - // end: (event) => { - // // Usually not needed - promise hasn't settled yet - // }, - - // asyncStart: (event) => { - // // Usually not needed - just marks promise settlement start - // }, - - asyncEnd: (event) => { - // Retrieve span using same context object - const span = spans.get(event); - // or... - // const span = event.span; - if (!span) return; - - span.log({ output: event.result }); - span.end(); - - // Clean up (optional for WeakMap, but recommended) - spans.delete(event); - }, - - error: (event) => { - const span = spans.get(event); - // or... - // const span = event.span; - if (!span) return; - - span.log({ error: event.error.message }); - span.end(); - - // Clean up (optional for WeakMap, but recommended) - spans.delete(event); - }, -}); -``` - -**Why WeakMap?** - -- Automatic garbage collection when context object is no longer referenced -- Does not tie lifetime of span to lifetime of event object -- No memory leaks from forgotten spans -- Fast O(1) lookup - -**Common Pattern:** - -- Use **start** to create the span and extract input -- Use **asyncEnd** to extract output and finalize the span (this fires before user code continues) -- Use **error** to handle failures - -### Extracting Data from Events - -#### From start - -Extract input data and metadata: - -```typescript -start: (event) => { - const self = event.self; // The `this` in the original call - const moduleVersion = event.moduleVersion; // The version of the module that triggered the event - const params = event.arguments[0]; // First argument - - // For OpenAI - const { messages, model, temperature, ...rest } = params || {}; - - span.log({ - input: messages, - metadata: { model, temperature, ...rest, provider: "openai" }, - }); -}; + { instrumentationName: "provider" }, +); ``` -#### From asyncEnd +Channel names must match the Orchestrion config exactly. Do not include the +`orchestrion:` prefix in the transform config; `defineChannels` and Orchestrion +construct it from the package and operation. -Extract output data and metrics: +## Subscribing -```typescript -asyncEnd: (event) => { - const result = event.result; +Prefer the shared tracing helpers: - // For OpenAI non-streaming - span.log({ - output: result.choices, - metrics: { - tokens: result.usage?.total_tokens, - prompt_tokens: result.usage?.prompt_tokens, - completion_tokens: result.usage?.completion_tokens, +```ts +this.register( + traceAsyncChannel(providerChannels.create, { + name: "provider.messages.create", + type: "llm", + extractInput(args) { + return { + input: args[0].messages, + metadata: { model: args[0].model }, + }; }, - }); -}; -``` - -**Important:** `asyncEnd` fires when the promise settles but **before** user code continues after the await. This is the perfect time to: - -- Extract the resolved value -- Patch streams (if the result is an async iterable) -- Finalize the span - ---- - -## Stream Handling - -### The Challenge - -When a function returns a stream (async iterable), the `asyncEnd` event gives you the stream object, but: - -- ❌ You cannot replace the return value -- ❌ You cannot iterate it (would consume it for the user) -- ❌ You cannot clone it (not all streams are cloneable) - -### The Solution: Stream Patching - -Instead of replacing the stream, **mutate it in-place** by patching its `Symbol.asyncIterator` method. - -**Key insight:** `asyncEnd` fires when the promise settles with the stream object, but **before** control returns to user code after the `await`. This gives us a window to patch the stream before the user iterates it. - -**Timing for streaming calls:** - -```typescript -// User code: -const stream = await client.chat.completions.create({ stream: true }); -// ↑ -// 1. start event (function called) -// 2. Synchronous portion executes -// 3. end event (promise returned) -// 4. Promise begins to settle -// 5. asyncStart event -// 6. Promise settles with stream object -// 7. asyncEnd event (WE PATCH HERE!) -// 8. Wrapped promise resolves -// ← User's await completes, gets patched stream -for await (const chunk of stream) { - // User iterates, our patched iterator collects chunks -} -``` - -### Using wrapStreamResult (High-Level API) - -The easiest way to handle streaming responses: - -```typescript -import { wrapStreamResult } from "./core"; - -channel.subscribe({ - asyncEnd: (event) => { - const { span, startTime } = spans.get(event); - - wrapStreamResult(event.result, { - // Process chunks (for streaming responses) - processChunks: (chunks) => { - const output = combineChunks(chunks); - const metrics = extractMetrics(chunks, startTime); - return { output, metrics }; - }, - - // Process result (for non-streaming responses) - onNonStream: (result) => { - const output = result.choices; - const metrics = extractMetrics(result, startTime); - return { output, metrics }; - }, - - // Called with processed result (both streaming and non-streaming) - onResult: (processed) => { - span.log(processed); - span.end(); - spans.delete(event); - }, - - // Error handling - onError: (error, chunks) => { - span.log({ - error: error.message, - partial_chunks: chunks.length, - }); - span.end(); - spans.delete(event); - }, - }); - }, -}); -``` - -### Using patchStreamIfNeeded (Low-Level API) - -For more control: - -```typescript -import { patchStreamIfNeeded, isAsyncIterable } from "./core"; - -channel.subscribe({ - asyncEnd: (event) => { - const { span } = spans.get(event); - - if (isAsyncIterable(event.result)) { - // It's a stream - patch it - patchStreamIfNeeded(event.result, { - // Called for each chunk as it's yielded - onChunk: (chunk) => { - console.log("Received chunk:", chunk); - }, - - // Called when stream completes - onComplete: (chunks) => { - const output = processChunks(chunks); - span.log({ output, metrics: { chunks: chunks.length } }); - span.end(); - }, - - // Called if stream errors - onError: (error, chunks) => { - span.log({ error: error.message, partial_chunks: chunks.length }); - span.end(); - }, - - // Optional: filter which chunks to collect - shouldCollect: (chunk) => { - // Only collect content chunks, skip metadata - return chunk.type !== "metadata"; - }, - }); - } else { - // Non-streaming response - span.log({ output: event.result }); - span.end(); - } - - spans.delete(event); - }, -}); -``` - -### How Stream Patching Works - -The helper mutates the stream object by replacing its `Symbol.asyncIterator` method: - -```typescript -// Original stream -const stream = event.result; - -// Patch the iterator method -const originalIterator = stream[Symbol.asyncIterator]; -stream[Symbol.asyncIterator] = function () { - const iterator = originalIterator.call(this); - const originalNext = iterator.next.bind(iterator); - const chunks = []; - - // Wrap the next() method - iterator.next = async function () { - const result = await originalNext(); - - if (!result.done) { - chunks.push(result.value); // Collect chunk - } else { - // Stream complete - log final output - span.log({ output: processChunks(chunks) }); - span.end(); - } - - return result; // Pass through to user unchanged - }; - - return iterator; -}; - -// User's code now uses our patched iterator -for await (const chunk of stream) { - console.log(chunk); // Works normally -} -``` - -**Result:** User code is unchanged, but we collect all chunks transparently. - -### Edge Cases Handled - -The stream patching helpers handle: - -1. **Frozen/Sealed Objects** - Detects and warns if stream cannot be patched -2. **Early Cancellation** - Patches `iterator.return()` to handle `break` -3. **Error Injection** - Patches `iterator.throw()` to handle errors -4. **Double-Patching** - Prevents patching the same stream twice - ---- - -## Complete Examples - -### Example 1: OpenAI Plugin (Full Implementation) - -```typescript -import { BasePlugin } from "./core"; -import { tracingChannel } from "dc-browser"; -import { startSpan, Span } from "../logger"; -import { wrapStreamResult } from "./core"; -import { SpanTypeAttribute } from "../../util/index"; -import { getCurrentUnixTimestamp } from "../util"; - -export class OpenAIPlugin extends BasePlugin { - private unsubscribers: Array<() => void> = []; - - protected onEnable(): void { - this.subscribeToOpenAI(); - } - - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - } - - private subscribeToOpenAI(): void { - // Chat completions - this.subscribeToChannel("orchestrion:openai:chat.completions.create", { - name: "Chat Completion", - type: SpanTypeAttribute.LLM, - extractInput: (args) => { - const params = args[0] || {}; - const { messages, ...metadata } = params; - return { - input: messages, - metadata: { ...metadata, provider: "openai" }, - }; - }, - processChunks: (chunks) => { - const content = chunks - .map((chunk: any) => chunk.choices?.[0]?.delta?.content || "") - .filter(Boolean) - .join(""); - - const lastChunk = chunks[chunks.length - 1]; - return { - output: [{ message: { content } }], - usage: lastChunk?.usage, - }; - }, - processNonStream: (result: any) => ({ - output: result.choices, - usage: result.usage, - }), - }); - - // Embeddings - this.subscribeToChannel("orchestrion:openai:embeddings.create", { - name: "Embedding", - type: SpanTypeAttribute.LLM, - extractInput: (args) => { - const params = args[0] || {}; - const { input, ...metadata } = params; - return { input, metadata: { ...metadata, provider: "openai" } }; - }, - processChunks: (chunks) => { - // Embeddings don't stream, but handle just in case - return { output: chunks, usage: chunks[0]?.usage }; - }, - processNonStream: (result: any) => ({ - output: result.data?.map((d: any) => d.embedding), - usage: result.usage, - }), - }); - } - - private subscribeToChannel( - channelName: string, - config: { - name: string; - type: string; - extractInput: (args: any[]) => { input: any; metadata: any }; - processChunks: (chunks: any[]) => { output: any; usage?: any }; - processNonStream: (result: any) => { output: any; usage?: any }; + extractOutput(result) { + return result.content; }, - ): void { - const channel = tracingChannel(channelName); - const spans = new WeakMap(); - - const handlers = { - start: (event: any) => { - const span = startSpan({ - name: config.name, - spanAttributes: { type: config.type }, - }); - - const startTime = getCurrentUnixTimestamp(); - spans.set(event, { span, startTime }); - - try { - const { input, metadata } = config.extractInput( - Array.from(event.arguments), - ); - span.log({ input, metadata }); - } catch (error) { - console.error(`Error extracting input for ${channelName}:`, error); - } - }, - - asyncEnd: (event: any) => { - const spanData = spans.get(event); - if (!spanData) return; - - const { span, startTime } = spanData; - - wrapStreamResult(event.result, { - processChunks: config.processChunks, - onNonStream: config.processNonStream, - onResult: (processed: any) => { - const metrics: any = {}; - - if (processed.usage) { - metrics.tokens = processed.usage.total_tokens; - metrics.prompt_tokens = processed.usage.prompt_tokens; - metrics.completion_tokens = processed.usage.completion_tokens; - } - - metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; - - span.log({ output: processed.output, metrics }); - span.end(); - spans.delete(event); - }, - onError: (error: Error) => { - span.log({ error: error.message }); - span.end(); - spans.delete(event); - }, - }); - }, - - error: (event: any) => { - const spanData = spans.get(event); - if (!spanData) return; - - spanData.span.log({ error: event.error.message }); - spanData.span.end(); - spans.delete(event); - }, - }; - - channel.subscribe(handlers); - this.unsubscribers.push(() => channel.unsubscribe(handlers)); - } -} -``` - -### Example 2: Anthropic Plugin - -```typescript -import { BasePlugin } from "./core"; -import { tracingChannel } from "dc-browser"; -import { startSpan } from "../logger"; -import { wrapStreamResult } from "./core"; -import { SpanTypeAttribute } from "../../util/index"; - -export class AnthropicPlugin extends BasePlugin { - private unsubscribers: Array<() => void> = []; - - protected onEnable(): void { - this.subscribeToAnthropic(); - } - - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - } - - private subscribeToAnthropic(): void { - const channel = tracingChannel("orchestrion:anthropic:messages.create"); - const spans = new WeakMap(); - - channel.subscribe({ - start: (event: any) => { - const span = startSpan({ - name: "Anthropic Message", - spanAttributes: { type: SpanTypeAttribute.LLM }, - }); - - const params = event.arguments[0] || {}; - - // Coalesce messages and system prompt - const input = [ - ...(params.system - ? [{ role: "system", content: params.system }] - : []), - ...(params.messages || []), - ]; - - const { messages, system, ...metadata } = params; - - span.log({ - input, - metadata: { ...metadata, provider: "anthropic" }, - }); - - spans.set(event, { span, startTime: Date.now() }); - }, - - asyncEnd: (event: any) => { - const spanData = spans.get(event); - if (!spanData) return; - - const { span, startTime } = spanData; - - wrapStreamResult(event.result, { - processChunks: (chunks: any[]) => { - // Extract text from content_block_delta events - const textChunks = chunks.filter( - (chunk) => - chunk.type === "content_block_delta" && - chunk.delta?.type === "text_delta", - ); - - const content = textChunks - .map((chunk) => chunk.delta.text) - .join(""); - - // Find usage in message_stop event - const stopEvent = chunks.find((c) => c.type === "message_stop"); - const usage = stopEvent?.message?.usage; - - return { - output: [{ type: "text", text: content }], - metrics: { - input_tokens: usage?.input_tokens, - output_tokens: usage?.output_tokens, - chunks: chunks.length, - time_to_first_token: Date.now() - startTime, - }, - }; - }, - - onNonStream: (result: any) => ({ - output: result.content, - metrics: { - input_tokens: result.usage?.input_tokens, - output_tokens: result.usage?.output_tokens, - time_to_first_token: Date.now() - startTime, - }, - }), - - onResult: (processed: any) => { - span.log(processed); - span.end(); - spans.delete(event); - }, - - onError: (error: Error) => { - span.log({ error: error.message }); - span.end(); - spans.delete(event); - }, - }); - }, - - error: (event: any) => { - const spanData = spans.get(event); - if (!spanData) return; - - spanData.span.log({ error: event.error.message }); - spanData.span.end(); - spans.delete(event); - }, - }); - - this.unsubscribers.push(() => - channel.unsubscribe({ - asyncStart: () => {}, - asyncEnd: () => {}, - error: () => {}, - }), - ); - } -} + extractMetrics(result) { + return { + prompt_tokens: result.usage.input_tokens, + completion_tokens: result.usage.output_tokens, + }; + }, + }), +); ``` -### Example 3: Simple Function Instrumentation +The helpers: -For standalone functions (not class methods): +- create and correlate spans with a `WeakMap` keyed by event context +- bind the current span store to `start` for async-context propagation +- contain extraction failures and log them through `debugLogger` +- patch streams without replacing their public semantics +- unsubscribe and unbind stores when a plugin is disabled -```typescript -// Instrumenting a standalone async function -const channel = tracingChannel("orchestrion:my-lib:fetchData"); +Use raw `IsoChannelHandlers` only when a provider requires lifecycle behavior +that the shared helpers cannot express. -channel.subscribe({ - asyncStart: (event: any) => { - const [url, options] = event.arguments; - console.log("Fetching:", url, options); - }, +## Manual Wrappers - asyncEnd: (event: any) => { - console.log("Fetched:", event.result); - }, +Manual wrappers call the same typed channel: - error: (event: any) => { - console.error("Fetch failed:", event.error); - }, +```ts +return providerChannels.create.tracePromise(() => originalCreate(params), { + arguments: [params], }); ``` ---- +Do not create spans directly inside wrappers. Keeping span creation in the +plugin prevents auto and manual instrumentation from drifting. -## Best Practices +## Promise and Stream Requirements -### 1. Always Clean Up Subscriptions +Instrumentation is non-invasive: -```typescript -class MyPlugin extends BasePlugin { - private unsubscribers: Array<() => void> = []; +- Native promises retain normal resolution and rejection behavior. +- Promise subclasses and other thenables are returned unchanged so helper + methods such as `withResponse()` remain available. +- A non-Promise value returned from an `Async` transform remains that value. +- Async iterables and event-emitter streams retain identity and public methods. +- Subscriber or extraction bugs must not alter provider calls. - protected onDisable(): void { - // ALWAYS unsubscribe to prevent memory leaks - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - } -} -``` +Stream patches must be idempotent and preserve cancellation, errors, early +termination, and async context. -### 2. Use WeakMap for Event Correlation +## Event and Span Safety -```typescript -// ✅ Good - automatic garbage collection -const spans = new WeakMap(); - -// ❌ Bad - manual cleanup required, risk of memory leaks -const spans = new Map(); -``` - -### 3. Handle Missing Span Data Gracefully - -```typescript -asyncEnd: (event) => { - const spanData = spans.get(event); - if (!spanData) { - // This can happen if asyncStart wasn't called - // or if the span was already cleaned up - return; - } - // ... rest of handler -}; -``` - -### 4. Wrap Data Extraction in Try-Catch - -```typescript -start: (event) => { - try { - const { input, metadata } = extractInput(event.arguments); - span.log({ input, metadata }); - } catch (error) { - console.error("Error extracting input:", error); - // Continue - don't let extraction errors break instrumentation - } -}; -``` - -### 5. Delete Span Data After Use - -```typescript -asyncEnd: (event) => { - const spanData = spans.get(event); - if (!spanData) return; - - span.log({ output: event.result }); - span.end(); - - // ALWAYS delete to prevent memory growth - spans.delete(event); -}; -``` - -### 6. Check for Stream Support - -```typescript -import { isAsyncIterable } from "./core"; - -asyncEnd: (event) => { - if (isAsyncIterable(event.result)) { - // Handle streaming - patchStreamIfNeeded(event.result, { ... }); - } else { - // Handle non-streaming - span.log({ output: event.result }); - span.end(); - } -} -``` - -### 7. Use Type Guards - -```typescript -interface OpenAIResponse { - choices: any[]; - usage?: { - total_tokens: number; - prompt_tokens: number; - completion_tokens: number; - }; -} - -function isOpenAIResponse(value: unknown): value is OpenAIResponse { - return ( - typeof value === "object" && - value !== null && - "choices" in value && - Array.isArray(value.choices) - ); -} - -asyncEnd: (event) => { - if (isOpenAIResponse(event.result)) { - // TypeScript knows event.result is OpenAIResponse - span.log({ output: event.result.choices }); - } -}; -``` - ---- +- Treat arguments, results, metadata, and headers as untrusted. +- Avoid prototype-sensitive merges and unnecessary mutation of provider data. +- Capture only fields permitted by the instrumentation specification. +- Pass `Error` objects directly to `span.log({ error })`. +- Use narrow vendored provider interfaces shared by wrappers and plugins. +- Keep enable, disable, subscription, and patching behavior idempotent. ## Testing -### Unit Testing Event Handlers - -```typescript -import { describe, it, expect, vi } from "vitest"; -import { tracingChannel } from "dc-browser"; - -describe("OpenAI Plugin", () => { - it("logs input on start", () => { - const channel = tracingChannel( - "orchestrion:openai:chat.completions.create", - ); - const logSpy = vi.fn(); - - channel.subscribe({ - start: (event) => { - const params = event.arguments[0]; - logSpy(params.messages); - }, - }); - - // Simulate orchestrion publishing an event - const mockEvent = { - arguments: [{ messages: [{ role: "user", content: "Hi" }] }], - self: {}, - moduleVersion: "5.0.0", - }; - - // Trigger the event - channel.traceSync(() => {}, mockEvent); - - expect(logSpy).toHaveBeenCalledWith([{ role: "user", content: "Hi" }]); - }); -}); -``` - -### Testing Stream Patching - -```typescript -import { patchStreamIfNeeded } from "./core"; - -describe("Stream patching", () => { - it("collects chunks from async iterator", async () => { - async function* mockStream() { - yield { delta: "Hello" }; - yield { delta: " " }; - yield { delta: "world" }; - } - - const stream = mockStream(); - const collected: any[] = []; - - patchStreamIfNeeded(stream, { - onComplete: (chunks) => { - collected.push(...chunks); - }, - }); - - // Consume the stream - const result: any[] = []; - for await (const chunk of stream) { - result.push(chunk); - } - - // Verify both user gets chunks AND we collected them - expect(result).toHaveLength(3); - expect(collected).toHaveLength(3); - expect(collected[0].delta).toBe("Hello"); - }); - - it("handles stream errors", async () => { - async function* errorStream() { - yield { data: "chunk1" }; - throw new Error("Stream error"); - } - - const stream = errorStream(); - let errorCaught = false; - let chunksBeforeError: any[] = []; - - patchStreamIfNeeded(stream, { - onComplete: () => {}, - onError: (error, chunks) => { - errorCaught = true; - chunksBeforeError = chunks; - }, - }); - - await expect(async () => { - for await (const chunk of stream) { - // consume - } - }).rejects.toThrow("Stream error"); - - expect(errorCaught).toBe(true); - expect(chunksBeforeError).toHaveLength(1); - }); - - it("handles early cancellation", async () => { - async function* longStream() { - for (let i = 0; i < 100; i++) { - yield { i }; - } - } - - const stream = longStream(); - let onCompleteCalled = false; - const collected: any[] = []; - - patchStreamIfNeeded(stream, { - onComplete: (chunks) => { - onCompleteCalled = true; - collected.push(...chunks); - }, - }); - - // Consume only 5 chunks then break - let count = 0; - for await (const chunk of stream) { - count++; - if (count === 5) break; - } - - // onComplete should still be called with partial chunks - expect(onCompleteCalled).toBe(true); - expect(collected).toHaveLength(5); - }); -}); -``` - -### Integration Testing with Real SDKs - -```typescript -import { describe, it, expect } from "vitest"; -import OpenAI from "openai"; -import { OpenAIPlugin } from "./openai-plugin"; -import { initLogger, currentSpan } from "../logger"; - -describe("OpenAI Plugin Integration", () => { - it("logs streaming chat completion", async () => { - // Set up plugin - const plugin = new OpenAIPlugin(); - plugin.enable(); - - // Initialize logger - initLogger({ projectName: "test" }); - - // Make actual API call - const client = new OpenAI({ - apiKey: process.env.OPENAI_API_KEY, - }); - - const stream = await client.chat.completions.create({ - model: "gpt-4", - messages: [{ role: "user", content: "Say hello" }], - stream: true, - }); - - // Consume stream - const chunks = []; - for await (const chunk of stream) { - chunks.push(chunk); - } - - // Verify span was logged - const span = currentSpan(); - expect(span).toBeDefined(); - // Add more assertions about span data - - plugin.disable(); - }); -}); -``` - ---- - -## Advanced Topics - -### Custom Event Context - -Add custom context to events: - -```typescript -start: (event) => { - const span = startSpan({ name: "Operation" }); - - // Store additional context - spans.set(event, { - span, - startTime: Date.now(), - requestId: generateRequestId(), - userId: getCurrentUserId(), - }); -}; -``` - -### Conditional Instrumentation - -Only instrument certain calls: - -```typescript -start: (event) => { - const params = event.arguments[0]; - - // Skip instrumentation for specific models - if (params.model === "gpt-3.5-turbo") { - return; // Don't create span - } - - const span = startSpan({ name: "Chat Completion" }); - spans.set(event, { span }); -}; -``` - -### Multiple Channel Subscription - -Subscribe to multiple channels with shared logic: - -```typescript -class MultiProviderPlugin extends BasePlugin { - protected onEnable(): void { - const providers = [ - { channel: "orchestrion:openai:chat.completions.create", name: "OpenAI" }, - { channel: "orchestrion:anthropic:messages.create", name: "Anthropic" }, - { channel: "orchestrion:vercel-ai:generateText", name: "Vercel" }, - ]; - - for (const { channel, name } of providers) { - this.subscribeToProvider(channel, name); - } - } - - private subscribeToProvider(channelName: string, providerName: string): void { - const channel = tracingChannel(channelName); - // ... shared subscription logic - } -} -``` - ---- - -## Troubleshooting - -### Events Not Firing - -**Problem:** Your subscriber isn't being called. - -**Solutions:** - -1. Verify orchestrion config is correct -2. Check channel name matches exactly -3. Ensure orchestrion transformation is running (check bundler plugin) -4. Verify function is actually being called -5. Check if `hasSubscribers` optimization is preventing instrumentation - -```typescript -// Debug: Force channel to always instrument -const channel = tracingChannel("orchestrion:openai:chat.completions.create"); -channel.hasSubscribers = true; // Force instrumentation -``` - -### Stream Not Being Patched - -**Problem:** Streaming output not collected. - -**Solutions:** - -1. Verify `event.result` is actually an async iterable -2. Check if stream is frozen/sealed -3. Ensure patching happens in `asyncEnd`, not `asyncStart` -4. Check browser console for warnings - -```typescript -import { isAsyncIterable } from "./core"; - -asyncEnd: (event) => { - console.log("Is stream?", isAsyncIterable(event.result)); - console.log("Is frozen?", Object.isFrozen(event.result)); -}; -``` - -### Memory Leaks - -**Problem:** Memory usage growing over time. - -**Solutions:** - -1. Always delete from WeakMap after span ends -2. Unsubscribe channels in `onDisable` -3. Use WeakMap, not Map -4. Check for dangling references - -```typescript -// ✅ Good -asyncEnd: (event) => { - span.end(); - spans.delete(event); // Always delete -}; - -// ❌ Bad -asyncEnd: (event) => { - span.end(); - // Forgot to delete - memory leak! -}; -``` - ---- - -## Summary - -**Key Takeaways:** - -1. **Orchestrion** instruments code at build/load time to publish diagnostics channel events -2. **Diagnostics Channels** provide asyncStart/asyncEnd/error events for instrumented functions -3. **BasePlugin** provides lifecycle management for subscribing to channels -4. **WeakMap** correlates events across asyncStart/asyncEnd/error -5. **Stream Patching** enables collecting streaming outputs by mutating the stream in-place -6. **wrapStreamResult** provides high-level API for handling both streaming and non-streaming results - -**Next Steps:** - -- Review the `BraintrustPlugin` implementation for a complete example -- Read `core/stream-patcher.ts` for stream patching implementation details -- Check orchestrion documentation for creating instrumentation configs -- Write tests for your plugin before deploying - ---- +Test at the narrowest useful layers: -## Additional Resources +1. Plugin unit tests for extraction and span handling. +2. Global hook/runtime tests for lifecycle and context behavior. +3. Orchestrion transformation tests for generated wrappers. +4. Bundler and loader tests for real transformed execution. +5. Provider e2e tests for wrapped and auto-hook parity. -- [Orchestrion Documentation](https://github.com/nodejs/orchestrion-js) -- [Node.js Diagnostics Channel API](https://nodejs.org/api/diagnostics_channel.html) -- [BraintrustPlugin Implementation](./braintrust-plugin.ts) -- [Stream Patcher Source](./core/stream-patcher.ts) +After instrumentation changes, run e2e tests in cassette replay mode. When an +e2e scenario itself changes, run it three times to catch flakes. diff --git a/js/src/instrumentation/braintrust-plugin.ts b/js/src/instrumentation/braintrust-plugin.ts index d5435f2b0..3851881d4 100644 --- a/js/src/instrumentation/braintrust-plugin.ts +++ b/js/src/instrumentation/braintrust-plugin.ts @@ -213,7 +213,7 @@ export class BraintrustPlugin extends BasePlugin { // ObservabilityExporter contract, and `BraintrustObservabilityExporter` // (wrappers/mastra.ts) is auto-installed by the loader patch in // `auto-instrumentations/loader/mastra-observability-patch.ts` rather than - // by a BasePlugin / tracingChannel subscription. + // by a BasePlugin / global hook subscription. } protected onDisable(): void { diff --git a/js/src/instrumentation/core/channel.ts b/js/src/instrumentation/core/channel.ts index d33a5dcc1..8cbf253d6 100644 --- a/js/src/instrumentation/core/channel.ts +++ b/js/src/instrumentation/core/channel.ts @@ -1,5 +1,5 @@ /** - * Utilities for diagnostics_channel naming and management. + * Utilities for instrumentation hook naming and management. */ /** diff --git a/js/src/instrumentation/core/plugin.ts b/js/src/instrumentation/core/plugin.ts index 41e1d7c3d..ae09a43f1 100644 --- a/js/src/instrumentation/core/plugin.ts +++ b/js/src/instrumentation/core/plugin.ts @@ -13,7 +13,7 @@ import { /** * Base class for creating instrumentation plugins. * - * Plugins subscribe to diagnostics_channel events and convert them + * Plugins subscribe to global instrumentation hook events and convert them * into spans, logs, or other observability data. */ export abstract class BasePlugin { diff --git a/js/src/instrumentation/core/stream-patcher.ts b/js/src/instrumentation/core/stream-patcher.ts index 8bb25d80b..13807d5ca 100644 --- a/js/src/instrumentation/core/stream-patcher.ts +++ b/js/src/instrumentation/core/stream-patcher.ts @@ -2,7 +2,7 @@ * Utilities for patching async iterables (streams) to collect chunks * without modifying the user-facing behavior. * - * This allows diagnostics channel subscribers to collect streaming outputs + * This allows global hook subscribers to collect streaming outputs * even though they cannot replace return values. */ diff --git a/js/src/instrumentation/core/types.ts b/js/src/instrumentation/core/types.ts index 7c34a0f74..dc2434dc3 100644 --- a/js/src/instrumentation/core/types.ts +++ b/js/src/instrumentation/core/types.ts @@ -1,7 +1,7 @@ /** - * Standard event types for diagnostics_channel-based instrumentation. + * Standard event types for global hook-based instrumentation. * - * These types follow the TracingChannel pattern from Node.js. + * These types retain the TracingChannel-compatible lifecycle. * For async functions (tracePromise): * - start: Called before the synchronous portion executes * - end: Called after the synchronous portion completes (promise returned) @@ -146,7 +146,7 @@ export type ErrorEventWith< > = TypedErrorEvent & TExtra; /** - * Subscription handlers for a TracingChannel. + * Subscription handlers for a tracing-compatible global hook. * * Common usage pattern: * - Use start to create spans and extract input diff --git a/js/src/instrumentation/index.ts b/js/src/instrumentation/index.ts index f398abbdb..338d62abd 100644 --- a/js/src/instrumentation/index.ts +++ b/js/src/instrumentation/index.ts @@ -1,8 +1,8 @@ /** * Instrumentation APIs for auto-instrumentation. * - * This module provides the core plugin infrastructure for converting - * diagnostics_channel events into Braintrust spans. + * This module provides the core plugin infrastructure for converting global + * instrumentation hook events into Braintrust spans. * * Following the OpenTelemetry pattern, BasePlugin (like InstrumentationBase) * lives in the core SDK, while individual instrumentation implementations diff --git a/js/src/instrumentation/plugins/genkit-plugin.test.ts b/js/src/instrumentation/plugins/genkit-plugin.test.ts index ee35f34c8..dcf41c0f7 100644 --- a/js/src/instrumentation/plugins/genkit-plugin.test.ts +++ b/js/src/instrumentation/plugins/genkit-plugin.test.ts @@ -1,4 +1,5 @@ -import { afterEach, describe, expect, it } from "vitest"; +import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; +import { _exportsForTestingOnly, initLogger } from "../../logger"; import { GenkitPlugin } from "./genkit-plugin"; import { genkitChannels } from "./genkit-channels"; @@ -37,8 +38,21 @@ async function collectAsync(stream: AsyncIterable): Promise { describe("GenkitPlugin stream patching", () => { const plugin = new GenkitPlugin(); + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ + projectName: "genkit-plugin.test.ts", + projectId: "test-project-id", + }); + }); + afterEach(() => { plugin.disable(); + _exportsForTestingOnly.clearTestBackgroundLogger(); }); it("does not consume generateStream chunks before user code reads them", async () => { diff --git a/js/src/instrumentation/plugins/langchain-plugin.test.ts b/js/src/instrumentation/plugins/langchain-plugin.test.ts index a098977fd..29ace20ab 100644 --- a/js/src/instrumentation/plugins/langchain-plugin.test.ts +++ b/js/src/instrumentation/plugins/langchain-plugin.test.ts @@ -1,14 +1,7 @@ -import * as diagnosticsChannel from "node:diagnostics_channel"; import { describe, expect, it } from "vitest"; -import iso from "../../isomorph"; import { LangChainPlugin } from "./langchain-plugin"; import { langChainChannels } from "./langchain-channels"; -iso.newTracingChannel = (nameOrChannels: string | object) => - diagnosticsChannel.tracingChannel( - nameOrChannels as string, - ) as never as ReturnType>; - function createManager(handlers: unknown[] = []) { return { handlers, diff --git a/js/src/instrumentation/registry.test.ts b/js/src/instrumentation/registry.test.ts index 4fa54be93..a0c41caa0 100644 --- a/js/src/instrumentation/registry.test.ts +++ b/js/src/instrumentation/registry.test.ts @@ -59,7 +59,7 @@ describe("Plugin Registry", () => { // Regression test for BT-5139: when the SDK is loaded from two different // module paths in the same process, each gets its own PluginRegistry // instance. Without cross-instance deduplication, both would subscribe to - // the same diagnostics_channel, causing every OpenAI call to produce two + // the same global hook, causing every OpenAI call to produce two // LLM spans. // // The dedup mechanism checks globalThis[Symbol.for("braintrust-state")], @@ -76,7 +76,7 @@ describe("Plugin Registry", () => { instanceA.enable(); expect(instanceA.isEnabled()).toBe(true); - // instanceB should be blocked — the channel is already subscribed + // instanceB should be blocked — the hook is already subscribed instanceB.enable(); expect(instanceB.isEnabled()).toBe(false); } finally { diff --git a/js/src/instrumentation/registry.ts b/js/src/instrumentation/registry.ts index a28f7964a..6ded5bd86 100644 --- a/js/src/instrumentation/registry.ts +++ b/js/src/instrumentation/registry.ts @@ -24,7 +24,7 @@ export type { InstrumentationConfig } from "./config"; // // - BT-5139 scenario: two SDK instances share the same state object → the // second instance sees the marker left by the first and skips subscription, -// preventing duplicate diagnostics_channel listeners. +// preventing duplicate global hook listeners. // // - vi.resetModules() test scenario: the test deletes the state from // globalThis between runs, so the next import creates a fresh state with no @@ -71,7 +71,7 @@ class PluginRegistry { } // If another SDK instance in the same process already registered plugins, - // skip to avoid duplicate diagnostics_channel subscriptions. + // skip to avoid duplicate global hook subscriptions. const sharedState = getSharedState(); if (sharedState) { if (sharedState[REGISTRY_STATE_KEY] !== undefined) { diff --git a/js/src/isomorph.ts b/js/src/isomorph.ts index 05e27d066..9837dd23e 100644 --- a/js/src/isomorph.ts +++ b/js/src/isomorph.ts @@ -2,6 +2,14 @@ import { type GitMetadataSettingsType as GitMetadataSettings, type RepoInfoType as RepoInfo, } from "./generated_types"; +import { + newGlobalTracingChannel, + type GlobalHookAsyncLocalStorage, + type GlobalHookChannel, + type GlobalHookHandlers, + type GlobalTracingChannel, + type GlobalTracingChannelCollection, +} from "./global-instrumentation-hooks"; export interface CallerLocation { caller_functionname: string; @@ -9,43 +17,11 @@ export interface CallerLocation { caller_lineno: number; } -export interface IsoAsyncLocalStorage { - enterWith(store: T): void; - run(store: T | undefined, callback: () => R): R; - getStore(): T | undefined; -} - -// eslint-disable-next-line @typescript-eslint/no-explicit-any -type IsoMessageFunction = ( - message: M, - name: N, -) => void; - -type IsoTransformFunction = (message: M) => S; - -/** - * Channel interface matching the shared node:diagnostics_channel and dc-browser API. - */ -// eslint-disable-next-line @typescript-eslint/no-explicit-any -export interface IsoChannel { - readonly name: N; - readonly hasSubscribers: boolean; - subscribe(subscription: IsoMessageFunction): void; - unsubscribe(subscription: IsoMessageFunction): boolean; - bindStore( - store: IsoAsyncLocalStorage, - transform?: IsoTransformFunction, - ): void; - unbindStore(store: IsoAsyncLocalStorage): boolean; - publish(message: M): void; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - runStores any>( - message: M, - fn: F, - thisArg?: ThisParameterType, - ...args: Parameters - ): ReturnType; -} +export type IsoAsyncLocalStorage = GlobalHookAsyncLocalStorage; +export type IsoChannel< + M = any, + N extends string | symbol = string, +> = GlobalHookChannel; class DefaultAsyncLocalStorage implements IsoAsyncLocalStorage { constructor() {} @@ -59,170 +35,10 @@ class DefaultAsyncLocalStorage implements IsoAsyncLocalStorage { } } -class DefaultChannel< - M, - N extends string | symbol = string, -> implements IsoChannel { - readonly hasSubscribers = false; - - constructor(public readonly name: N) {} - - subscribe(_subscription: IsoMessageFunction): void {} - - unsubscribe(_subscription: IsoMessageFunction): boolean { - return false; - } - - bindStore( - _store: IsoAsyncLocalStorage, - _transform?: IsoTransformFunction, - ): void {} - - unbindStore(_store: IsoAsyncLocalStorage): boolean { - return false; - } - - publish(_message: M): void {} - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - runStores any>( - _message: M, - fn: F, - thisArg?: ThisParameterType, - ...args: Parameters - ): ReturnType { - return fn.apply(thisArg, args); - } -} - -/** - * TracingChannel interface matching both node:diagnostics_channel and dc-browser. - * A composite of the five tracing subchannels used to instrument sync/async operations. - */ -// eslint-disable-next-line @typescript-eslint/no-explicit-any -export interface IsoTracingChannelCollection { - readonly start?: IsoChannel; - readonly end?: IsoChannel; - readonly asyncStart?: IsoChannel; - readonly asyncEnd?: IsoChannel; - readonly error?: IsoChannel; -} - -// eslint-disable-next-line @typescript-eslint/no-explicit-any -export interface IsoTracingChannel< - M = any, -> extends IsoTracingChannelCollection { - readonly hasSubscribers: boolean; - subscribe(handlers: IsoChannelHandlers): void; - unsubscribe(handlers: IsoChannelHandlers): boolean; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - traceSync any>( - fn: F, - message?: M, - thisArg?: ThisParameterType, - ...args: Parameters - ): ReturnType; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - tracePromise PromiseLike>( - fn: F, - message?: M, - thisArg?: ThisParameterType, - ...args: Parameters - ): ReturnType; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - traceCallback any>( - fn: F, - position?: number, - message?: M, - thisArg?: ThisParameterType, - ...args: Parameters - ): ReturnType; -} - -// eslint-disable-next-line @typescript-eslint/no-explicit-any -export interface IsoChannelHandlers { - start?: (context: M, name: string) => void; - end?: (context: M, name: string) => void; - asyncStart?: (context: M, name: string) => void; - asyncEnd?: (context: M, name: string) => void; - error?: (context: M, name: string) => void; -} - -/** - * Default no-op TracingChannel implementation. - */ -class DefaultTracingChannel implements IsoTracingChannel { - readonly start: IsoChannel; - readonly end: IsoChannel; - readonly asyncStart: IsoChannel; - readonly asyncEnd: IsoChannel; - readonly error: IsoChannel; - - constructor(nameOrChannels: string | IsoTracingChannelCollection) { - if (typeof nameOrChannels === "string") { - this.start = new DefaultChannel(`tracing:${nameOrChannels}:start`); - this.end = new DefaultChannel(`tracing:${nameOrChannels}:end`); - this.asyncStart = new DefaultChannel( - `tracing:${nameOrChannels}:asyncStart`, - ); - this.asyncEnd = new DefaultChannel(`tracing:${nameOrChannels}:asyncEnd`); - this.error = new DefaultChannel(`tracing:${nameOrChannels}:error`); - return; - } - - this.start = nameOrChannels.start ?? new DefaultChannel("tracing:start"); - this.end = nameOrChannels.end ?? new DefaultChannel("tracing:end"); - this.asyncStart = - nameOrChannels.asyncStart ?? new DefaultChannel("tracing:asyncStart"); - this.asyncEnd = - nameOrChannels.asyncEnd ?? new DefaultChannel("tracing:asyncEnd"); - this.error = nameOrChannels.error ?? new DefaultChannel("tracing:error"); - } - - get hasSubscribers(): boolean { - return ( - this.start.hasSubscribers || - this.end.hasSubscribers || - this.asyncStart.hasSubscribers || - this.asyncEnd.hasSubscribers || - this.error.hasSubscribers - ); - } - - subscribe(_handlers: IsoChannelHandlers): void {} - unsubscribe(_handlers: IsoChannelHandlers): boolean { - return false; - } - // eslint-disable-next-line @typescript-eslint/no-explicit-any - traceSync any>( - fn: F, - _message?: M, - thisArg?: ThisParameterType, - ...args: Parameters - ): ReturnType { - return fn.apply(thisArg, args); - } - // eslint-disable-next-line @typescript-eslint/no-explicit-any - tracePromise PromiseLike>( - fn: F, - _message?: M, - thisArg?: ThisParameterType, - ...args: Parameters - ): ReturnType { - // eslint-disable-next-line @typescript-eslint/consistent-type-assertions, @typescript-eslint/no-explicit-any - return fn.apply(thisArg, args) as any; - } - // eslint-disable-next-line @typescript-eslint/no-explicit-any - traceCallback any>( - fn: F, - _position?: number, - _message?: M, - thisArg?: ThisParameterType, - ...args: Parameters - ): ReturnType { - return fn.apply(thisArg, args); - } -} +export type IsoTracingChannelCollection = + GlobalTracingChannelCollection; +export type IsoTracingChannel = GlobalTracingChannel; +export type IsoChannelHandlers = GlobalHookHandlers; interface Common { buildType: @@ -304,7 +120,7 @@ const iso: Common = { // eslint-disable-next-line @typescript-eslint/no-explicit-any newTracingChannel: ( nameOrChannels: string | IsoTracingChannelCollection, - ) => new DefaultTracingChannel(nameOrChannels), + ) => newGlobalTracingChannel(nameOrChannels), processOn: (_0, _1) => {}, basename: (filepath: string) => filepath.split(/[\\/]/).pop() || filepath, // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. diff --git a/js/src/logger.ts b/js/src/logger.ts index fab3b3ead..633010ade 100644 --- a/js/src/logger.ts +++ b/js/src/logger.ts @@ -501,7 +501,7 @@ export const BRAINTRUST_CURRENT_SPAN_STORE = Symbol.for( * - Default (`BraintrustContextManager`): stores a `Span` * - OTEL compat (`OtelContextManager`): stores an OTEL `Context` object * - * TracingChannel's `bindStore` transform (via `wrapSpanForStore`) produces the + * The global hook's `bindStore` transform (via `wrapSpanForStore`) produces the * correct value type for whichever mode is active. */ export type CurrentSpanStore = IsoAsyncLocalStorage; @@ -512,7 +512,7 @@ export abstract class ContextManager { abstract getCurrentSpan(): Span | undefined; /** - * Returns the value to store in the ALS bound to a TracingChannel's start event. + * Returns the value to store in the ALS bound to a global hook's start event. * In default mode this is the Span itself; in OTEL mode it is the OTEL Context * containing the span so that OTEL's own ALS stores a proper Context object. */ diff --git a/js/src/node/apply-auto-instrumentation-entry.ts b/js/src/node/apply-auto-instrumentation-entry.ts index 0328688dd..11b995bf8 100644 --- a/js/src/node/apply-auto-instrumentation-entry.ts +++ b/js/src/node/apply-auto-instrumentation-entry.ts @@ -1,9 +1,7 @@ -import * as diagnostics_channel from "node:diagnostics_channel"; import { register } from "node:module"; import { pathToFileURL } from "node:url"; import { getDefaultAutoInstrumentationConfigs } from "../auto-instrumentations/configs/all"; import { ModulePatch } from "../auto-instrumentations/loader/cjs-patch"; -import { patchTracingChannel } from "../auto-instrumentations/patch-tracing-channel"; interface ApplyAutoInstrumentationState { applied?: boolean; @@ -30,8 +28,6 @@ if (state !== existingState) { } if (!state.applied) { - patchTracingChannel(diagnostics_channel.tracingChannel); - const allConfigs = getDefaultAutoInstrumentationConfigs(); const currentModuleUrl = getCurrentModuleUrl(); diff --git a/js/src/node/config.ts b/js/src/node/config.ts index f9743837a..6b3f37802 100644 --- a/js/src/node/config.ts +++ b/js/src/node/config.ts @@ -1,7 +1,5 @@ import { AsyncLocalStorage } from "node:async_hooks"; -import * as diagnostics_channel from "node:diagnostics_channel"; import * as path from "node:path"; -import { patchTracingChannel } from "../auto-instrumentations/patch-tracing-channel"; import * as fs from "node:fs/promises"; import * as os from "node:os"; import * as fsSync from "node:fs"; @@ -111,12 +109,6 @@ export function configureNode() { }; iso.getCallerLocation = getCallerLocation; iso.newAsyncLocalStorage = () => new AsyncLocalStorage(); - iso.newTracingChannel = <_M = any>(nameOrChannels: string | object) => - (diagnostics_channel as any).tracingChannel(nameOrChannels) as any; - - // Patch TracingChannel.prototype.tracePromise to handle APIPromise and other - // Promise subclasses (mirrors the fix in hook.mts for the --import loader path). - patchTracingChannel((diagnostics_channel as any).tracingChannel); iso.processOn = (event: string, handler: (code: unknown) => void) => { process.on(event, handler); }; diff --git a/js/src/workerd/config.ts b/js/src/workerd/config.ts index cdf516d96..bdd813520 100644 --- a/js/src/workerd/config.ts +++ b/js/src/workerd/config.ts @@ -1,12 +1,6 @@ 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; @@ -27,13 +21,6 @@ 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; diff --git a/js/tests/auto-instrumentations/error-handling.test.ts b/js/tests/auto-instrumentations/error-handling.test.ts index 62f4430e8..2c31588a4 100644 --- a/js/tests/auto-instrumentations/error-handling.test.ts +++ b/js/tests/auto-instrumentations/error-handling.test.ts @@ -7,7 +7,7 @@ * - Error event contains correct error information * - Both sync and async errors are handled * - * Tests run against bundled code with actual diagnostics_channel subscriptions. + * Tests run against bundled code with actual global hook subscriptions. */ import { describe, it, expect, beforeAll, afterAll, beforeEach } from "vitest"; @@ -33,8 +33,6 @@ describe("Error Handling", () => { if (!fs.existsSync(outputDir)) { fs.mkdirSync(outputDir, { recursive: true }); } - - // Note: dc-browser is now an npm package, no symlinks needed }); beforeEach(() => { diff --git a/js/tests/auto-instrumentations/event-content.test.ts b/js/tests/auto-instrumentations/event-content.test.ts index 40e426a97..5a2fdd100 100644 --- a/js/tests/auto-instrumentations/event-content.test.ts +++ b/js/tests/auto-instrumentations/event-content.test.ts @@ -1,14 +1,14 @@ /** * EVENT CONTENT VALIDATION TESTS * - * These tests verify that diagnostics_channel events contain the correct information: + * These tests verify that global hook events contain the correct information: * - Start events contain correct arguments * - End events contain correct results * - `self` context is captured correctly * - moduleVersion is populated (if available) * - Event timing and ordering * - * Tests run against bundled code with actual diagnostics_channel subscriptions. + * Tests run against bundled code with actual global hook subscriptions. */ import { describe, it, expect, beforeAll, afterAll, beforeEach } from "vitest"; @@ -35,8 +35,6 @@ describe("Event Content Validation", () => { if (!fs.existsSync(outputDir)) { fs.mkdirSync(outputDir, { recursive: true }); } - - // Note: dc-browser is now an npm package, no symlinks needed }); beforeEach(() => { diff --git a/js/tests/auto-instrumentations/fixtures/global-hook-listener.cjs b/js/tests/auto-instrumentations/fixtures/global-hook-listener.cjs new file mode 100644 index 000000000..98b11e0c3 --- /dev/null +++ b/js/tests/auto-instrumentations/fixtures/global-hook-listener.cjs @@ -0,0 +1,176 @@ +const hooks = + globalThis.__braintrust_instrumentation_hooks ?? + (() => { + const registry = new Map(); + Object.defineProperty(globalThis, "__braintrust_instrumentation_hooks", { + configurable: false, + enumerable: false, + value: registry, + writable: false, + }); + return registry; + })(); + +function phase(name) { + const subscribers = []; + const stores = new Map(); + return { + name, + get hasSubscribers() { + return subscribers.length > 0 || stores.size > 0; + }, + bindStore(store, transform) { + stores.set(store, transform); + }, + unbindStore(store) { + return stores.delete(store); + }, + publish(message) { + for (const subscriber of subscribers) subscriber(message, name); + }, + runStores(message, fn) { + let run = () => { + this.publish(message); + return fn(); + }; + for (const [store, transform] of stores) { + const next = run; + run = () => store.run(transform ? transform(message) : message, next); + } + return run(); + }, + subscribe(subscriber) { + subscribers.push(subscriber); + }, + unsubscribe(subscriber) { + const index = subscribers.indexOf(subscriber); + if (index === -1) return false; + subscribers.splice(index, 1); + return true; + }, + }; +} + +function getTracingHook(channelName) { + let hook = hooks.get(channelName); + if (hook) return hook; + + hook = { + start: phase(`tracing:${channelName}:start`), + end: phase(`tracing:${channelName}:end`), + asyncStart: phase(`tracing:${channelName}:asyncStart`), + asyncEnd: phase(`tracing:${channelName}:asyncEnd`), + error: phase(`tracing:${channelName}:error`), + get hasSubscribers() { + return ( + this.start.hasSubscribers || + this.end.hasSubscribers || + this.asyncStart.hasSubscribers || + this.asyncEnd.hasSubscribers || + this.error.hasSubscribers + ); + }, + subscribe(handlers) { + for (const name of ["start", "end", "asyncStart", "asyncEnd", "error"]) { + if (handlers[name]) this[name].subscribe(handlers[name]); + } + }, + unsubscribe(handlers) { + let done = true; + for (const name of ["start", "end", "asyncStart", "asyncEnd", "error"]) { + if (handlers[name] && !this[name].unsubscribe(handlers[name])) { + done = false; + } + } + return done; + }, + traceSync(fn, message) { + return this.start.runStores(message, () => { + try { + message.result = fn(); + return message.result; + } catch (error) { + message.error = error; + this.error.publish(message); + throw error; + } finally { + this.end.publish(message); + } + }); + }, + tracePromise(fn, message) { + return this.start.runStores(message, () => { + try { + const result = fn(); + if (typeof result?.then !== "function") { + message.result = result; + this.end.publish(message); + return result; + } + this.end.publish(message); + return result.then( + (value) => { + message.result = value; + this.asyncStart.publish(message); + this.asyncEnd.publish(message); + return value; + }, + (error) => { + message.error = error; + this.error.publish(message); + this.asyncStart.publish(message); + this.asyncEnd.publish(message); + throw error; + }, + ); + } catch (error) { + message.error = error; + this.error.publish(message); + this.end.publish(message); + throw error; + } + }); + }, + traceCallback(fn, position, message) { + const callback = Array.prototype.at.call(message.arguments, position); + if (typeof callback !== "function") return fn(); + const currentHook = this; + function wrappedCallback(error, result) { + if (error) { + message.error = error; + currentHook.error.publish(message); + } else { + message.result = result; + } + return currentHook.asyncStart.runStores(message, () => { + try { + return callback.apply(this, arguments); + } finally { + currentHook.asyncEnd.publish(message); + } + }); + } + Array.prototype.splice.call( + message.arguments, + position, + 1, + wrappedCallback, + ); + return this.start.runStores(message, () => { + try { + return fn(); + } catch (error) { + message.error = error; + this.error.publish(message); + throw error; + } finally { + this.end.publish(message); + } + }); + }, + }; + hooks.set(channelName, hook); + return hook; +} + +module.exports = { getTracingHook }; diff --git a/js/tests/auto-instrumentations/fixtures/listener-cjs.cjs b/js/tests/auto-instrumentations/fixtures/listener-cjs.cjs index e2c42e693..825c04309 100644 --- a/js/tests/auto-instrumentations/fixtures/listener-cjs.cjs +++ b/js/tests/auto-instrumentations/fixtures/listener-cjs.cjs @@ -1,12 +1,12 @@ -const { tracingChannel } = require("diagnostics_channel"); const { parentPort } = require("worker_threads"); +const { getTracingHook } = require("./global-hook-listener.cjs"); const events = { start: [], end: [], error: [] }; // NOTE: code-transformer prepends "orchestrion:openai:" to the channel name const expectedChannel = "orchestrion:openai:chat.completions.create"; -// Subscribe to the channel and accumulate events -const channel = tracingChannel(expectedChannel); +// Subscribe to the global hook and accumulate events +const channel = getTracingHook(expectedChannel); channel.subscribe({ start: (ctx) => { // Convert arguments to array for serialization diff --git a/js/tests/auto-instrumentations/fixtures/listener-esm.mjs b/js/tests/auto-instrumentations/fixtures/listener-esm.mjs index 4ab2f726a..c735b77fe 100644 --- a/js/tests/auto-instrumentations/fixtures/listener-esm.mjs +++ b/js/tests/auto-instrumentations/fixtures/listener-esm.mjs @@ -1,23 +1,21 @@ -import { tracingChannel } from "node:diagnostics_channel"; +import { createRequire } from "node:module"; import { parentPort } from "node:worker_threads"; -import * as dc from "node:diagnostics_channel"; + +const { getTracingHook } = createRequire(import.meta.url)( + "./global-hook-listener.cjs", +); const events = { start: [], end: [], error: [] }; // NOTE: code-transformer prepends "orchestrion:openai:" to the channel name const expectedChannel = "orchestrion:openai:chat.completions.create"; -// Get the kStoreKey symbol to access the store -const kStoreKey = dc.kStoreKey || Symbol.for("diagnostics_channel.store"); - -// Subscribe to the channel and accumulate events -const channel = tracingChannel(expectedChannel); +// Subscribe to the global hook and accumulate events +const channel = getTracingHook(expectedChannel); channel.subscribe({ start: (ctx) => { - // Arguments are stored in the Symbol(diagnostics_channel.store) - const store = ctx[kStoreKey]; events.start.push({ - args: store?.arguments ? Array.from(store.arguments) : [], - self: !!store?.self, + args: ctx.arguments ? Array.from(ctx.arguments) : [], + self: !!ctx.self, }); }, asyncEnd: (ctx) => { diff --git a/js/tests/auto-instrumentations/fixtures/orchestrion-js/arguments_mutation/test.js b/js/tests/auto-instrumentations/fixtures/orchestrion-js/arguments_mutation/test.js index 6406b99a5..9a7917da9 100644 --- a/js/tests/auto-instrumentations/fixtures/orchestrion-js/arguments_mutation/test.js +++ b/js/tests/auto-instrumentations/fixtures/orchestrion-js/arguments_mutation/test.js @@ -3,8 +3,7 @@ * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. **/ const { fetch_simple, fetch_complex } = require("./instrumented.js"); -const assert = require("node:assert"); -const { tracingChannel } = require("node:diagnostics_channel"); +const { assert, getTracingHook } = require("../common/preamble.js"); const handler = { start(message) { @@ -21,8 +20,8 @@ const handler = { }, }; -tracingChannel("orchestrion:undici:fetch_simple").subscribe(handler); -tracingChannel("orchestrion:undici:fetch.complex").subscribe(handler); +getTracingHook("orchestrion:undici:fetch_simple").subscribe(handler); +getTracingHook("orchestrion:undici:fetch.complex").subscribe(handler); assert.strictEqual(fetch_simple.length, 2); assert.strictEqual(fetch_complex.length, 2); diff --git a/js/tests/auto-instrumentations/fixtures/orchestrion-js/common/preamble.js b/js/tests/auto-instrumentations/fixtures/orchestrion-js/common/preamble.js index 3ec3b4abb..d232f0aae 100644 --- a/js/tests/auto-instrumentations/fixtures/orchestrion-js/common/preamble.js +++ b/js/tests/auto-instrumentations/fixtures/orchestrion-js/common/preamble.js @@ -2,10 +2,157 @@ * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. **/ -const { tracingChannel } = require("node:diagnostics_channel"); const assert = require("node:assert"); + +const hooks = new Map(); +Object.defineProperty(globalThis, "__braintrust_instrumentation_hooks", { + configurable: false, + enumerable: false, + value: hooks, + writable: false, +}); + +function phase() { + const subscribers = []; + return { + get hasSubscribers() { + return subscribers.length > 0; + }, + publish(message) { + for (const subscriber of subscribers) subscriber(message); + }, + runStores(message, fn) { + this.publish(message); + return fn(); + }, + subscribe(subscriber) { + subscribers.push(subscriber); + }, + }; +} + +function getTracingHook(channelName) { + let hook = hooks.get(channelName); + if (hook) return hook; + + hook = { + start: phase(), + end: phase(), + asyncStart: phase(), + asyncEnd: phase(), + error: phase(), + get hasSubscribers() { + return ( + this.start.hasSubscribers || + this.end.hasSubscribers || + this.asyncStart.hasSubscribers || + this.asyncEnd.hasSubscribers || + this.error.hasSubscribers + ); + }, + subscribe(handlers) { + for (const name of ["start", "end", "asyncStart", "asyncEnd", "error"]) { + if (handlers[name]) this[name].subscribe(handlers[name]); + } + }, + traceSync(fn, message) { + return this.start.runStores(message, () => { + try { + message.result = fn(); + return message.result; + } catch (error) { + message.error = error; + this.error.publish(message); + throw error; + } finally { + this.end.publish(message); + } + }); + }, + tracePromise(fn, message) { + return this.start.runStores(message, () => { + try { + const result = fn(); + if (typeof result?.then !== "function") { + message.result = result; + this.end.publish(message); + return result; + } + this.end.publish(message); + const resolve = (value) => { + message.result = value; + this.asyncStart.publish(message); + this.asyncEnd.publish(message); + return value; + }; + const reject = (error) => { + message.error = error; + this.error.publish(message); + this.asyncStart.publish(message); + this.asyncEnd.publish(message); + throw error; + }; + if (result instanceof Promise && result.constructor === Promise) { + return result.then(resolve, reject); + } + result.then(resolve, (error) => { + try { + reject(error); + } catch {} + }); + return result; + } catch (error) { + message.error = error; + this.error.publish(message); + this.end.publish(message); + throw error; + } + }); + }, + traceCallback(fn, position, message) { + const callback = Array.prototype.at.call(message.arguments, position); + if (typeof callback !== "function") return fn(); + const currentHook = this; + function wrappedCallback(error, result) { + if (error) { + message.error = error; + currentHook.error.publish(message); + } else { + message.result = result; + } + return currentHook.asyncStart.runStores(message, () => { + try { + return callback.apply(this, arguments); + } finally { + currentHook.asyncEnd.publish(message); + } + }); + } + Array.prototype.splice.call( + message.arguments, + position, + 1, + wrappedCallback, + ); + return this.start.runStores(message, () => { + try { + return fn(); + } catch (error) { + message.error = error; + this.error.publish(message); + throw error; + } finally { + this.end.publish(message); + } + }); + }, + }; + hooks.set(channelName, hook); + return hook; +} + function getContext(channelName) { - const channel = tracingChannel(channelName); + const channel = getTracingHook(channelName); const context = {}; channel.subscribe({ start(message) { @@ -26,4 +173,4 @@ function getContext(channelName) { }); return context; } -module.exports = { assert, getContext }; +module.exports = { assert, getContext, getTracingHook }; diff --git a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/mod.js b/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/mod.js deleted file mode 100644 index 17b5ed007..000000000 --- a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/mod.js +++ /dev/null @@ -1,9 +0,0 @@ -/** - * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. - * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. - **/ -async function fetch(url) { - return 42; -} - -exports.fetch = fetch; diff --git a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/polyfill.js b/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/polyfill.js deleted file mode 100644 index a42b11f2f..000000000 --- a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/polyfill.js +++ /dev/null @@ -1,8 +0,0 @@ -/** - * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. - * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. - **/ -module.exports = require("node:diagnostics_channel"); -const tracingChannel = module.exports.tracingChannel; -tracingChannel.polyfilled = true; -module.exports.tracingChannel = tracingChannel; diff --git a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/test.js b/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/test.js deleted file mode 100644 index e0dfcfb29..000000000 --- a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_cjs/test.js +++ /dev/null @@ -1,19 +0,0 @@ -/** - * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. - * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. - **/ -const { fetch } = require("./instrumented.js"); -const { assert, getContext } = require("../common/preamble.js"); -const { tracingChannel } = require("node:diagnostics_channel"); -const context = getContext("orchestrion:undici:fetch_decl"); -(async () => { - const result = await fetch("https://example.com"); - assert.strictEqual(result, 42); - assert.deepStrictEqual(context, { - start: true, - end: true, - asyncStart: 42, - asyncEnd: 42, - }); - assert.ok(tracingChannel.polyfilled); -})(); diff --git a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/mod.mjs b/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/mod.mjs deleted file mode 100644 index 8e57364a2..000000000 --- a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/mod.mjs +++ /dev/null @@ -1,9 +0,0 @@ -/** - * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. - * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. - **/ -async function fetch(url) { - return 42; -} - -export { fetch }; diff --git a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/polyfill.js b/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/polyfill.js deleted file mode 100644 index 5ef01bf44..000000000 --- a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/polyfill.js +++ /dev/null @@ -1,9 +0,0 @@ -/** - * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. - * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. - **/ -const dc = require("node:diagnostics_channel"); -const api = {}; -api.tracingChannel = dc.tracingChannel; -api.tracingChannel.polyfilled = true; -module.exports = api; diff --git a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/test.mjs b/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/test.mjs deleted file mode 100644 index 5916b9e61..000000000 --- a/js/tests/auto-instrumentations/fixtures/orchestrion-js/polyfill_mjs/test.mjs +++ /dev/null @@ -1,17 +0,0 @@ -/** - * Unless explicitly stated otherwise all files in this repository are licensed under the Apache-2.0 License. - * This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2025 Datadog, Inc. - **/ -import { fetch } from "./instrumented.mjs"; -import { assert, getContext } from "../common/preamble.js"; -import { tracingChannel } from "diagnostics_channel"; -const context = getContext("orchestrion:undici:fetch_decl"); -const result = await fetch("https://example.com"); -assert.strictEqual(result, 42); -assert.deepStrictEqual(context, { - start: true, - end: true, - asyncStart: 42, - asyncEnd: 42, -}); -assert.ok(tracingChannel.polyfilled); diff --git a/js/tests/auto-instrumentations/fixtures/test-api-promise-preservation.mjs b/js/tests/auto-instrumentations/fixtures/test-api-promise-preservation.mjs deleted file mode 100644 index cf11a448a..000000000 --- a/js/tests/auto-instrumentations/fixtures/test-api-promise-preservation.mjs +++ /dev/null @@ -1,45 +0,0 @@ -import { tracingChannel } from "node:diagnostics_channel"; -import { parentPort } from "node:worker_threads"; - -class HelperPromise extends Promise { - #innerPromise; - - constructor(responsePromise) { - let resolveOuter; - super((resolve) => { - resolveOuter = resolve; - }); - this.#innerPromise = Promise.resolve(responsePromise); - this.#innerPromise.then(resolveOuter); - } - - then(onfulfilled, onrejected) { - return this.#innerPromise.then(onfulfilled, onrejected); - } - - async withResponse() { - return { - data: await this.#innerPromise, - response: { ok: true }, - }; - } -} - -const channel = tracingChannel("braintrust:test:helper-promise"); -const traced = channel.tracePromise( - () => new HelperPromise(Promise.resolve("ok")), - {}, -); - -const withResponse = await traced.withResponse(); - -parentPort?.postMessage({ - type: "helper-result", - result: { - awaitedValue: await traced, - constructorName: traced.constructor.name, - hasWithResponse: typeof traced.withResponse === "function", - withResponseData: withResponse.data, - withResponseOk: withResponse.response.ok, - }, -}); diff --git a/js/tests/auto-instrumentations/function-behavior.test.ts b/js/tests/auto-instrumentations/function-behavior.test.ts index e04d9db0f..1797252b5 100644 --- a/js/tests/auto-instrumentations/function-behavior.test.ts +++ b/js/tests/auto-instrumentations/function-behavior.test.ts @@ -31,8 +31,6 @@ describe("Function Behavior Preservation", () => { if (!fs.existsSync(outputDir)) { fs.mkdirSync(outputDir, { recursive: true }); } - - // Note: dc-browser is now an npm package, no symlinks needed }); afterAll(() => { diff --git a/js/tests/auto-instrumentations/loader-hook.test.ts b/js/tests/auto-instrumentations/loader-hook.test.ts index 51a20be41..48c28a4e3 100644 --- a/js/tests/auto-instrumentations/loader-hook.test.ts +++ b/js/tests/auto-instrumentations/loader-hook.test.ts @@ -16,10 +16,6 @@ const hookPath = path.join( const listenerPath = path.join(fixturesDir, "listener-esm.mjs"); const testAppEsmPath = path.join(fixturesDir, "test-app-esm.mjs"); const testAppCjsPath = path.join(fixturesDir, "test-app-cjs.cjs"); -const helperPromisePath = path.join( - fixturesDir, - "test-api-promise-preservation.mjs", -); const runtimeApplyAutoSideEffectEsmPath = path.join( fixturesDir, "runtime-apply-auto-side-effect-esm.mjs", @@ -43,7 +39,7 @@ describe("Unified Loader Hook Integration Tests", () => { }); describe("Unified hook (--import) handles both ESM and CJS", () => { - it("should emit diagnostics_channel events for ESM OpenAI calls", async () => { + it("should invoke global hooks for ESM OpenAI calls", async () => { const result = await runWithWorker({ execArgv: ["--import", listenerPath, "--import", hookPath], script: testAppEsmPath, @@ -54,7 +50,7 @@ describe("Unified Loader Hook Integration Tests", () => { expect(result.events.start[0].args).toBeDefined(); }); - it("should emit diagnostics_channel events for CJS OpenAI calls", async () => { + it("should invoke global hooks for CJS OpenAI calls", async () => { const result = await runWithWorker({ execArgv: ["--import", listenerPath, "--import", hookPath], script: testAppCjsPath, @@ -63,26 +59,6 @@ describe("Unified Loader Hook Integration Tests", () => { expect(result.events.start.length).toBeGreaterThan(0); expect(result.events.end.length).toBeGreaterThan(0); }); - - it("should preserve helper methods on promise subclasses", async () => { - const result = await runWithWorkerMessage<{ - awaitedValue: string; - constructorName: string; - hasWithResponse: boolean; - withResponseData: string; - withResponseOk: boolean; - }>({ - execArgv: ["--import", hookPath], - messageType: "helper-result", - script: helperPromisePath, - }); - - expect(result.hasWithResponse).toBe(true); - expect(result.awaitedValue).toBe("ok"); - expect(result.withResponseData).toBe("ok"); - expect(result.withResponseOk).toBe(true); - expect(result.constructorName).toBe("HelperPromise"); - }); }); describe("apply-auto-instrumentation side-effect runtime setup", () => { diff --git a/js/tests/auto-instrumentations/multiple-instrumentations.test.ts b/js/tests/auto-instrumentations/multiple-instrumentations.test.ts index 88d3a39d7..2b48ff0a6 100644 --- a/js/tests/auto-instrumentations/multiple-instrumentations.test.ts +++ b/js/tests/auto-instrumentations/multiple-instrumentations.test.ts @@ -33,8 +33,6 @@ describe("Multiple Instrumentations", () => { fs.mkdirSync(outputDir, { recursive: true }); } - // Note: dc-browser is now an npm package, no symlinks needed - // Create mock embeddings module for testing multiple methods const embeddingsDir = path.join(nodeModulesDir, "openai/resources"); if (!fs.existsSync(embeddingsDir)) { diff --git a/js/tests/auto-instrumentations/next-config.test.ts b/js/tests/auto-instrumentations/next-config.test.ts index 0333faab9..91d96d401 100644 --- a/js/tests/auto-instrumentations/next-config.test.ts +++ b/js/tests/auto-instrumentations/next-config.test.ts @@ -131,7 +131,7 @@ describe("wrapNextjsConfigWithBraintrust", () => { rule.loaders[0].loader.includes("webpack-loader"), ), ).toBe(true); - expect(config.turbopack.resolveAlias["dc-browser"]).toContain("dc-browser"); + expect(config.turbopack.resolveAlias).toBeUndefined(); expect(config.turbopack.rules["*.css"]).toEqual({ loaders: ["css-loader"], }); @@ -143,16 +143,12 @@ describe("wrapNextjsConfigWithBraintrust", () => { const config = wrapNextjsConfigWithBraintrust({ turbopack: { resolveAlias: { - "dc-browser": "/custom/dc-browser.js", "user-module": "/custom/user-module.js", }, rules: {}, }, }) as any; - expect(config.turbopack.resolveAlias["dc-browser"]).toBe( - "/custom/dc-browser.js", - ); expect(config.turbopack.resolveAlias["user-module"]).toBe( "/custom/user-module.js", ); @@ -191,14 +187,6 @@ describe("wrapNextjsConfigWithBraintrust", () => { return "/braintrust/webpack-loader.cjs"; } - if (specifier === "braintrust/package.json") { - return "/braintrust/package.json"; - } - - if (specifier === "dc-browser") { - return "/braintrust/node_modules/dc-browser/dist/index.js"; - } - throw new Error(`Cannot resolve module ${specifier}`); }, }, diff --git a/js/tests/auto-instrumentations/orchestrion-js-upstream.test.ts b/js/tests/auto-instrumentations/orchestrion-js-upstream.test.ts index c124cb7bf..633b20c64 100644 --- a/js/tests/auto-instrumentations/orchestrion-js-upstream.test.ts +++ b/js/tests/auto-instrumentations/orchestrion-js-upstream.test.ts @@ -30,7 +30,6 @@ interface UpstreamFixtureCase { mjs?: boolean; filePath?: string; transformerFilePath?: string; - dcModule?: string; } let outputRoot: string; @@ -57,7 +56,6 @@ function runFixture({ mjs = false, filePath = TEST_MODULE_PATH, transformerFilePath = filePath, - dcModule, }: UpstreamFixtureCase): void { const ext = mjs ? "mjs" : "js"; const sourceDir = path.join(fixtureRoot, name); @@ -66,7 +64,7 @@ function runFixture({ fs.rmSync(runDir, { recursive: true, force: true }); fs.cpSync(sourceDir, runDir, { recursive: true }); - const matcher = create(configs, dcModule); + const matcher = create(configs); const transformer = matcher.getTransformer( TEST_MODULE_NAME, TEST_MODULE_VERSION, @@ -321,19 +319,6 @@ const fixtureCases: UpstreamFixtureCase[] = [ }), ], }, - { - name: "polyfill_cjs", - title: "supports a custom diagnostics channel module in CJS", - configs: [config("fetch_decl", { functionName: "fetch", kind: "Async" })], - dcModule: "./polyfill.js", - }, - { - name: "polyfill_mjs", - title: "supports a custom diagnostics channel module in ESM", - configs: [config("fetch_decl", { functionName: "fetch", kind: "Async" })], - dcModule: "./polyfill.js", - mjs: true, - }, { name: "promise_subclass", title: "preserves Promise subclass return values", diff --git a/js/tests/auto-instrumentations/runtime-execution.test.ts b/js/tests/auto-instrumentations/runtime-execution.test.ts index eb5e349ac..eb882dd3f 100644 --- a/js/tests/auto-instrumentations/runtime-execution.test.ts +++ b/js/tests/auto-instrumentations/runtime-execution.test.ts @@ -33,8 +33,6 @@ describe("Runtime Execution of Bundled Code", () => { if (!fs.existsSync(outputDir)) { fs.mkdirSync(outputDir, { recursive: true }); } - - // Note: dc-browser is now an npm package, no symlinks needed }); afterAll(() => { @@ -249,9 +247,6 @@ describe("Runtime Execution of Bundled Code", () => { outDir, emptyOutDir: true, minify: false, - rollupOptions: { - external: ["diagnostics_channel"], - }, }, plugins: [vitePlugin({ browser: false })], logLevel: "error", @@ -311,9 +306,6 @@ describe("Runtime Execution of Bundled Code", () => { outDir, emptyOutDir: true, minify: false, - rollupOptions: { - external: ["diagnostics_channel"], - }, }, plugins: [vitePlugin({ browser: false })], logLevel: "error", diff --git a/js/tests/auto-instrumentations/test-helpers.ts b/js/tests/auto-instrumentations/test-helpers.ts index cbe8f5c81..f2ddd23e0 100644 --- a/js/tests/auto-instrumentations/test-helpers.ts +++ b/js/tests/auto-instrumentations/test-helpers.ts @@ -2,12 +2,11 @@ * Test helpers for functional testing of instrumented code. */ -import * as diagnostics_channel from "diagnostics_channel"; - -// Use type assertion for tracingChannel which may not be in older @types/node -const tracingChannel = (diagnostics_channel as any).tracingChannel as ( - name: string, -) => any; +import { + newGlobalTracingChannel, + type GlobalHookHandlers, + type GlobalTracingChannel, +} from "../../src/global-instrumentation-hooks"; export interface CapturedEvent { arguments?: any[]; @@ -29,9 +28,13 @@ export interface EventCollector { } /** - * Creates an event collector for capturing diagnostics_channel events. + * Creates an event collector for capturing global instrumentation hook events. */ export function createEventCollector(): EventCollector { + const subscriptions: Array<{ + channel: GlobalTracingChannel; + handlers: GlobalHookHandlers; + }> = []; const collector: EventCollector = { start: [], end: [], @@ -46,8 +49,8 @@ export function createEventCollector(): EventCollector { this.error = []; }, subscribe(channelName: string) { - const channel = tracingChannel(channelName); - channel.subscribe({ + const channel = newGlobalTracingChannel(channelName); + const handlers = { start: (ctx: any) => { this.start.push({ arguments: ctx.arguments ? Array.from(ctx.arguments) : undefined, @@ -78,11 +81,14 @@ export function createEventCollector(): EventCollector { timestamp: Date.now(), }); }, - }); + }; + channel.subscribe(handlers); + subscriptions.push({ channel, handlers }); }, unsubscribe() { - // Note: diagnostics_channel doesn't provide a direct unsubscribe API - // In a real scenario, you'd store the channel reference and manage subscriptions + for (const { channel, handlers } of subscriptions.splice(0)) { + channel.unsubscribe(handlers); + } }, }; diff --git a/js/tests/auto-instrumentations/transformation.test.ts b/js/tests/auto-instrumentations/transformation.test.ts index fc402096b..c8d6e46dc 100644 --- a/js/tests/auto-instrumentations/transformation.test.ts +++ b/js/tests/auto-instrumentations/transformation.test.ts @@ -2,10 +2,9 @@ * ORCHESTRION TRANSFORMATION TESTS * * These tests verify that the internal Orchestrion-JS fork correctly - * transforms code to inject tracingChannel calls at build time. + * transforms code to invoke global instrumentation hooks at build time. * * IMPORTANT: Tests use a mock OpenAI package structure in test/fixtures/node_modules/openai. - * IMPORTANT: dc-browser is now an npm package dependency. */ import { describe, it, expect, beforeAll, afterAll } from "vitest"; @@ -18,6 +17,7 @@ import { create, type InstrumentationConfig, } from "../../src/auto-instrumentations/orchestrion-js"; +import { newGlobalTracingChannel } from "../../src/global-instrumentation-hooks"; const __dirname = path.dirname(fileURLToPath(import.meta.url)); const fixturesDir = path.join(__dirname, "fixtures"); @@ -58,14 +58,20 @@ function transformTestCode( return transformer!.transform(code, moduleType); } +function expectGlobalHookTransform(output: string): void { + expect(output).toContain("__braintrust_instrumentation_hooks"); + expect(output).toContain("orchestrion:openai:chat.completions.create"); + expect(output).toContain("__apm$hook.tracePromise"); + expect(output).not.toContain("diagnostics_channel"); + expect(output).not.toContain("dc-browser"); +} + describe("Orchestrion Transformation Tests", () => { beforeAll(() => { // Create output directory if (!fs.existsSync(outputDir)) { fs.mkdirSync(outputDir, { recursive: true }); } - - // Note: dc-browser is now an npm package, no symlinks needed }); afterAll(() => { @@ -89,7 +95,7 @@ describe("Orchestrion Transformation Tests", () => { ); expect(result.code).toContain("orchestrion:test-sdk:test"); - expect(result.code).toContain("tr_ch_apm$test.start.runStores"); + expect(result.code).toContain("__apm$hook.tracePromise"); }); it("supports method-only configs", () => { @@ -105,7 +111,7 @@ describe("Orchestrion Transformation Tests", () => { ); expect(result.code).toContain("orchestrion:test-sdk:test"); - expect(result.code).toContain("tr_ch_apm$test.start.runStores"); + expect(result.code).toContain("__apm$hook.tracePromise"); }); it("supports function declaration configs", () => { @@ -119,7 +125,7 @@ describe("Orchestrion Transformation Tests", () => { ); expect(result.code).toContain("orchestrion:test-sdk:test"); - expect(result.code).toContain("tr_ch_apm$test.start.runStores"); + expect(result.code).toContain("__apm$hook.traceSync"); }); it("supports export-alias function configs", () => { @@ -134,7 +140,7 @@ describe("Orchestrion Transformation Tests", () => { ); expect(result.code).toContain("orchestrion:test-sdk:test"); - expect(result.code).toContain("tr_ch_apm$test.start.runStores"); + expect(result.code).toContain("__apm$hook.traceSync"); }); it("supports export-alias class method configs", () => { @@ -156,7 +162,7 @@ describe("Orchestrion Transformation Tests", () => { ); expect(result.code).toContain("orchestrion:test-sdk:test"); - expect(result.code).toContain("tr_ch_apm$test.start.runStores"); + expect(result.code).toContain("__apm$hook.tracePromise"); }); it("supports private class method configs", () => { @@ -178,7 +184,7 @@ describe("Orchestrion Transformation Tests", () => { ); expect(result.code).toContain("orchestrion:test-sdk:test"); - expect(result.code).toContain("tr_ch_apm$test.start.runStores"); + expect(result.code).toContain("__apm$hook.tracePromise"); }); it("supports object/property configs", () => { @@ -196,7 +202,7 @@ describe("Orchestrion Transformation Tests", () => { ); expect(result.code).toContain("orchestrion:test-sdk:test"); - expect(result.code).toContain("tr_ch_apm$test.start.runStores"); + expect(result.code).toContain("__apm$hook.tracePromise"); }); it("supports callback configs", () => { @@ -210,8 +216,7 @@ describe("Orchestrion Transformation Tests", () => { ); expect(result.code).toContain("orchestrion:test-sdk:test"); - expect(result.code).toContain("Array.prototype.splice.call"); - expect(result.code).toContain("tr_ch_apm$test.asyncStart.runStores"); + expect(result.code).toContain("__apm$hook.traceCallback"); }); it("supports raw AST query configs", () => { @@ -227,7 +232,7 @@ describe("Orchestrion Transformation Tests", () => { ); expect(result.code).toContain("orchestrion:test-sdk:test"); - expect(result.code).toContain("tr_ch_apm$test.start.runStores"); + expect(result.code).toContain("__apm$hook.tracePromise"); }); it("supports index selection", () => { @@ -247,12 +252,10 @@ describe("Orchestrion Transformation Tests", () => { `, ); - const wrapperCount = result.code.match( - /tr_ch_apm\$test\.start\.runStores/g, - ); + const wrapperCount = result.code.match(/__apm\$hook\.tracePromise/g); expect(wrapperCount).toHaveLength(1); expect(result.code.indexOf("secondCreate")).toBeLessThan( - result.code.indexOf("tr_ch_apm$test.start.runStores"), + result.code.indexOf("__apm$hook.tracePromise"), ); }); @@ -272,10 +275,41 @@ describe("Orchestrion Transformation Tests", () => { file: "test-sdk/index.mjs", }); }); + + it("observes hooks registered after the transformed module loads", () => { + const result = transformTestCode( + { functionName: "query", kind: "Sync" }, + ` + function query(input) { + return input; + } + module.exports = { query }; + `, + "cjs", + ); + const loadedModule = { + exports: {} as { query: (input: string) => string }, + }; + Function( + "module", + "exports", + result.code, + )(loadedModule, loadedModule.exports); + + expect(loadedModule.exports.query("before")).toBe("before"); + + const events: unknown[] = []; + newGlobalTracingChannel("orchestrion:test-sdk:test").subscribe({ + start: (event) => events.push(event), + }); + + expect(loadedModule.exports.query("after")).toBe("after"); + expect(events).toHaveLength(1); + }); }); describe("esbuild", () => { - it("should transform OpenAI SDK code with tracingChannel", async () => { + it("should transform OpenAI SDK code with global hooks", async () => { const { braintrustEsbuildPlugin } = await import("../../src/auto-instrumentations/bundler/esbuild.js"); @@ -292,7 +326,7 @@ describe("Orchestrion Transformation Tests", () => { logLevel: "error", absWorkingDir: fixturesDir, preserveSymlinks: true, // CRITICAL: Don't dereference symlinks! - platform: "node", // Allow Node.js built-ins like diagnostics_channel + platform: "node", }); expect(result.errors).toHaveLength(0); @@ -300,13 +334,10 @@ describe("Orchestrion Transformation Tests", () => { const output = fs.readFileSync(outfile, "utf-8"); - // Verify orchestrion transformed the code - expect(output).toContain("tracingChannel"); - expect(output).toContain("orchestrion:openai:chat.completions.create"); - expect(output).not.toContain("TracingChannel"); + expectGlobalHookTransform(output); }); - it("should bundle dc-browser module when useDiagnosticChannelCompatShim is true", async () => { + it("should use global hooks when the legacy compatibility option is true", async () => { const { braintrustEsbuildPlugin } = await import("../../src/auto-instrumentations/bundler/esbuild.js"); @@ -333,19 +364,12 @@ describe("Orchestrion Transformation Tests", () => { const output = fs.readFileSync(outfile, "utf-8"); - // Verify orchestrion transformed the code - expect(output).toContain("tracingChannel"); - expect(output).toContain("orchestrion:openai:chat.completions.create"); - - // Verify dc-browser module is bundled (should contain TracingChannel class implementation) - expect(output).toContain("TracingChannel"); - // Should NOT import from external diagnostics_channel - expect(output).not.toMatch(/from\s+["']diagnostics_channel["']/); + expectGlobalHookTransform(output); }); }); describe("vite", () => { - it("should transform OpenAI SDK code with tracingChannel", async () => { + it("should transform OpenAI SDK code with global hooks", async () => { const { braintrustVitePlugin } = await import("../../src/auto-instrumentations/bundler/vite.js"); @@ -363,9 +387,6 @@ describe("Orchestrion Transformation Tests", () => { outDir, emptyOutDir: true, minify: false, - rollupOptions: { - external: ["diagnostics_channel"], // Mark Node built-ins as external, don't try to bundle them - }, }, plugins: [braintrustVitePlugin()], logLevel: "error", @@ -379,13 +400,10 @@ describe("Orchestrion Transformation Tests", () => { const output = fs.readFileSync(bundlePath, "utf-8"); - // Verify orchestrion transformed the code - expect(output).toContain("tracingChannel"); - expect(output).toContain("orchestrion:openai:chat.completions.create"); - expect(output).not.toContain("TracingChannel"); + expectGlobalHookTransform(output); }); - it("should bundle dc-browser module when useDiagnosticChannelCompatShim is true", async () => { + it("should use global hooks when the legacy compatibility option is true", async () => { const { braintrustVitePlugin } = await import("../../src/auto-instrumentations/bundler/vite.js"); @@ -418,14 +436,7 @@ describe("Orchestrion Transformation Tests", () => { const output = fs.readFileSync(bundlePath, "utf-8"); - // Verify orchestrion transformed the code - expect(output).toContain("tracingChannel"); - expect(output).toContain("orchestrion:openai:chat.completions.create"); - - // Verify dc-browser module is bundled (should contain TracingChannel class implementation) - expect(output).toContain("TracingChannel"); - // Should NOT import from external diagnostics_channel - expect(output).not.toMatch(/from\s+["']diagnostics_channel["']/); + expectGlobalHookTransform(output); }); }); @@ -456,7 +467,7 @@ describe("Orchestrion Transformation Tests", () => { }); } - it("should transform OpenAI SDK code with tracingChannel", async () => { + it("should transform OpenAI SDK code with global hooks", async () => { const { braintrustWebpackPlugin } = await import("../../src/auto-instrumentations/bundler/webpack.js"); @@ -470,17 +481,14 @@ describe("Orchestrion Transformation Tests", () => { experiments: { outputModule: true }, mode: "development", resolve: { modules: [nodeModulesDir, "node_modules"] }, - externals: { diagnostics_channel: "module diagnostics_channel" }, plugins: [braintrustWebpackPlugin()], }); expect(errors).toHaveLength(0); - expect(output).toContain("tracingChannel"); - expect(output).toContain("orchestrion:openai:chat.completions.create"); - expect(output).not.toContain("TracingChannel"); + expectGlobalHookTransform(output); }); - it("should bundle dc-browser module when useDiagnosticChannelCompatShim is true", async () => { + it("should use global hooks when the legacy compatibility option is true", async () => { const { braintrustWebpackPlugin } = await import("../../src/auto-instrumentations/bundler/webpack.js"); @@ -500,10 +508,7 @@ describe("Orchestrion Transformation Tests", () => { }); expect(errors).toHaveLength(0); - expect(output).toContain("tracingChannel"); - expect(output).toContain("orchestrion:openai:chat.completions.create"); - expect(output).toContain("TracingChannel"); - expect(output).not.toMatch(/require\(["']diagnostics_channel["']\)/); + expectGlobalHookTransform(output); }); }); @@ -538,7 +543,7 @@ describe("Orchestrion Transformation Tests", () => { }); } - it("should transform OpenAI SDK code with tracingChannel (turbopack loader-only mode)", async () => { + it("should transform OpenAI SDK code with global hooks (turbopack loader-only mode)", async () => { const { errors, output } = await runWebpackWithLoader({ entry: path.join(fixturesDir, "test-app.js"), output: { @@ -549,7 +554,6 @@ describe("Orchestrion Transformation Tests", () => { experiments: { outputModule: true }, mode: "development", resolve: { modules: [nodeModulesDir, "node_modules"] }, - externals: { diagnostics_channel: "module diagnostics_channel" }, // No plugins — only the loader, mirroring turbopack's constraint module: { rules: [ @@ -566,11 +570,10 @@ describe("Orchestrion Transformation Tests", () => { }); expect(errors).toHaveLength(0); - expect(output).toContain("tracingChannel"); - expect(output).toContain("orchestrion:openai:chat.completions.create"); + expectGlobalHookTransform(output); }); - it("should bundle dc-browser polyfill when browser: true (turbopack loader-only mode)", async () => { + it("should use global hooks when browser mode is true (turbopack loader-only mode)", async () => { const { errors, output } = await runWebpackWithLoader({ entry: path.join(fixturesDir, "test-app.js"), output: { @@ -597,15 +600,12 @@ describe("Orchestrion Transformation Tests", () => { }); expect(errors).toHaveLength(0); - expect(output).toContain("tracingChannel"); - expect(output).toContain("orchestrion:openai:chat.completions.create"); - expect(output).toContain("TracingChannel"); - expect(output).not.toMatch(/require\(["']diagnostics_channel["']\)/); + expectGlobalHookTransform(output); }); }); describe("rollup", () => { - it("should transform OpenAI SDK code with tracingChannel", async () => { + it("should transform OpenAI SDK code with global hooks", async () => { const { rollup } = await import("rollup"); const { braintrustRollupPlugin } = await import("../../src/auto-instrumentations/bundler/rollup.js"); @@ -645,13 +645,10 @@ describe("Orchestrion Transformation Tests", () => { const output = fs.readFileSync(outfile, "utf-8"); - // Verify orchestrion transformed the code - expect(output).toContain("tracingChannel"); - expect(output).toContain("orchestrion:openai:chat.completions.create"); - expect(output).not.toContain("TracingChannel"); + expectGlobalHookTransform(output); }); - it("should bundle dc-browser module when useDiagnosticChannelCompatShim is true", async () => { + it("should use global hooks when the legacy compatibility option is true", async () => { const { rollup } = await import("rollup"); const { braintrustRollupPlugin } = await import("../../src/auto-instrumentations/bundler/rollup.js"); @@ -669,15 +666,6 @@ describe("Orchestrion Transformation Tests", () => { .resolve(fixturesDir, "node_modules", source) .replace(/\\/g, "/"); } - if (source === "dc-browser") { - // Bundler resolveId always returns posix-style paths - return path - .resolve( - __dirname, - "../../node_modules/dc-browser/dist/index.mjs", - ) - .replace(/\\/g, "/"); - } return null; }, }; @@ -703,14 +691,7 @@ describe("Orchestrion Transformation Tests", () => { const output = fs.readFileSync(outfile, "utf-8"); - // Verify orchestrion transformed the code - expect(output).toContain("tracingChannel"); - expect(output).toContain("orchestrion:openai:chat.completions.create"); - - // Verify dc-browser module is bundled (should contain TracingChannel class implementation) - expect(output).toContain("TracingChannel"); - // Should NOT import from external diagnostics_channel - expect(output).not.toMatch(/from\s+["']diagnostics_channel["']/); + expectGlobalHookTransform(output); }); }); }); diff --git a/js/tsup.config.ts b/js/tsup.config.ts index 9780ad447..14bca6b36 100644 --- a/js/tsup.config.ts +++ b/js/tsup.config.ts @@ -98,7 +98,7 @@ export default defineConfig([ entry: ["src/instrumentation/index.ts"], format: ["cjs", "esm"], outDir: "dist/instrumentation", - external: ["dc-browser", "@braintrust/instrumentation-core", "zod"], + external: ["@braintrust/instrumentation-core", "zod"], dts: { compilerOptions: { skipLibCheck: true, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index e51840b13..b9d1af0bf 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -375,9 +375,6 @@ importers: cors: specifier: ^2.8.5 version: 2.8.5 - dc-browser: - specifier: ^1.0.4 - version: 1.0.4 dotenv: specifier: ^16.4.5 version: 16.4.5 @@ -2497,9 +2494,6 @@ packages: dataloader@1.4.0: resolution: {integrity: sha512-68s5jYdlvasItOJnCuI2Q9s4q98g0pCyL3HrcKJu8KNugUl8ahgmZYg38ysLTgQjjXX3H8CJLkAvWrclWfcalw==} - dc-browser@1.0.4: - resolution: {integrity: sha512-7oEtnzNlcE+hr4OvO3GR6Gndgw8BhW+wKOEwMqSleyY7N29jbAxzyW5BaJl7qBCw+6OIxfMWtY0T+6dxq8RWLw==} - debug@4.4.3: resolution: {integrity: sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA==} engines: {node: '>=6.0'} @@ -6500,8 +6494,6 @@ snapshots: dataloader@1.4.0: {} - dc-browser@1.0.4: {} - debug@4.4.3: dependencies: ms: 2.1.3 From 87ade206c1377ece6629b2d69fd333252b8b2aeb Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 24 Jul 2026 12:47:13 +0000 Subject: [PATCH 2/2] Update PR #2275 --- .changeset/global-hooks-cutover.md | 5 +++++ js/src/global-instrumentation-hooks.ts | 10 +++++----- js/src/isomorph.ts | 8 +------- 3 files changed, 11 insertions(+), 12 deletions(-) create mode 100644 .changeset/global-hooks-cutover.md diff --git a/.changeset/global-hooks-cutover.md b/.changeset/global-hooks-cutover.md new file mode 100644 index 000000000..8bedf9548 --- /dev/null +++ b/.changeset/global-hooks-cutover.md @@ -0,0 +1,5 @@ +--- +"braintrust": patch +--- + +wip ref: Replace diagnostic channels with proprietary global hooks diff --git a/js/src/global-instrumentation-hooks.ts b/js/src/global-instrumentation-hooks.ts index 510230e73..214212c67 100644 --- a/js/src/global-instrumentation-hooks.ts +++ b/js/src/global-instrumentation-hooks.ts @@ -13,12 +13,12 @@ export interface GlobalHookAsyncLocalStorage { getStore(): T | undefined; } -export type GlobalHookMessageFunction< - M = any, - N extends string | symbol = string, -> = (message: M, name: N) => void; +type GlobalHookMessageFunction = ( + message: M, + name: N, +) => void; -export type GlobalHookTransformFunction = (message: M) => S; +type GlobalHookTransformFunction = (message: M) => S; export interface GlobalHookChannel< M = any, diff --git a/js/src/isomorph.ts b/js/src/isomorph.ts index 9837dd23e..231d7e6f2 100644 --- a/js/src/isomorph.ts +++ b/js/src/isomorph.ts @@ -5,7 +5,6 @@ import { import { newGlobalTracingChannel, type GlobalHookAsyncLocalStorage, - type GlobalHookChannel, type GlobalHookHandlers, type GlobalTracingChannel, type GlobalTracingChannelCollection, @@ -18,10 +17,6 @@ export interface CallerLocation { } export type IsoAsyncLocalStorage = GlobalHookAsyncLocalStorage; -export type IsoChannel< - M = any, - N extends string | symbol = string, -> = GlobalHookChannel; class DefaultAsyncLocalStorage implements IsoAsyncLocalStorage { constructor() {} @@ -35,8 +30,7 @@ class DefaultAsyncLocalStorage implements IsoAsyncLocalStorage { } } -export type IsoTracingChannelCollection = - GlobalTracingChannelCollection; +type IsoTracingChannelCollection = GlobalTracingChannelCollection; export type IsoTracingChannel = GlobalTracingChannel; export type IsoChannelHandlers = GlobalHookHandlers;