diff --git a/.changeset/remove-plugins-public-api.md b/.changeset/remove-plugins-public-api.md new file mode 100644 index 000000000..21275cc2e --- /dev/null +++ b/.changeset/remove-plugins-public-api.md @@ -0,0 +1,5 @@ +--- +"braintrust": major +--- + +ref!: Remove plugins as part of public API diff --git a/js/src/auto-instrumentations/configs/ai-sdk.ts b/js/src/auto-instrumentations/configs/ai-sdk.ts index cdbb0c1f3..d38234e5e 100644 --- a/js/src/auto-instrumentations/configs/ai-sdk.ts +++ b/js/src/auto-instrumentations/configs/ai-sdk.ts @@ -2,7 +2,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { aiSDKChannels, harnessAgentChannels, -} from "../../instrumentation/plugins/ai-sdk-channels"; +} from "../../instrumentation/providers/ai-sdk-channels"; /** * Instrumentation configurations for the Vercel AI SDK. @@ -222,7 +222,7 @@ export const aiSDKConfigs: InstrumentationConfig[] = [ }, // AI SDK v7 exposes its telemetry lifecycle through a dispatcher created for - // each operation. We patch that dispatcher in the plugin instead of rewriting + // each operation. We patch that dispatcher in the consumer instead of rewriting // the module to call registerTelemetry(). { channelName: aiSDKChannels.v7CreateTelemetryDispatcher.channelName, diff --git a/js/src/auto-instrumentations/configs/anthropic.ts b/js/src/auto-instrumentations/configs/anthropic.ts index 3acaa4d75..862e3528b 100644 --- a/js/src/auto-instrumentations/configs/anthropic.ts +++ b/js/src/auto-instrumentations/configs/anthropic.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { anthropicChannels } from "../../instrumentation/plugins/anthropic-channels"; +import { anthropicChannels } from "../../instrumentation/providers/anthropic-channels"; /** * Instrumentation configurations for the Anthropic SDK. diff --git a/js/src/auto-instrumentations/configs/bedrock-runtime.ts b/js/src/auto-instrumentations/configs/bedrock-runtime.ts index 29b4d0f0d..91dddcfae 100644 --- a/js/src/auto-instrumentations/configs/bedrock-runtime.ts +++ b/js/src/auto-instrumentations/configs/bedrock-runtime.ts @@ -2,7 +2,7 @@ import type { InstrumentationConfig } from "@apm-js-collab/code-transformer"; import { smithyClientChannels, smithyCoreChannels, -} from "../../instrumentation/plugins/bedrock-runtime-channels"; +} from "../../instrumentation/providers/bedrock-runtime-channels"; export const bedrockRuntimeConfigs: InstrumentationConfig[] = [ { diff --git a/js/src/auto-instrumentations/configs/claude-agent-sdk.ts b/js/src/auto-instrumentations/configs/claude-agent-sdk.ts index 1cf2c6537..6b663df61 100644 --- a/js/src/auto-instrumentations/configs/claude-agent-sdk.ts +++ b/js/src/auto-instrumentations/configs/claude-agent-sdk.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { claudeAgentSDKChannels } from "../../instrumentation/plugins/claude-agent-sdk-channels"; +import { claudeAgentSDKChannels } from "../../instrumentation/providers/claude-agent-sdk-channels"; /** * Instrumentation configuration for the Claude Agent SDK. diff --git a/js/src/auto-instrumentations/configs/cloudflare-agents.ts b/js/src/auto-instrumentations/configs/cloudflare-agents.ts index 0b655bb78..b18372cb1 100644 --- a/js/src/auto-instrumentations/configs/cloudflare-agents.ts +++ b/js/src/auto-instrumentations/configs/cloudflare-agents.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { cloudflareAgentsChannels } from "../../instrumentation/plugins/cloudflare-agents-channels"; +import { cloudflareAgentsChannels } from "../../instrumentation/providers/cloudflare-agents-channels"; const cloudflareAgentsVersionRange = ">=0.17.0 <0.18.0"; diff --git a/js/src/auto-instrumentations/configs/cloudflare-ai-chat.ts b/js/src/auto-instrumentations/configs/cloudflare-ai-chat.ts index 02f081612..1a603bffc 100644 --- a/js/src/auto-instrumentations/configs/cloudflare-ai-chat.ts +++ b/js/src/auto-instrumentations/configs/cloudflare-ai-chat.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { cloudflareAIChatChannels } from "../../instrumentation/plugins/cloudflare-ai-chat-channels"; +import { cloudflareAIChatChannels } from "../../instrumentation/providers/cloudflare-ai-chat-channels"; const cloudflareAIChatVersionRange = ">=0.9.0 <0.10.0"; diff --git a/js/src/auto-instrumentations/configs/cloudflare-think.ts b/js/src/auto-instrumentations/configs/cloudflare-think.ts index 61fb4e8ae..d0e270b0f 100644 --- a/js/src/auto-instrumentations/configs/cloudflare-think.ts +++ b/js/src/auto-instrumentations/configs/cloudflare-think.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { cloudflareThinkChannels } from "../../instrumentation/plugins/cloudflare-think-channels"; +import { cloudflareThinkChannels } from "../../instrumentation/providers/cloudflare-think-channels"; const cloudflareThinkVersionRange = ">=0.13.0 <0.14.0"; diff --git a/js/src/auto-instrumentations/configs/cohere.ts b/js/src/auto-instrumentations/configs/cohere.ts index 106980884..1ad936ac7 100644 --- a/js/src/auto-instrumentations/configs/cohere.ts +++ b/js/src/auto-instrumentations/configs/cohere.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { cohereChannels } from "../../instrumentation/plugins/cohere-channels"; +import { cohereChannels } from "../../instrumentation/providers/cohere-channels"; export const cohereConfigs: InstrumentationConfig[] = [ { diff --git a/js/src/auto-instrumentations/configs/cursor-sdk.ts b/js/src/auto-instrumentations/configs/cursor-sdk.ts index ecb845b4b..0b0eddb32 100644 --- a/js/src/auto-instrumentations/configs/cursor-sdk.ts +++ b/js/src/auto-instrumentations/configs/cursor-sdk.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { cursorSDKChannels } from "../../instrumentation/plugins/cursor-sdk-channels"; +import { cursorSDKChannels } from "../../instrumentation/providers/cursor-sdk-channels"; const cursorSDKVersionRange = ">=1.0.7 <2.0.0"; diff --git a/js/src/auto-instrumentations/configs/genkit.ts b/js/src/auto-instrumentations/configs/genkit.ts index 349635b01..5422a08fe 100644 --- a/js/src/auto-instrumentations/configs/genkit.ts +++ b/js/src/auto-instrumentations/configs/genkit.ts @@ -2,7 +2,7 @@ import type { InstrumentationConfig } from "../orchestrion-js"; import { genkitChannels, genkitCoreChannels, -} from "../../instrumentation/plugins/genkit-channels"; +} from "../../instrumentation/providers/genkit-channels"; const genkitVersionRange = ">=1.0.0 <2.0.0"; diff --git a/js/src/auto-instrumentations/configs/github-copilot.ts b/js/src/auto-instrumentations/configs/github-copilot.ts index 901f9e682..aed9eb7ed 100644 --- a/js/src/auto-instrumentations/configs/github-copilot.ts +++ b/js/src/auto-instrumentations/configs/github-copilot.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { gitHubCopilotChannels } from "../../instrumentation/plugins/github-copilot-channels"; +import { gitHubCopilotChannels } from "../../instrumentation/providers/github-copilot-channels"; export const gitHubCopilotConfigs: InstrumentationConfig[] = [ // ESM: CopilotClient.createSession diff --git a/js/src/auto-instrumentations/configs/google-adk.ts b/js/src/auto-instrumentations/configs/google-adk.ts index aeb72c44b..8c07fdbd7 100644 --- a/js/src/auto-instrumentations/configs/google-adk.ts +++ b/js/src/auto-instrumentations/configs/google-adk.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { googleADKChannels } from "../../instrumentation/plugins/google-adk-channels"; +import { googleADKChannels } from "../../instrumentation/providers/google-adk-channels"; const googleADKVersionRange = ">=0.1.0"; const googleADKBundledIndexV06VersionRange = ">=0.6.1 <0.7.0"; diff --git a/js/src/auto-instrumentations/configs/google-genai.ts b/js/src/auto-instrumentations/configs/google-genai.ts index 16edd1d5c..53fb24d43 100644 --- a/js/src/auto-instrumentations/configs/google-genai.ts +++ b/js/src/auto-instrumentations/configs/google-genai.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { googleGenAIChannels } from "../../instrumentation/plugins/google-genai-channels"; +import { googleGenAIChannels } from "../../instrumentation/providers/google-genai-channels"; /** * Instrumentation configurations for the Google GenAI SDK. diff --git a/js/src/auto-instrumentations/configs/groq.ts b/js/src/auto-instrumentations/configs/groq.ts index 864b2a166..5db99b980 100644 --- a/js/src/auto-instrumentations/configs/groq.ts +++ b/js/src/auto-instrumentations/configs/groq.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { groqChannels } from "../../instrumentation/plugins/groq-channels"; +import { groqChannels } from "../../instrumentation/providers/groq-channels"; export const groqConfigs: InstrumentationConfig[] = [ { diff --git a/js/src/auto-instrumentations/configs/huggingface-transformers.ts b/js/src/auto-instrumentations/configs/huggingface-transformers.ts index dd951ac25..781aab990 100644 --- a/js/src/auto-instrumentations/configs/huggingface-transformers.ts +++ b/js/src/auto-instrumentations/configs/huggingface-transformers.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { huggingFaceTransformersChannels } from "../../instrumentation/plugins/huggingface-transformers-channels"; +import { huggingFaceTransformersChannels } from "../../instrumentation/providers/huggingface-transformers-channels"; const moduleFiles = [ { diff --git a/js/src/auto-instrumentations/configs/huggingface.ts b/js/src/auto-instrumentations/configs/huggingface.ts index f4c9bec06..53b4d011f 100644 --- a/js/src/auto-instrumentations/configs/huggingface.ts +++ b/js/src/auto-instrumentations/configs/huggingface.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { huggingFaceChannels } from "../../instrumentation/plugins/huggingface-channels"; +import { huggingFaceChannels } from "../../instrumentation/providers/huggingface-channels"; export const huggingFaceConfigs: InstrumentationConfig[] = [ { diff --git a/js/src/auto-instrumentations/configs/langchain.ts b/js/src/auto-instrumentations/configs/langchain.ts index e8108ed02..c1263a285 100644 --- a/js/src/auto-instrumentations/configs/langchain.ts +++ b/js/src/auto-instrumentations/configs/langchain.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { langChainChannels } from "../../instrumentation/plugins/langchain-channels"; +import { langChainChannels } from "../../instrumentation/providers/langchain-channels"; const langChainCoreVersionRange = ">=0.3.42"; const langChainCallbackManagerFilePath = "dist/callbacks/manager.js"; diff --git a/js/src/auto-instrumentations/configs/langsmith.ts b/js/src/auto-instrumentations/configs/langsmith.ts index 5ecc54f87..0904bc64c 100644 --- a/js/src/auto-instrumentations/configs/langsmith.ts +++ b/js/src/auto-instrumentations/configs/langsmith.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { langSmithChannels } from "../../instrumentation/plugins/langsmith-channels"; +import { langSmithChannels } from "../../instrumentation/providers/langsmith-channels"; const versionRange = ">=0.3.30 <1.0.0"; diff --git a/js/src/auto-instrumentations/configs/mistral.ts b/js/src/auto-instrumentations/configs/mistral.ts index e61fc24be..dab2493c3 100644 --- a/js/src/auto-instrumentations/configs/mistral.ts +++ b/js/src/auto-instrumentations/configs/mistral.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { mistralChannels } from "../../instrumentation/plugins/mistral-channels"; +import { mistralChannels } from "../../instrumentation/providers/mistral-channels"; export const mistralConfigs: InstrumentationConfig[] = [ { diff --git a/js/src/auto-instrumentations/configs/ollama.ts b/js/src/auto-instrumentations/configs/ollama.ts index 5ef28eb64..8403bb652 100644 --- a/js/src/auto-instrumentations/configs/ollama.ts +++ b/js/src/auto-instrumentations/configs/ollama.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { ollamaChannels } from "../../instrumentation/plugins/ollama-channels"; +import { ollamaChannels } from "../../instrumentation/providers/ollama-channels"; const methods = [ ["chat", ollamaChannels.chat.channelName], diff --git a/js/src/auto-instrumentations/configs/openai-agents.ts b/js/src/auto-instrumentations/configs/openai-agents.ts index 3647edbc5..8bf95e99f 100644 --- a/js/src/auto-instrumentations/configs/openai-agents.ts +++ b/js/src/auto-instrumentations/configs/openai-agents.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { openAIAgentsCoreChannels } from "../../instrumentation/plugins/openai-agents-channels"; +import { openAIAgentsCoreChannels } from "../../instrumentation/providers/openai-agents-channels"; const lifecycleMethods = [ ["onTraceStart", openAIAgentsCoreChannels.onTraceStart.channelName], diff --git a/js/src/auto-instrumentations/configs/openai-codex.ts b/js/src/auto-instrumentations/configs/openai-codex.ts index 40ee60784..a725caeb9 100644 --- a/js/src/auto-instrumentations/configs/openai-codex.ts +++ b/js/src/auto-instrumentations/configs/openai-codex.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { openAICodexChannels } from "../../instrumentation/plugins/openai-codex-channels"; +import { openAICodexChannels } from "../../instrumentation/providers/openai-codex-channels"; const openAICodexVersionRange = ">=0.128.0 <1.0.0"; diff --git a/js/src/auto-instrumentations/configs/openai.ts b/js/src/auto-instrumentations/configs/openai.ts index 28d85b7f6..411bd2654 100644 --- a/js/src/auto-instrumentations/configs/openai.ts +++ b/js/src/auto-instrumentations/configs/openai.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { openAIChannels } from "../../instrumentation/plugins/openai-channels"; +import { openAIChannels } from "../../instrumentation/providers/openai-channels"; /** * Instrumentation configurations for the OpenAI SDK. diff --git a/js/src/auto-instrumentations/configs/openrouter-agent.ts b/js/src/auto-instrumentations/configs/openrouter-agent.ts index f9f6191e1..d228c1ed8 100644 --- a/js/src/auto-instrumentations/configs/openrouter-agent.ts +++ b/js/src/auto-instrumentations/configs/openrouter-agent.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { openRouterAgentChannels } from "../../instrumentation/plugins/openrouter-agent-channels"; +import { openRouterAgentChannels } from "../../instrumentation/providers/openrouter-agent-channels"; export const openRouterAgentConfigs: InstrumentationConfig[] = [ { diff --git a/js/src/auto-instrumentations/configs/openrouter.ts b/js/src/auto-instrumentations/configs/openrouter.ts index c7b2ce7e4..2607b3da4 100644 --- a/js/src/auto-instrumentations/configs/openrouter.ts +++ b/js/src/auto-instrumentations/configs/openrouter.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { openRouterChannels } from "../../instrumentation/plugins/openrouter-channels"; +import { openRouterChannels } from "../../instrumentation/providers/openrouter-channels"; export const openRouterConfigs: InstrumentationConfig[] = [ { diff --git a/js/src/auto-instrumentations/configs/pi-coding-agent.ts b/js/src/auto-instrumentations/configs/pi-coding-agent.ts index 0bf6e2a56..3401aec0d 100644 --- a/js/src/auto-instrumentations/configs/pi-coding-agent.ts +++ b/js/src/auto-instrumentations/configs/pi-coding-agent.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { piCodingAgentChannels } from "../../instrumentation/plugins/pi-coding-agent-channels"; +import { piCodingAgentChannels } from "../../instrumentation/providers/pi-coding-agent-channels"; const piCodingAgentVersionRange = ">=0.79.0 <0.82.0"; diff --git a/js/src/auto-instrumentations/configs/strands-agent-sdk.ts b/js/src/auto-instrumentations/configs/strands-agent-sdk.ts index 7a7f7a5b9..f02439a19 100644 --- a/js/src/auto-instrumentations/configs/strands-agent-sdk.ts +++ b/js/src/auto-instrumentations/configs/strands-agent-sdk.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "@apm-js-collab/code-transformer"; -import { strandsAgentSDKChannels } from "../../instrumentation/plugins/strands-agent-sdk-channels"; +import { strandsAgentSDKChannels } from "../../instrumentation/providers/strands-agent-sdk-channels"; const strandsAgentSDKVersionRange = ">=1.0.0 <2.0.0"; diff --git a/js/src/auto-instrumentations/configs/voyageai.ts b/js/src/auto-instrumentations/configs/voyageai.ts index f54edf84e..9fbe886be 100644 --- a/js/src/auto-instrumentations/configs/voyageai.ts +++ b/js/src/auto-instrumentations/configs/voyageai.ts @@ -1,5 +1,5 @@ import type { InstrumentationConfig } from "../orchestrion-js"; -import { voyageAIChannels } from "../../instrumentation/plugins/voyageai-channels"; +import { voyageAIChannels } from "../../instrumentation/providers/voyageai-channels"; const GENERATED_CLIENT_MODULES = [ "dist/cjs/Client.js", diff --git a/js/src/auto-instrumentations/loader/cjs-patch.ts b/js/src/auto-instrumentations/loader/cjs-patch.ts index df5c448ae..d94511563 100644 --- a/js/src/auto-instrumentations/loader/cjs-patch.ts +++ b/js/src/auto-instrumentations/loader/cjs-patch.ts @@ -50,7 +50,7 @@ export class ModulePatch { // Per-package source patches (see loader/special-case-patches.ts). // Anti-pattern intentionally isolated in its own module — do not - // expand inline here; new integrations belong in the standard plugin + // expand inline here; new integrations belong in the standard consumer // pipeline. const patched = applySpecialCasePatch({ packageName, diff --git a/js/src/auto-instrumentations/loader/special-case-patches.ts b/js/src/auto-instrumentations/loader/special-case-patches.ts index f42215592..bf94b5926 100644 --- a/js/src/auto-instrumentations/loader/special-case-patches.ts +++ b/js/src/auto-instrumentations/loader/special-case-patches.ts @@ -6,8 +6,8 @@ * Every entry in this file represents a target SDK that doesn't expose a * 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`. + * **prefer the standard channel-consumer pattern** used by the integrations in + * `js/src/instrumentation/providers/*-plugin.ts`. * Only add an entry here when the target SDK gives us no other option (e.g. * the SDK relies on side-effectful module initialization, exposes its API * exclusively as re-exports from content-hashed chunks, or otherwise can't diff --git a/js/src/exports.ts b/js/src/exports.ts index 9ed496ceb..fe6cbd7f5 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -196,7 +196,7 @@ export { BraintrustMiddleware, wrapAISDKModel, } from "./wrappers/ai-sdk"; -export { braintrustEveInstrumentation } from "./instrumentation/plugins/eve-plugin"; +export { braintrustEveInstrumentation } from "./instrumentation"; export { collectAnthropicSession } from "./wrappers/anthropic-session-collector"; export { wrapAnthropic } from "./wrappers/anthropic"; export { diff --git a/js/src/imports.test.ts b/js/src/imports.test.ts index c2d53c767..e00c08223 100644 --- a/js/src/imports.test.ts +++ b/js/src/imports.test.ts @@ -133,11 +133,13 @@ describe("Import restrictions", () => { // Check for dynamic import() statements // Match import(...) but not static import statements - // Exception: allow dynamic import in anthropic-plugin for APIPromise patching + // Exception: allow dynamic import in anthropic-instrumentation for APIPromise patching if ( /\bimport\s*\(/.test(line) && !/^import\s+/.test(line.trim()) && - !relativePath.includes("instrumentation/plugins/anthropic-plugin.ts") + !relativePath.includes( + "instrumentation/providers/anthropic-instrumentation.ts", + ) ) { violations.push( `${relativePath}:${index + 1} - Found dynamic import() statement: "${line.trim()}"`, diff --git a/js/src/instrumentation/README.md b/js/src/instrumentation/README.md index 26487ec26..1fc368594 100644 --- a/js/src/instrumentation/README.md +++ b/js/src/instrumentation/README.md @@ -1,7 +1,8 @@ -# Writing Braintrust Instrumentation Plugins +# Braintrust Instrumentation Consumers -Braintrust instrumentation plugins wrap provider calls through typed invocation -hooks or consume tracing-compatible events from the internal global registry. +Braintrust instrumentation consumers wrap provider calls through typed +invocation hooks or consume tracing-compatible events from the internal global +registry. Auto-instrumented provider code and manual wrappers use the same typed channels, so extraction, stream handling, and span behavior stay aligned. @@ -13,15 +14,15 @@ An instrumentation has four parts: transformation. 2. A typed channel defines its arguments, result, extra event fields, and stable `orchestrion::` identifier. -3. A plugin intercepts that channel, or subscribes to its legacy tracing - lifecycle, and maps the call into Braintrust spans. +3. An internal consumer intercepts that channel, or subscribes to its legacy + tracing lifecycle, and maps the call into Braintrust spans. 4. A manual wrapper invokes the same typed channel when transformation is not available. -The global hook transport is internal. New and migrated plugins should prefer -the typed channel's `intercept` API. Existing plugins can continue using -`traceAsyncChannel`, `traceStreamingChannel`, `traceSyncStreamChannel`, or -`BasePlugin` helpers during the gradual migration. +The global hook transport and its consumers are internal. New and migrated +consumers should prefer the typed channel's `intercept` API. Existing consumers +can continue using `traceAsyncChannel`, `traceStreamingChannel`, or +`traceSyncStreamChannel` during the gradual migration. ## Invocation Hooks @@ -100,27 +101,25 @@ construct it from the package and operation. Prefer the shared tracing helpers: ```ts -this.register( - traceAsyncChannel(providerChannels.create, { - name: "provider.messages.create", - type: "llm", - extractInput(args) { - return { - input: args[0].messages, - metadata: { model: args[0].model }, - }; - }, - extractOutput(result) { - return result.content; - }, - extractMetrics(result) { - return { - prompt_tokens: result.usage.input_tokens, - completion_tokens: result.usage.output_tokens, - }; - }, - }), -); +traceAsyncChannel(providerChannels.create, { + name: "provider.messages.create", + type: "llm", + extractInput(args) { + return { + input: args[0].messages, + metadata: { model: args[0].model }, + }; + }, + extractOutput(result) { + return result.content; + }, + extractMetrics(result) { + return { + prompt_tokens: result.usage.input_tokens, + completion_tokens: result.usage.output_tokens, + }; + }, +}); ``` The helpers: @@ -129,7 +128,7 @@ The helpers: - 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 +- install process-lifetime subscriptions and span-store bindings Use raw `IsoChannelHandlers` only when a provider requires lifecycle behavior that the shared helpers cannot express. @@ -146,7 +145,7 @@ return providerChannels.create.invoke(originalCreate, this, [params], { ``` Legacy wrappers can continue calling the tracing-compatible operators until -their plugin is migrated: +their consumer is migrated: ```ts return providerChannels.create.tracePromise(() => originalCreate(params), { @@ -155,7 +154,7 @@ return providerChannels.create.tracePromise(() => originalCreate(params), { ``` Do not create spans directly inside wrappers. Keeping span creation in the -plugin prevents auto and manual instrumentation from drifting. +internal consumer prevents auto and manual instrumentation from drifting. ## Promise and Stream Requirements @@ -177,14 +176,14 @@ termination, and async context. - 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. +- Use narrow vendored provider interfaces shared by wrappers and consumers. +- Keep enable, subscription, and patching behavior idempotent. ## Testing Test at the narrowest useful layers: -1. Plugin unit tests for extraction and span handling. +1. Consumer 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. diff --git a/js/src/instrumentation/auto-instrumentation-suppression.ts b/js/src/instrumentation/auto-instrumentation-suppression.ts index 2b31044ef..6002e0d1f 100644 --- a/js/src/instrumentation/auto-instrumentation-suppression.ts +++ b/js/src/instrumentation/auto-instrumentation-suppression.ts @@ -45,10 +45,10 @@ export function runWithAutoInstrumentationSuppressed(callback: () => R): R { export function bindAutoInstrumentationSuppressionToStart( tracingChannel: Pick, "start">, -): (() => void) | undefined { +): void { const startChannel = tracingChannel.start; if (!startChannel) { - return undefined; + return; } const store = suppressionStore(); @@ -61,10 +61,6 @@ export function bindAutoInstrumentationSuppressionToStart( }, ], })); - - return () => { - startChannel.unbindStore(store); - }; } export function enterAutoInstrumentationAllowed(): () => void { diff --git a/js/src/instrumentation/braintrust-plugin.test.ts b/js/src/instrumentation/braintrust-plugin.test.ts deleted file mode 100644 index 3365ec279..000000000 --- a/js/src/instrumentation/braintrust-plugin.test.ts +++ /dev/null @@ -1,1791 +0,0 @@ -import { describe, it, expect, vi, beforeEach } from "vitest"; -import { BraintrustPlugin } from "./braintrust-plugin"; -import { OpenAIPlugin } from "./plugins/openai-plugin"; -import { OpenAICodexPlugin } from "./plugins/openai-codex-plugin"; -import { AnthropicPlugin } from "./plugins/anthropic-plugin"; -import { AISDKPlugin } from "./plugins/ai-sdk-plugin"; -import { ClaudeAgentSDKPlugin } from "./plugins/claude-agent-sdk-plugin"; -import { CloudflareThinkPlugin } from "./plugins/cloudflare-think-plugin"; -import { OpenAIAgentsPlugin } from "./plugins/openai-agents-plugin"; -import { GoogleGenAIPlugin } from "./plugins/google-genai-plugin"; -import { HuggingFacePlugin } from "./plugins/huggingface-plugin"; -import { HuggingFaceTransformersPlugin } from "./plugins/huggingface-transformers-plugin"; -import { OpenRouterAgentPlugin } from "./plugins/openrouter-agent-plugin"; -import { OpenRouterPlugin } from "./plugins/openrouter-plugin"; -import { MistralPlugin } from "./plugins/mistral-plugin"; -import { OllamaPlugin } from "./plugins/ollama-plugin"; -import { CoherePlugin } from "./plugins/cohere-plugin"; -import { GroqPlugin } from "./plugins/groq-plugin"; -import { GitHubCopilotPlugin } from "./plugins/github-copilot-plugin"; -import { LangChainPlugin } from "./plugins/langchain-plugin"; -import { LangSmithPlugin } from "./plugins/langsmith-plugin"; -import { PiCodingAgentPlugin } from "./plugins/pi-coding-agent-plugin"; -import { StrandsAgentSDKPlugin } from "./plugins/strands-agent-sdk-plugin"; -import { VoyageAIPlugin } from "./plugins/voyageai-plugin"; -import { CloudflareAIChatPlugin } from "./plugins/cloudflare-ai-chat-plugin"; -import { CloudflareAgentsPlugin } from "./plugins/cloudflare-agents-plugin"; - -function createPluginClassMock() { - return vi.fn(function MockPlugin(this: { - enable: ReturnType; - disable: ReturnType; - }) { - this.enable = vi.fn(); - this.disable = vi.fn(); - }); -} - -// Preserve the re-exported utility functions (parseMetricsFromUsage, etc.) -// while mocking out the OpenAIPlugin class — those utilities are also -// imported and tested at the bottom of this file via braintrust-plugin's -// re-exports. -vi.mock("./plugins/openai-plugin", async () => { - const actual = await vi.importActual< - typeof import("./plugins/openai-plugin") - >("./plugins/openai-plugin"); - return { - ...actual, - OpenAIPlugin: createPluginClassMock(), - }; -}); - -vi.mock("./plugins/anthropic-plugin", () => ({ - AnthropicPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/openai-codex-plugin", () => ({ - OpenAICodexPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/ai-sdk-plugin", () => ({ - AISDKPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/claude-agent-sdk-plugin", () => ({ - ClaudeAgentSDKPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/cloudflare-think-plugin", () => ({ - CloudflareThinkPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/openai-agents-plugin", () => ({ - OpenAIAgentsPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/google-genai-plugin", () => ({ - GoogleGenAIPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/huggingface-plugin", () => ({ - HuggingFacePlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/huggingface-transformers-plugin", () => ({ - HuggingFaceTransformersPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/openrouter-plugin", () => ({ - OpenRouterPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/openrouter-agent-plugin", () => ({ - OpenRouterAgentPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/mistral-plugin", () => ({ - MistralPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/ollama-plugin", () => ({ - OllamaPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/cohere-plugin", () => ({ - CoherePlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/groq-plugin", () => ({ - GroqPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/github-copilot-plugin", () => ({ - GitHubCopilotPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/langchain-plugin", () => ({ - LangChainPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/langsmith-plugin", () => ({ - LangSmithPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/pi-coding-agent-plugin", () => ({ - PiCodingAgentPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/strands-agent-sdk-plugin", () => ({ - StrandsAgentSDKPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/voyageai-plugin", () => ({ - VoyageAIPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/cloudflare-ai-chat-plugin", () => ({ - CloudflareAIChatPlugin: createPluginClassMock(), -})); - -vi.mock("./plugins/cloudflare-agents-plugin", () => ({ - CloudflareAgentsPlugin: createPluginClassMock(), -})); - -describe("BraintrustPlugin", () => { - beforeEach(() => { - vi.clearAllMocks(); - }); - - describe("sub-plugin creation (default behavior)", () => { - it("should create and enable OpenAI plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(OpenAIPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable Anthropic plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(AnthropicPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable OpenAI Codex plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(OpenAICodexPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(OpenAICodexPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable AI SDK plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(AISDKPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable Claude Agent SDK plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - const mockInstance = - vi.mocked(ClaudeAgentSDKPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable Cloudflare Think plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(CloudflareThinkPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(CloudflareThinkPlugin).mock.results[0] - .value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable OpenAI Agents plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(OpenAIAgentsPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(OpenAIAgentsPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable Google GenAI plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(GoogleGenAIPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable HuggingFace plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(HuggingFacePlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - expect(HuggingFaceTransformersPlugin).toHaveBeenCalledTimes(1); - const transformersMockInstance = vi.mocked(HuggingFaceTransformersPlugin) - .mock.results[0].value; - expect(transformersMockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable OpenRouter plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(OpenRouterPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable OpenRouter Agent plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(OpenRouterAgentPlugin).mock.results[0] - .value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable Mistral plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(MistralPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(MistralPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable Cohere plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(CoherePlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(CoherePlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable Voyage AI plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(VoyageAIPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(VoyageAIPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable Groq plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(GroqPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(GroqPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable GitHubCopilot plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(GitHubCopilotPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(GitHubCopilotPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable LangChain plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(LangChainPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(LangChainPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable Strands Agent SDK plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(StrandsAgentSDKPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(StrandsAgentSDKPlugin).mock.results[0] - .value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable Cloudflare AI Chat plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(CloudflareAIChatPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(CloudflareAIChatPlugin).mock.results[0] - .value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create and enable Cloudflare Agents plugin by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(CloudflareAgentsPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(CloudflareAgentsPlugin).mock.results[0] - .value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should create all plugins when enabled with no config", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(OpenAICodexPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(OpenAIAgentsPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - expect(CoherePlugin).toHaveBeenCalledTimes(1); - expect(GroqPlugin).toHaveBeenCalledTimes(1); - expect(GitHubCopilotPlugin).toHaveBeenCalledTimes(1); - expect(StrandsAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(LangChainPlugin).toHaveBeenCalledTimes(1); - }); - - it("should create all plugins when enabled with empty config", () => { - const plugin = new BraintrustPlugin({}); - plugin.enable(); - - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(OpenAICodexPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(OpenAIAgentsPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - expect(CoherePlugin).toHaveBeenCalledTimes(1); - expect(GroqPlugin).toHaveBeenCalledTimes(1); - expect(GitHubCopilotPlugin).toHaveBeenCalledTimes(1); - expect(StrandsAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(LangChainPlugin).toHaveBeenCalledTimes(1); - }); - - it("should create all plugins when enabled with empty integrations config", () => { - const plugin = new BraintrustPlugin({ integrations: {} }); - plugin.enable(); - - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(OpenAICodexPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(OpenAIAgentsPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - expect(CoherePlugin).toHaveBeenCalledTimes(1); - expect(GroqPlugin).toHaveBeenCalledTimes(1); - expect(StrandsAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(LangChainPlugin).toHaveBeenCalledTimes(1); - }); - }); - - describe("config-based disabling", () => { - it("should not create OpenAI plugin when openai: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { openai: false }, - }); - plugin.enable(); - - expect(OpenAIPlugin).not.toHaveBeenCalled(); - // Other plugins should still be created - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(OpenAIAgentsPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Anthropic plugin when anthropic: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { anthropic: false }, - }); - plugin.enable(); - - expect(AnthropicPlugin).not.toHaveBeenCalled(); - // Other plugins should still be created - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(OpenAICodexPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(OpenAIAgentsPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create OpenAI Codex plugin when openaiCodexSDK: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { openaiCodexSDK: false }, - }); - plugin.enable(); - - expect(OpenAICodexPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create AI SDK plugin when aisdk: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { aisdk: false }, - }); - plugin.enable(); - - expect(AISDKPlugin).not.toHaveBeenCalled(); - // Other plugins should still be created - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Claude Agent SDK plugin when claudeAgentSDK: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { claudeAgentSDK: false }, - }); - plugin.enable(); - - expect(ClaudeAgentSDKPlugin).not.toHaveBeenCalled(); - // Other plugins should still be created - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(OpenAIAgentsPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Cloudflare Think plugin when cloudflareThink: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { cloudflareThink: false }, - }); - plugin.enable(); - - expect(CloudflareThinkPlugin).not.toHaveBeenCalled(); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create OpenAI Agents plugin when openAIAgents: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { openAIAgents: false }, - }); - plugin.enable(); - - expect(OpenAIAgentsPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create HuggingFace plugin when huggingface: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { huggingface: false }, - }); - plugin.enable(); - - expect(HuggingFacePlugin).not.toHaveBeenCalled(); - expect(HuggingFaceTransformersPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Google GenAI plugin when googleGenAI: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { googleGenAI: false }, - }); - plugin.enable(); - - expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); - // Other plugins should still be created - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create OpenRouter plugin when openrouter: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { openrouter: false }, - }); - plugin.enable(); - - expect(OpenRouterPlugin).not.toHaveBeenCalled(); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Mistral plugin when mistral: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { mistral: false }, - }); - plugin.enable(); - - expect(MistralPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Ollama plugin when ollama: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { ollama: false }, - }); - plugin.enable(); - - expect(OllamaPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Cohere plugin when cohere: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { cohere: false }, - }); - plugin.enable(); - - expect(CoherePlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Groq plugin when groq: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { groq: false }, - }); - plugin.enable(); - - expect(GroqPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create GitHubCopilot plugin when gitHubCopilot: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { gitHubCopilot: false }, - }); - plugin.enable(); - - expect(GitHubCopilotPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(GroqPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create LangChain plugin when langchain: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { langchain: false }, - }); - plugin.enable(); - - expect(LangChainPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - }); - - it("should create LangSmith with LangChain deduplication by default", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - expect(LangSmithPlugin).toHaveBeenCalledWith({ - skipLangChainRuns: true, - }); - const mockInstance = vi.mocked(LangSmithPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should disable LangSmith or let it capture LangChain runs", () => { - const disabled = new BraintrustPlugin({ - integrations: { langsmith: false }, - }); - disabled.enable(); - expect(LangSmithPlugin).not.toHaveBeenCalled(); - - const withoutLangChain = new BraintrustPlugin({ - integrations: { langchain: false }, - }); - withoutLangChain.enable(); - expect(LangSmithPlugin).toHaveBeenCalledWith({ - skipLangChainRuns: false, - }); - }); - - it("should not create OpenRouter Agent plugin when openrouterAgent: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { openrouterAgent: false }, - }); - plugin.enable(); - - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).not.toHaveBeenCalled(); - }); - - it("should not create any plugins when all are disabled", () => { - const plugin = new BraintrustPlugin({ - integrations: { - openai: false, - openaiCodexSDK: false, - anthropic: false, - aisdk: false, - claudeAgentSDK: false, - cloudflareThink: false, - openAIAgents: false, - googleGenAI: false, - huggingface: false, - openrouter: false, - openrouterAgent: false, - mistral: false, - ollama: false, - cohere: false, - groq: false, - gitHubCopilot: false, - langchain: false, - langsmith: false, - piCodingAgent: false, - strandsAgentSDK: false, - cloudflareAIChat: false, - cloudflareAgents: false, - }, - }); - plugin.enable(); - - expect(OpenAIPlugin).not.toHaveBeenCalled(); - expect(OpenAICodexPlugin).not.toHaveBeenCalled(); - expect(AnthropicPlugin).not.toHaveBeenCalled(); - expect(AISDKPlugin).not.toHaveBeenCalled(); - expect(ClaudeAgentSDKPlugin).not.toHaveBeenCalled(); - expect(CloudflareThinkPlugin).not.toHaveBeenCalled(); - expect(OpenAIAgentsPlugin).not.toHaveBeenCalled(); - expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); - expect(HuggingFacePlugin).not.toHaveBeenCalled(); - expect(OpenRouterPlugin).not.toHaveBeenCalled(); - expect(OpenRouterAgentPlugin).not.toHaveBeenCalled(); - expect(MistralPlugin).not.toHaveBeenCalled(); - expect(OllamaPlugin).not.toHaveBeenCalled(); - expect(CoherePlugin).not.toHaveBeenCalled(); - expect(GroqPlugin).not.toHaveBeenCalled(); - expect(GitHubCopilotPlugin).not.toHaveBeenCalled(); - expect(LangChainPlugin).not.toHaveBeenCalled(); - expect(LangSmithPlugin).not.toHaveBeenCalled(); - expect(PiCodingAgentPlugin).not.toHaveBeenCalled(); - expect(StrandsAgentSDKPlugin).not.toHaveBeenCalled(); - expect(CloudflareAIChatPlugin).not.toHaveBeenCalled(); - expect(CloudflareAgentsPlugin).not.toHaveBeenCalled(); - }); - - it("should not create Pi Coding Agent plugin when piCodingAgent: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { piCodingAgent: false }, - }); - plugin.enable(); - - expect(PiCodingAgentPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Strands Agent SDK plugin when strandsAgentSDK: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { strandsAgentSDK: false }, - }); - plugin.enable(); - - expect(StrandsAgentSDKPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Cloudflare AI Chat plugin when cloudflareAIChat: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { cloudflareAIChat: false }, - }); - plugin.enable(); - - expect(CloudflareAIChatPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Cloudflare Agents plugin when cloudflareAgents: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { cloudflareAgents: false }, - }); - plugin.enable(); - - expect(CloudflareAgentsPlugin).not.toHaveBeenCalled(); - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - }); - - it("should allow selective enabling of plugins", () => { - const plugin = new BraintrustPlugin({ - integrations: { - openai: true, - anthropic: false, - aisdk: false, - claudeAgentSDK: true, - openAIAgents: true, - googleGenAI: false, - huggingface: true, - openrouter: true, - mistral: false, - }, - }); - plugin.enable(); - - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(OpenAIAgentsPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).not.toHaveBeenCalled(); - expect(AISDKPlugin).not.toHaveBeenCalled(); - expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); - expect(MistralPlugin).not.toHaveBeenCalled(); - }); - }); - - describe("backward compatibility (legacy config keys)", () => { - it("should not create AI SDK plugin when vercel: false (legacy)", () => { - const plugin = new BraintrustPlugin({ - integrations: { vercel: false }, - }); - plugin.enable(); - - expect(AISDKPlugin).not.toHaveBeenCalled(); - // Other plugins should still be created - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(OpenAIAgentsPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create Google GenAI plugin when google: false (legacy)", () => { - const plugin = new BraintrustPlugin({ - integrations: { google: false }, - }); - plugin.enable(); - - expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); - // Other plugins should still be created - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - }); - - it("should not create LangChain plugin when langgraph: false (alias)", () => { - const plugin = new BraintrustPlugin({ - integrations: { langgraph: false }, - }); - plugin.enable(); - - expect(LangChainPlugin).not.toHaveBeenCalled(); - expect(LangSmithPlugin).toHaveBeenCalledWith({ - skipLangChainRuns: true, - }); - }); - - it("should not create AI SDK plugin when both aisdk and vercel are false", () => { - const plugin = new BraintrustPlugin({ - integrations: { aisdk: false, vercel: false }, - }); - plugin.enable(); - - expect(AISDKPlugin).not.toHaveBeenCalled(); - }); - - it("should not create Google GenAI plugin when both googleGenAI and google are false", () => { - const plugin = new BraintrustPlugin({ - integrations: { googleGenAI: false, google: false }, - }); - plugin.enable(); - - expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); - }); - - it("should not create AI SDK plugin when aisdk is true but vercel is false", () => { - const plugin = new BraintrustPlugin({ - integrations: { aisdk: true, vercel: false }, - }); - plugin.enable(); - - expect(AISDKPlugin).not.toHaveBeenCalled(); - }); - - it("should not create Google GenAI plugin when googleGenAI is true but google is false", () => { - const plugin = new BraintrustPlugin({ - integrations: { googleGenAI: true, google: false }, - }); - plugin.enable(); - - expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); - }); - - it("should create AI SDK plugin when vercel is true and aisdk is not set", () => { - const plugin = new BraintrustPlugin({ - integrations: { vercel: true }, - }); - plugin.enable(); - - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - }); - - it("should create Google GenAI plugin when google is true and googleGenAI is not set", () => { - const plugin = new BraintrustPlugin({ - integrations: { google: true }, - }); - plugin.enable(); - - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - }); - }); - - describe("enable/disable lifecycle", () => { - it("should enable all sub-plugins when enabled", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - const openaiMock = vi.mocked(OpenAIPlugin).mock.results[0].value; - const openAICodexMock = - vi.mocked(OpenAICodexPlugin).mock.results[0].value; - const anthropicMock = vi.mocked(AnthropicPlugin).mock.results[0].value; - const aiSDKMock = vi.mocked(AISDKPlugin).mock.results[0].value; - const claudeAgentSDKMock = - vi.mocked(ClaudeAgentSDKPlugin).mock.results[0].value; - const openAIAgentsMock = - vi.mocked(OpenAIAgentsPlugin).mock.results[0].value; - const googleGenAIMock = - vi.mocked(GoogleGenAIPlugin).mock.results[0].value; - const huggingFaceMock = - vi.mocked(HuggingFacePlugin).mock.results[0].value; - const huggingFaceTransformersMock = vi.mocked( - HuggingFaceTransformersPlugin, - ).mock.results[0].value; - const openRouterMock = vi.mocked(OpenRouterPlugin).mock.results[0].value; - const openRouterAgentMock = vi.mocked(OpenRouterAgentPlugin).mock - .results[0].value; - const mistralMock = vi.mocked(MistralPlugin).mock.results[0].value; - const ollamaMock = vi.mocked(OllamaPlugin).mock.results[0].value; - const cohereMock = vi.mocked(CoherePlugin).mock.results[0].value; - const groqMock = vi.mocked(GroqPlugin).mock.results[0].value; - const piCodingAgentMock = - vi.mocked(PiCodingAgentPlugin).mock.results[0].value; - const strandsAgentSDKMock = vi.mocked(StrandsAgentSDKPlugin).mock - .results[0].value; - const cloudflareAgentsMock = vi.mocked(CloudflareAgentsPlugin).mock - .results[0].value; - const langChainMock = vi.mocked(LangChainPlugin).mock.results[0].value; - - expect(openaiMock.enable).toHaveBeenCalledTimes(1); - expect(openAICodexMock.enable).toHaveBeenCalledTimes(1); - expect(anthropicMock.enable).toHaveBeenCalledTimes(1); - expect(aiSDKMock.enable).toHaveBeenCalledTimes(1); - expect(claudeAgentSDKMock.enable).toHaveBeenCalledTimes(1); - expect(openAIAgentsMock.enable).toHaveBeenCalledTimes(1); - expect(googleGenAIMock.enable).toHaveBeenCalledTimes(1); - expect(huggingFaceMock.enable).toHaveBeenCalledTimes(1); - expect(huggingFaceTransformersMock.enable).toHaveBeenCalledTimes(1); - expect(openRouterMock.enable).toHaveBeenCalledTimes(1); - expect(openRouterAgentMock.enable).toHaveBeenCalledTimes(1); - expect(mistralMock.enable).toHaveBeenCalledTimes(1); - expect(ollamaMock.enable).toHaveBeenCalledTimes(1); - expect(cohereMock.enable).toHaveBeenCalledTimes(1); - expect(groqMock.enable).toHaveBeenCalledTimes(1); - expect(piCodingAgentMock.enable).toHaveBeenCalledTimes(1); - expect(strandsAgentSDKMock.enable).toHaveBeenCalledTimes(1); - expect(cloudflareAgentsMock.enable).toHaveBeenCalledTimes(1); - expect(langChainMock.enable).toHaveBeenCalledTimes(1); - }); - - it("should disable and nullify all sub-plugins when disabled", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - const openaiMock = vi.mocked(OpenAIPlugin).mock.results[0].value; - const openAICodexMock = - vi.mocked(OpenAICodexPlugin).mock.results[0].value; - const anthropicMock = vi.mocked(AnthropicPlugin).mock.results[0].value; - const aiSDKMock = vi.mocked(AISDKPlugin).mock.results[0].value; - const claudeAgentSDKMock = - vi.mocked(ClaudeAgentSDKPlugin).mock.results[0].value; - const openAIAgentsMock = - vi.mocked(OpenAIAgentsPlugin).mock.results[0].value; - const googleGenAIMock = - vi.mocked(GoogleGenAIPlugin).mock.results[0].value; - const huggingFaceMock = - vi.mocked(HuggingFacePlugin).mock.results[0].value; - const huggingFaceTransformersMock = vi.mocked( - HuggingFaceTransformersPlugin, - ).mock.results[0].value; - const openRouterMock = vi.mocked(OpenRouterPlugin).mock.results[0].value; - const openRouterAgentMock = vi.mocked(OpenRouterAgentPlugin).mock - .results[0].value; - const mistralMock = vi.mocked(MistralPlugin).mock.results[0].value; - const ollamaMock = vi.mocked(OllamaPlugin).mock.results[0].value; - const cohereMock = vi.mocked(CoherePlugin).mock.results[0].value; - const groqMock = vi.mocked(GroqPlugin).mock.results[0].value; - const piCodingAgentMock = - vi.mocked(PiCodingAgentPlugin).mock.results[0].value; - const strandsAgentSDKMock = vi.mocked(StrandsAgentSDKPlugin).mock - .results[0].value; - const cloudflareAgentsMock = vi.mocked(CloudflareAgentsPlugin).mock - .results[0].value; - const langChainMock = vi.mocked(LangChainPlugin).mock.results[0].value; - - plugin.disable(); - - expect(openaiMock.disable).toHaveBeenCalledTimes(1); - expect(openAICodexMock.disable).toHaveBeenCalledTimes(1); - expect(anthropicMock.disable).toHaveBeenCalledTimes(1); - expect(aiSDKMock.disable).toHaveBeenCalledTimes(1); - expect(claudeAgentSDKMock.disable).toHaveBeenCalledTimes(1); - expect(openAIAgentsMock.disable).toHaveBeenCalledTimes(1); - expect(googleGenAIMock.disable).toHaveBeenCalledTimes(1); - expect(huggingFaceMock.disable).toHaveBeenCalledTimes(1); - expect(huggingFaceTransformersMock.disable).toHaveBeenCalledTimes(1); - expect(openRouterMock.disable).toHaveBeenCalledTimes(1); - expect(openRouterAgentMock.disable).toHaveBeenCalledTimes(1); - expect(mistralMock.disable).toHaveBeenCalledTimes(1); - expect(ollamaMock.disable).toHaveBeenCalledTimes(1); - expect(cohereMock.disable).toHaveBeenCalledTimes(1); - expect(groqMock.disable).toHaveBeenCalledTimes(1); - expect(piCodingAgentMock.disable).toHaveBeenCalledTimes(1); - expect(strandsAgentSDKMock.disable).toHaveBeenCalledTimes(1); - expect(cloudflareAgentsMock.disable).toHaveBeenCalledTimes(1); - expect(langChainMock.disable).toHaveBeenCalledTimes(1); - }); - - it("should be idempotent on multiple enable calls", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - plugin.enable(); - plugin.enable(); - - // Should only create and enable plugins once - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - const mockInstance = vi.mocked(OpenAIPlugin).mock.results[0].value; - expect(mockInstance.enable).toHaveBeenCalledTimes(1); - }); - - it("should be idempotent on multiple disable calls", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - - const openaiMock = vi.mocked(OpenAIPlugin).mock.results[0].value; - - plugin.disable(); - plugin.disable(); - plugin.disable(); - - // Should only disable plugins once - expect(openaiMock.disable).toHaveBeenCalledTimes(1); - }); - - it("should not disable plugins if never enabled", () => { - const plugin = new BraintrustPlugin(); - plugin.disable(); - - // Should not create any plugins - expect(OpenAIPlugin).not.toHaveBeenCalled(); - expect(OpenAICodexPlugin).not.toHaveBeenCalled(); - expect(AnthropicPlugin).not.toHaveBeenCalled(); - expect(AISDKPlugin).not.toHaveBeenCalled(); - expect(ClaudeAgentSDKPlugin).not.toHaveBeenCalled(); - expect(OpenAIAgentsPlugin).not.toHaveBeenCalled(); - expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); - expect(HuggingFacePlugin).not.toHaveBeenCalled(); - expect(OpenRouterPlugin).not.toHaveBeenCalled(); - expect(OpenRouterAgentPlugin).not.toHaveBeenCalled(); - expect(MistralPlugin).not.toHaveBeenCalled(); - expect(OllamaPlugin).not.toHaveBeenCalled(); - expect(CoherePlugin).not.toHaveBeenCalled(); - expect(GroqPlugin).not.toHaveBeenCalled(); - expect(PiCodingAgentPlugin).not.toHaveBeenCalled(); - expect(StrandsAgentSDKPlugin).not.toHaveBeenCalled(); - }); - - it("should allow re-enabling after disable", () => { - const plugin = new BraintrustPlugin(); - plugin.enable(); - plugin.disable(); - - vi.clearAllMocks(); - - plugin.enable(); - - expect(OpenAIPlugin).toHaveBeenCalledTimes(1); - expect(OpenAICodexPlugin).toHaveBeenCalledTimes(1); - expect(AnthropicPlugin).toHaveBeenCalledTimes(1); - expect(AISDKPlugin).toHaveBeenCalledTimes(1); - expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(OpenAIAgentsPlugin).toHaveBeenCalledTimes(1); - expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); - expect(HuggingFacePlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); - expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); - expect(MistralPlugin).toHaveBeenCalledTimes(1); - expect(OllamaPlugin).toHaveBeenCalledTimes(1); - expect(CoherePlugin).toHaveBeenCalledTimes(1); - expect(GroqPlugin).toHaveBeenCalledTimes(1); - expect(PiCodingAgentPlugin).toHaveBeenCalledTimes(1); - expect(StrandsAgentSDKPlugin).toHaveBeenCalledTimes(1); - expect(LangChainPlugin).toHaveBeenCalledTimes(1); - expect(LangSmithPlugin).toHaveBeenCalledTimes(1); - }); - - it("should only disable plugins that were enabled", () => { - const plugin = new BraintrustPlugin({ - integrations: { - openai: true, - anthropic: false, - aisdk: true, - claudeAgentSDK: false, - openAIAgents: true, - googleGenAI: true, - huggingface: true, - openrouter: true, - openrouterAgent: true, - mistral: false, - ollama: false, - cohere: false, - groq: true, - langchain: true, - }, - }); - plugin.enable(); - - const openaiMock = vi.mocked(OpenAIPlugin).mock.results[0].value; - const aiSDKMock = vi.mocked(AISDKPlugin).mock.results[0].value; - const openAIAgentsMock = - vi.mocked(OpenAIAgentsPlugin).mock.results[0].value; - const googleGenAIMock = - vi.mocked(GoogleGenAIPlugin).mock.results[0].value; - const huggingFaceMock = - vi.mocked(HuggingFacePlugin).mock.results[0].value; - const openRouterMock = vi.mocked(OpenRouterPlugin).mock.results[0].value; - const openRouterAgentMock = vi.mocked(OpenRouterAgentPlugin).mock - .results[0].value; - const groqMock = vi.mocked(GroqPlugin).mock.results[0].value; - const langChainMock = vi.mocked(LangChainPlugin).mock.results[0].value; - - plugin.disable(); - - expect(openaiMock.disable).toHaveBeenCalledTimes(1); - expect(aiSDKMock.disable).toHaveBeenCalledTimes(1); - expect(openAIAgentsMock.disable).toHaveBeenCalledTimes(1); - expect(googleGenAIMock.disable).toHaveBeenCalledTimes(1); - expect(huggingFaceMock.disable).toHaveBeenCalledTimes(1); - expect(openRouterMock.disable).toHaveBeenCalledTimes(1); - expect(openRouterAgentMock.disable).toHaveBeenCalledTimes(1); - expect(groqMock.disable).toHaveBeenCalledTimes(1); - expect(langChainMock.disable).toHaveBeenCalledTimes(1); - expect(MistralPlugin).not.toHaveBeenCalled(); - expect(OllamaPlugin).not.toHaveBeenCalled(); - expect(CoherePlugin).not.toHaveBeenCalled(); - }); - - it("should not create Voyage AI plugin when voyageai: false", () => { - const plugin = new BraintrustPlugin({ - integrations: { voyageai: false }, - }); - plugin.enable(); - - expect(VoyageAIPlugin).not.toHaveBeenCalled(); - }); - }); -}); - -// Re-exported utility function tests from OpenAI plugin -import { - parseMetricsFromUsage, - processImagesInOutput, - aggregateChatCompletionChunks, -} from "./braintrust-plugin"; -import { Attachment } from "../logger"; - -describe("parseMetricsFromUsage", () => { - it("should return empty object for null usage", () => { - expect(parseMetricsFromUsage(null)).toEqual({}); - }); - - it("should return empty object for undefined usage", () => { - expect(parseMetricsFromUsage(undefined)).toEqual({}); - }); - - it("should parse legacy token names", () => { - const usage = { - prompt_tokens: 10, - completion_tokens: 20, - total_tokens: 30, - }; - - const result = parseMetricsFromUsage(usage); - - expect(result).toEqual({ - prompt_tokens: 10, - completion_tokens: 20, - tokens: 30, // total_tokens is mapped to tokens - }); - }); - - it("should parse new API token names and map them", () => { - const usage = { - input_tokens: 15, - output_tokens: 25, - total_tokens: 40, - }; - - const result = parseMetricsFromUsage(usage); - - expect(result).toEqual({ - prompt_tokens: 15, // input_tokens mapped to prompt_tokens - completion_tokens: 25, // output_tokens mapped to completion_tokens - tokens: 40, // total_tokens mapped to tokens - }); - }); - - it("should parse token details fields", () => { - const usage = { - input_tokens: 100, - input_tokens_details: { - cached_tokens: 50, - audio_tokens: 10, - }, - output_tokens: 80, - output_tokens_details: { - reasoning_tokens: 20, - }, - }; - - const result = parseMetricsFromUsage(usage); - - expect(result).toEqual({ - prompt_tokens: 100, - prompt_cached_tokens: 50, // input -> prompt prefix - prompt_audio_tokens: 10, - completion_tokens: 80, - completion_reasoning_tokens: 20, // output -> completion prefix - }); - }); - - it("should handle mixed legacy and new token names", () => { - const usage = { - prompt_tokens: 10, - input_tokens: 15, - completion_tokens: 20, - total_tokens: 30, - }; - - const result = parseMetricsFromUsage(usage); - - expect(result).toEqual({ - prompt_tokens: 15, // Last one wins (input_tokens mapped) - completion_tokens: 20, - tokens: 30, // total_tokens is mapped to tokens - }); - }); - - it("should ignore non-number token values", () => { - const usage = { - prompt_tokens: "not a number", - completion_tokens: 20, - total_tokens: null, - }; - - const result = parseMetricsFromUsage(usage); - - expect(result).toEqual({ - completion_tokens: 20, - }); - }); - - it("should ignore non-object token details", () => { - const usage = { - input_tokens: 100, - input_tokens_details: "not an object", - output_tokens: 80, - }; - - const result = parseMetricsFromUsage(usage); - - expect(result).toEqual({ - prompt_tokens: 100, - completion_tokens: 80, - }); - }); - - it("should ignore non-number values in token details", () => { - const usage = { - input_tokens: 100, - input_tokens_details: { - cached_tokens: 50, - invalid: "not a number", - also_invalid: null, - }, - }; - - const result = parseMetricsFromUsage(usage); - - expect(result).toEqual({ - prompt_tokens: 100, - prompt_cached_tokens: 50, - }); - }); - - it("should handle empty usage object", () => { - expect(parseMetricsFromUsage({})).toEqual({}); - }); - - it("should handle unknown token prefix in details", () => { - const usage = { - custom_tokens: 100, - custom_tokens_details: { - special: 25, - }, - }; - - const result = parseMetricsFromUsage(usage); - - expect(result).toEqual({ - custom_tokens: 100, - custom_special: 25, // custom prefix preserved - }); - }); -}); - -describe("processImagesInOutput", () => { - it("should convert image_generation_call to attachment", () => { - // Create a small 1x1 red PNG base64 - const base64Image = - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="; - - const output = { - type: "image_generation_call", - result: base64Image, - output_format: "png", - revised_prompt: "A red pixel", - }; - - const result = processImagesInOutput(output); - - expect(result.type).toBe("image_generation_call"); - expect(result.output_format).toBe("png"); - expect(result.revised_prompt).toBe("A red pixel"); - expect(result.result).toBeInstanceOf(Attachment); - - const attachment = result.result as Attachment; - expect(attachment.reference.filename).toContain(".png"); - expect(attachment.reference.content_type).toBe("image/png"); - }); - - it("should handle image_generation_call with long revised_prompt", () => { - const longPrompt = - "This is a very long prompt that should be truncated to 50 characters when used as filename"; - const base64Image = - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="; - - const output = { - type: "image_generation_call", - result: base64Image, - revised_prompt: longPrompt, - }; - - const result = processImagesInOutput(output); - - const attachment = result.result as Attachment; - // Should be truncated to 50 chars + sanitized + extension - expect(attachment.reference.filename.length).toBeLessThanOrEqual(54); // 50 + ".png" - }); - - it("should handle image_generation_call without revised_prompt", () => { - const base64Image = - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="; - - const output = { - type: "image_generation_call", - result: base64Image, - output_format: "jpg", - }; - - const result = processImagesInOutput(output); - - const attachment = result.result as Attachment; - expect(attachment.reference.filename).toBe("generated_image.jpg"); - expect(attachment.reference.content_type).toBe("image/jpg"); - }); - - it("should use png as default format", () => { - const base64Image = - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="; - - const output = { - type: "image_generation_call", - result: base64Image, - }; - - const result = processImagesInOutput(output); - - const attachment = result.result as Attachment; - expect(attachment.reference.filename).toBe("generated_image.png"); - expect(attachment.reference.content_type).toBe("image/png"); - }); - - it("should process arrays recursively", () => { - const base64Image = - "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="; - - const output = [ - { type: "text", content: "Hello" }, - { - type: "image_generation_call", - result: base64Image, - }, - { type: "text", content: "World" }, - ]; - - const result = processImagesInOutput(output); - - expect(Array.isArray(result)).toBe(true); - expect(result[0]).toEqual({ type: "text", content: "Hello" }); - expect(result[1].result).toBeInstanceOf(Attachment); - expect(result[2]).toEqual({ type: "text", content: "World" }); - }); - - it("should pass through non-image objects unchanged", () => { - const output = { - type: "text", - content: "Hello world", - }; - - const result = processImagesInOutput(output); - expect(result).toEqual(output); - }); - - it("should pass through primitive values unchanged", () => { - expect(processImagesInOutput("string")).toBe("string"); - expect(processImagesInOutput(42)).toBe(42); - expect(processImagesInOutput(true)).toBe(true); - expect(processImagesInOutput(null)).toBe(null); - }); - - it("should not process image_generation_call with non-string result", () => { - const output = { - type: "image_generation_call", - result: 12345, // Not a string - }; - - const result = processImagesInOutput(output); - expect(result).toEqual(output); - expect(result.result).toBe(12345); - }); -}); - -describe("aggregateChatCompletionChunks", () => { - it("should aggregate simple text chunks", () => { - const chunks = [ - { - choices: [{ delta: { role: "assistant", content: "Hello" } }], - }, - { - choices: [{ delta: { content: " world" } }], - }, - { - choices: [{ delta: { content: "!" } }], - }, - ]; - - const result = aggregateChatCompletionChunks(chunks as any); - - expect(result.output).toEqual([ - { - index: 0, - message: { - role: "assistant", - content: "Hello world!", - tool_calls: undefined, - }, - logprobs: null, - finish_reason: undefined, - }, - ]); - expect(result.metrics).toEqual({}); - }); - - it("should extract role from first chunk", () => { - const chunks = [ - { - choices: [{ delta: { role: "assistant" } }], - }, - { - choices: [{ delta: { content: "Hi" } }], - }, - ]; - - const result = aggregateChatCompletionChunks(chunks as any); - - expect(result.output[0].message.role).toBe("assistant"); - }); - - it("should extract finish_reason from last chunk with it", () => { - const chunks = [ - { - choices: [{ delta: { role: "assistant", content: "Done" } }], - }, - { - choices: [{ delta: { finish_reason: "stop" } }], - }, - ]; - - const result = aggregateChatCompletionChunks(chunks); - - expect(result.output[0].finish_reason).toBe("stop"); - }); - - it("should aggregate tool calls by id", () => { - const chunks = [ - { - choices: [ - { - delta: { - role: "assistant", - tool_calls: [ - { - id: "call_1", - type: "function", - function: { name: "get_weather", arguments: '{"loc' }, - }, - ], - }, - }, - ], - }, - { - choices: [ - { - delta: { - tool_calls: [ - { - function: { arguments: 'ation":"' }, - }, - ], - }, - }, - ], - }, - { - choices: [ - { - delta: { - tool_calls: [ - { - function: { arguments: 'NYC"}' }, - }, - ], - }, - }, - ], - }, - ]; - - const result = aggregateChatCompletionChunks(chunks); - - expect(result.output[0].message.tool_calls).toEqual([ - { - id: "call_1", - type: "function", - function: { name: "get_weather", arguments: '{"location":"NYC"}' }, - }, - ]); - }); - - it("should handle multiple tool calls", () => { - const chunks = [ - { - choices: [ - { - delta: { - role: "assistant", - tool_calls: [ - { - id: "call_1", - type: "function", - function: { name: "tool1", arguments: '{"a":' }, - }, - ], - }, - }, - ], - }, - { - choices: [ - { - delta: { - tool_calls: [ - { - function: { arguments: "1}" }, - }, - ], - }, - }, - ], - }, - { - choices: [ - { - delta: { - tool_calls: [ - { - id: "call_2", - type: "function", - function: { name: "tool2", arguments: '{"b":' }, - }, - ], - }, - }, - ], - }, - { - choices: [ - { - delta: { - tool_calls: [ - { - function: { arguments: "2}" }, - }, - ], - }, - }, - ], - }, - ]; - - const result = aggregateChatCompletionChunks(chunks as any); - const output = result.output as any[]; - - expect(output[0].message.tool_calls).toHaveLength(2); - expect(output[0].message.tool_calls[0]).toEqual({ - id: "call_1", - type: "function", - function: { name: "tool1", arguments: '{"a":1}' }, - }); - expect(output[0].message.tool_calls[1]).toEqual({ - id: "call_2", - type: "function", - function: { name: "tool2", arguments: '{"b":2}' }, - }); - }); - - it("should parse usage metrics from chunks", () => { - const chunks = [ - { - choices: [{ delta: { role: "assistant", content: "Hi" } }], - }, - { - choices: [{ delta: { content: "!" } }], - usage: { - prompt_tokens: 10, - completion_tokens: 2, - total_tokens: 12, - }, - }, - ]; - - const result = aggregateChatCompletionChunks(chunks as any); - - expect(result.metrics).toEqual({ - prompt_tokens: 10, - completion_tokens: 2, - tokens: 12, // total_tokens is mapped to tokens - }); - }); - - it("should merge usage from multiple chunks", () => { - const chunks = [ - { - choices: [{ delta: { role: "assistant" } }], - usage: { - prompt_tokens: 10, - }, - }, - { - choices: [{ delta: { content: "Hi" } }], - usage: { - completion_tokens: 5, - total_tokens: 15, - }, - }, - ]; - - const result = aggregateChatCompletionChunks(chunks); - - expect(result.metrics).toEqual({ - prompt_tokens: 10, - completion_tokens: 5, - tokens: 15, // total_tokens is mapped to tokens - }); - }); - - it("should handle empty chunks array", () => { - const result = aggregateChatCompletionChunks([]); - - expect(result.output).toEqual([ - { - index: 0, - message: { - role: undefined, - content: undefined, - tool_calls: undefined, - }, - logprobs: null, - finish_reason: undefined, - }, - ]); - expect(result.metrics).toEqual({}); - }); - - it("should handle chunks without choices", () => { - const chunks = [ - {}, - { choices: null }, - { choices: [] }, - { choices: [{ delta: { content: "Hi" } }] }, - ]; - - const result = aggregateChatCompletionChunks(chunks as any); - - expect(result.output[0].message.content).toBe("Hi"); - }); - - it("should handle chunks with only usage", () => { - const chunks = [ - { - usage: { - prompt_tokens: 10, - completion_tokens: 5, - }, - }, - ]; - - const result = aggregateChatCompletionChunks(chunks); - - expect(result.metrics).toEqual({ - prompt_tokens: 10, - completion_tokens: 5, - }); - }); - - it("should handle mixed content and tool calls", () => { - const chunks = [ - { - choices: [ - { - delta: { - role: "assistant", - content: "Let me check", - }, - }, - ], - }, - { - choices: [ - { - delta: { - tool_calls: [ - { - id: "call_1", - type: "function", - function: { name: "check", arguments: "{}" }, - }, - ], - }, - }, - ], - }, - { - choices: [{ delta: { finish_reason: "tool_calls" } }], - }, - ]; - - const result = aggregateChatCompletionChunks(chunks); - - expect(result.output[0].message.content).toBe("Let me check"); - expect(result.output[0].message.tool_calls).toHaveLength(1); - expect(result.output[0].finish_reason).toBe("tool_calls"); - }); -}); diff --git a/js/src/instrumentation/braintrust-plugin.ts b/js/src/instrumentation/braintrust-plugin.ts deleted file mode 100644 index 53bebc00d..000000000 --- a/js/src/instrumentation/braintrust-plugin.ts +++ /dev/null @@ -1,402 +0,0 @@ -import { BasePlugin } from "./core"; -import { OpenAIPlugin } from "./plugins/openai-plugin"; -import { OpenAICodexPlugin } from "./plugins/openai-codex-plugin"; -import { AnthropicPlugin } from "./plugins/anthropic-plugin"; -import { AISDKPlugin } from "./plugins/ai-sdk-plugin"; -import { ClaudeAgentSDKPlugin } from "./plugins/claude-agent-sdk-plugin"; -import { CloudflareThinkPlugin } from "./plugins/cloudflare-think-plugin"; -import { CursorSDKPlugin } from "./plugins/cursor-sdk-plugin"; -import { OpenAIAgentsPlugin } from "./plugins/openai-agents-plugin"; -import { GoogleGenAIPlugin } from "./plugins/google-genai-plugin"; -import { HuggingFacePlugin } from "./plugins/huggingface-plugin"; -import { HuggingFaceTransformersPlugin } from "./plugins/huggingface-transformers-plugin"; -import { OpenRouterAgentPlugin } from "./plugins/openrouter-agent-plugin"; -import { OpenRouterPlugin } from "./plugins/openrouter-plugin"; -import { MistralPlugin } from "./plugins/mistral-plugin"; -import { OllamaPlugin } from "./plugins/ollama-plugin"; -import { GoogleADKPlugin } from "./plugins/google-adk-plugin"; -import { CoherePlugin } from "./plugins/cohere-plugin"; -import { GroqPlugin } from "./plugins/groq-plugin"; -import { BedrockRuntimePlugin } from "./plugins/bedrock-runtime-plugin"; -import { GenkitPlugin } from "./plugins/genkit-plugin"; -import { GitHubCopilotPlugin } from "./plugins/github-copilot-plugin"; -import { LangChainPlugin } from "./plugins/langchain-plugin"; -import { LangSmithPlugin } from "./plugins/langsmith-plugin"; -import { PiCodingAgentPlugin } from "./plugins/pi-coding-agent-plugin"; -import { StrandsAgentSDKPlugin } from "./plugins/strands-agent-sdk-plugin"; -import { VoyageAIPlugin } from "./plugins/voyageai-plugin"; -import { CloudflareAIChatPlugin } from "./plugins/cloudflare-ai-chat-plugin"; -import { CloudflareAgentsPlugin } from "./plugins/cloudflare-agents-plugin"; -import type { InstrumentationIntegrationsConfig } from "./config"; - -export interface BraintrustPluginConfig { - integrations?: InstrumentationIntegrationsConfig; -} - -/** - * Default Braintrust plugin that manages all AI provider instrumentation plugins. - * - * This plugin orchestrates: - * - OpenAI SDK (chat completions, embeddings, etc.) - * - Anthropic SDK (messages) - * - Claude Agent SDK (agent interactions) - * - Vercel AI SDK (generateText, streamText, etc.) - * - Google GenAI SDK - * - HuggingFace Inference SDK - * - LangChain.js and LangGraph - * - Mistral SDK - * - Ollama SDK - * - Cohere SDK - * - Voyage AI SDK - * - * The plugin is automatically enabled when the Braintrust library is loaded. - * Individual integrations can be disabled via configuration. - */ -export class BraintrustPlugin extends BasePlugin { - private config: BraintrustPluginConfig; - private openaiPlugin: OpenAIPlugin | null = null; - private openAICodexPlugin: OpenAICodexPlugin | null = null; - private anthropicPlugin: AnthropicPlugin | null = null; - private aiSDKPlugin: AISDKPlugin | null = null; - private claudeAgentSDKPlugin: ClaudeAgentSDKPlugin | null = null; - private cloudflareThinkPlugin: CloudflareThinkPlugin | null = null; - private cursorSDKPlugin: CursorSDKPlugin | null = null; - private openAIAgentsPlugin: OpenAIAgentsPlugin | null = null; - private googleGenAIPlugin: GoogleGenAIPlugin | null = null; - private huggingFacePlugin: HuggingFacePlugin | null = null; - private huggingFaceTransformersPlugin: HuggingFaceTransformersPlugin | null = - null; - private openRouterPlugin: OpenRouterPlugin | null = null; - private openRouterAgentPlugin: OpenRouterAgentPlugin | null = null; - private mistralPlugin: MistralPlugin | null = null; - private ollamaPlugin: OllamaPlugin | null = null; - private googleADKPlugin: GoogleADKPlugin | null = null; - private coherePlugin: CoherePlugin | null = null; - private groqPlugin: GroqPlugin | null = null; - private bedrockRuntimePlugin: BedrockRuntimePlugin | null = null; - private genkitPlugin: GenkitPlugin | null = null; - private gitHubCopilotPlugin: GitHubCopilotPlugin | null = null; - private langChainPlugin: LangChainPlugin | null = null; - private langSmithPlugin: LangSmithPlugin | null = null; - private piCodingAgentPlugin: PiCodingAgentPlugin | null = null; - private strandsAgentSDKPlugin: StrandsAgentSDKPlugin | null = null; - private voyageAIPlugin: VoyageAIPlugin | null = null; - private cloudflareAIChatPlugin: CloudflareAIChatPlugin | null = null; - private cloudflareAgentsPlugin: CloudflareAgentsPlugin | null = null; - - constructor(config: BraintrustPluginConfig = {}) { - super(); - this.config = config; - } - - protected onEnable(): void { - const integrations = this.config.integrations ?? {}; - - // Enable OpenAI integration (default: true) - if (integrations.openai !== false) { - this.openaiPlugin = new OpenAIPlugin(); - this.openaiPlugin.enable(); - } - - if (integrations.openaiCodexSDK !== false) { - this.openAICodexPlugin = new OpenAICodexPlugin(); - this.openAICodexPlugin.enable(); - } - - // Enable Anthropic integration (default: true) - if (integrations.anthropic !== false) { - this.anthropicPlugin = new AnthropicPlugin(); - this.anthropicPlugin.enable(); - } - - // Enable AI SDK integration (default: true) - // Support both 'aisdk' and legacy 'vercel' config keys - if (integrations.aisdk !== false && integrations.vercel !== false) { - this.aiSDKPlugin = new AISDKPlugin(); - this.aiSDKPlugin.enable(); - } - - // Enable Claude Agent SDK integration (default: true) - if (integrations.claudeAgentSDK !== false) { - this.claudeAgentSDKPlugin = new ClaudeAgentSDKPlugin(); - this.claudeAgentSDKPlugin.enable(); - } - - if (integrations.cloudflareThink !== false) { - this.cloudflareThinkPlugin = new CloudflareThinkPlugin(); - this.cloudflareThinkPlugin.enable(); - } - - if (integrations.cursorSDK !== false && integrations.cursor !== false) { - this.cursorSDKPlugin = new CursorSDKPlugin(); - this.cursorSDKPlugin.enable(); - } - - // Enable OpenAI Agents SDK integration (default: true) - if (integrations.openAIAgents !== false) { - this.openAIAgentsPlugin = new OpenAIAgentsPlugin(); - this.openAIAgentsPlugin.enable(); - } - - // Enable Google GenAI integration (default: true) - // Support both 'googleGenAI' and legacy 'google' config keys - if (integrations.googleGenAI !== false && integrations.google !== false) { - this.googleGenAIPlugin = new GoogleGenAIPlugin(); - this.googleGenAIPlugin.enable(); - } - - if (integrations.huggingface !== false) { - this.huggingFacePlugin = new HuggingFacePlugin(); - this.huggingFacePlugin.enable(); - this.huggingFaceTransformersPlugin = new HuggingFaceTransformersPlugin(); - this.huggingFaceTransformersPlugin.enable(); - } - - if (integrations.openrouter !== false) { - this.openRouterPlugin = new OpenRouterPlugin(); - this.openRouterPlugin.enable(); - } - - if (integrations.openrouterAgent !== false) { - this.openRouterAgentPlugin = new OpenRouterAgentPlugin(); - this.openRouterAgentPlugin.enable(); - } - - if (integrations.mistral !== false) { - this.mistralPlugin = new MistralPlugin(); - this.mistralPlugin.enable(); - } - - if (integrations.ollama !== false) { - this.ollamaPlugin = new OllamaPlugin(); - this.ollamaPlugin.enable(); - } - - // Enable Google ADK integration (default: true) - if (integrations.googleADK !== false) { - this.googleADKPlugin = new GoogleADKPlugin(); - this.googleADKPlugin.enable(); - } - - if (integrations.cohere !== false) { - this.coherePlugin = new CoherePlugin(); - this.coherePlugin.enable(); - } - - if (integrations.voyageai !== false) { - this.voyageAIPlugin = new VoyageAIPlugin(); - this.voyageAIPlugin.enable(); - } - - if (integrations.groq !== false) { - this.groqPlugin = new GroqPlugin(); - this.groqPlugin.enable(); - } - - if ( - integrations.bedrock !== false && - integrations.awsBedrock !== false && - integrations.awsBedrockRuntime !== false - ) { - this.bedrockRuntimePlugin = new BedrockRuntimePlugin(); - this.bedrockRuntimePlugin.enable(); - } - - if (integrations.genkit !== false) { - this.genkitPlugin = new GenkitPlugin(); - this.genkitPlugin.enable(); - } - - if (integrations.gitHubCopilot !== false) { - this.gitHubCopilotPlugin = new GitHubCopilotPlugin(); - this.gitHubCopilotPlugin.enable(); - } - - if (integrations.piCodingAgent !== false) { - this.piCodingAgentPlugin = new PiCodingAgentPlugin(); - this.piCodingAgentPlugin.enable(); - } - - if (integrations.strandsAgentSDK !== false) { - this.strandsAgentSDKPlugin = new StrandsAgentSDKPlugin(); - this.strandsAgentSDKPlugin.enable(); - } - - if (integrations.cloudflareAIChat !== false) { - this.cloudflareAIChatPlugin = new CloudflareAIChatPlugin(); - this.cloudflareAIChatPlugin.enable(); - } - - if (integrations.cloudflareAgents !== false) { - this.cloudflareAgentsPlugin = new CloudflareAgentsPlugin(); - this.cloudflareAgentsPlugin.enable(); - } - - if (integrations.langchain !== false && integrations.langgraph !== false) { - this.langChainPlugin = new LangChainPlugin(); - this.langChainPlugin.enable(); - } - - if (integrations.langsmith !== false) { - this.langSmithPlugin = new LangSmithPlugin({ - skipLangChainRuns: integrations.langchain !== false, - }); - this.langSmithPlugin.enable(); - } - - // Mastra is intentionally not wired here: `@mastra/core` ships its own - // 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 / global hook subscription. - } - - protected onDisable(): void { - if (this.openaiPlugin) { - this.openaiPlugin.disable(); - this.openaiPlugin = null; - } - - if (this.openAICodexPlugin) { - this.openAICodexPlugin.disable(); - this.openAICodexPlugin = null; - } - - if (this.anthropicPlugin) { - this.anthropicPlugin.disable(); - this.anthropicPlugin = null; - } - - if (this.aiSDKPlugin) { - this.aiSDKPlugin.disable(); - this.aiSDKPlugin = null; - } - - if (this.claudeAgentSDKPlugin) { - this.claudeAgentSDKPlugin.disable(); - this.claudeAgentSDKPlugin = null; - } - - if (this.cloudflareThinkPlugin) { - this.cloudflareThinkPlugin.disable(); - this.cloudflareThinkPlugin = null; - } - - if (this.cursorSDKPlugin) { - this.cursorSDKPlugin.disable(); - this.cursorSDKPlugin = null; - } - - if (this.openAIAgentsPlugin) { - this.openAIAgentsPlugin.disable(); - this.openAIAgentsPlugin = null; - } - - if (this.googleGenAIPlugin) { - this.googleGenAIPlugin.disable(); - this.googleGenAIPlugin = null; - } - - if (this.huggingFacePlugin) { - this.huggingFacePlugin.disable(); - this.huggingFacePlugin = null; - } - - if (this.huggingFaceTransformersPlugin) { - this.huggingFaceTransformersPlugin.disable(); - this.huggingFaceTransformersPlugin = null; - } - - if (this.openRouterPlugin) { - this.openRouterPlugin.disable(); - this.openRouterPlugin = null; - } - - if (this.openRouterAgentPlugin) { - this.openRouterAgentPlugin.disable(); - this.openRouterAgentPlugin = null; - } - - if (this.mistralPlugin) { - this.mistralPlugin.disable(); - this.mistralPlugin = null; - } - - if (this.ollamaPlugin) { - this.ollamaPlugin.disable(); - this.ollamaPlugin = null; - } - - if (this.googleADKPlugin) { - this.googleADKPlugin.disable(); - this.googleADKPlugin = null; - } - - if (this.coherePlugin) { - this.coherePlugin.disable(); - this.coherePlugin = null; - } - - if (this.voyageAIPlugin) { - this.voyageAIPlugin.disable(); - this.voyageAIPlugin = null; - } - - if (this.groqPlugin) { - this.groqPlugin.disable(); - this.groqPlugin = null; - } - - if (this.bedrockRuntimePlugin) { - this.bedrockRuntimePlugin.disable(); - this.bedrockRuntimePlugin = null; - } - - if (this.genkitPlugin) { - this.genkitPlugin.disable(); - this.genkitPlugin = null; - } - - if (this.gitHubCopilotPlugin) { - this.gitHubCopilotPlugin.disable(); - this.gitHubCopilotPlugin = null; - } - - if (this.piCodingAgentPlugin) { - this.piCodingAgentPlugin.disable(); - this.piCodingAgentPlugin = null; - } - - if (this.strandsAgentSDKPlugin) { - this.strandsAgentSDKPlugin.disable(); - this.strandsAgentSDKPlugin = null; - } - - if (this.cloudflareAIChatPlugin) { - this.cloudflareAIChatPlugin.disable(); - this.cloudflareAIChatPlugin = null; - } - - if (this.cloudflareAgentsPlugin) { - this.cloudflareAgentsPlugin.disable(); - this.cloudflareAgentsPlugin = null; - } - - if (this.langChainPlugin) { - this.langChainPlugin.disable(); - this.langChainPlugin = null; - } - - if (this.langSmithPlugin) { - this.langSmithPlugin.disable(); - this.langSmithPlugin = null; - } - } -} - -// Re-export utility functions from OpenAIPlugin for backward compatibility -export { - parseMetricsFromUsage, - processImagesInOutput, - aggregateChatCompletionChunks, -} from "./plugins/openai-plugin"; diff --git a/js/src/instrumentation/core/channel-tracing.test.ts b/js/src/instrumentation/core/channel-tracing.test.ts index 55c08f05b..319d24344 100644 --- a/js/src/instrumentation/core/channel-tracing.test.ts +++ b/js/src/instrumentation/core/channel-tracing.test.ts @@ -27,12 +27,32 @@ import { traceAsyncChannel, traceStreamingChannel } from "./channel-tracing"; const testChannels = defineChannels( "channel-tracing-test", { - asyncCall: channel<[Record], { ok: true }>({ - channelName: "async.call", + asyncBinding: channel<[Record], { ok: true }>({ + channelName: "async.binding", kind: "async", }), - streamingCall: channel<[Record], { ok: true }>({ - channelName: "streaming.call", + provenance: channel<[Record], { ok: true }>({ + channelName: "async.provenance", + kind: "async", + }), + skipped: channel<[Record], { ok: true }>({ + channelName: "async.skipped", + kind: "async", + }), + throwingPredicate: channel<[Record], { ok: true }>({ + channelName: "async.throwing-predicate", + kind: "async", + }), + suppressed: channel<[Record], { ok: true }>({ + channelName: "async.suppressed", + kind: "async", + }), + streamingCleanup: channel<[Record], { ok: true }>({ + channelName: "streaming.cleanup", + kind: "async", + }), + streamingCancellation: channel<[Record], { ok: true }>({ + channelName: "streaming.cancellation", kind: "async", }), }, @@ -60,7 +80,7 @@ describe("traceAsyncChannel current span binding", () => { }); it("binds the created span into the traced async execution context", async () => { - const unsubscribe = traceAsyncChannel(testChannels.asyncCall, { + traceAsyncChannel(testChannels.asyncBinding, { name: "channel-tracing-test", type: "function", extractInput: () => ({ @@ -73,21 +93,16 @@ describe("traceAsyncChannel current span binding", () => { const seenSpanIds: string[] = []; - try { - await testChannels.asyncCall.tracePromise( - async () => { - seenSpanIds.push(currentSpan().spanId); - await Promise.resolve(); - seenSpanIds.push(currentSpan().spanId); - - return { ok: true as const }; - }, - { arguments: [{}] } as any, - ); - } finally { - unsubscribe(); - } + await testChannels.asyncBinding.tracePromise( + async () => { + seenSpanIds.push(currentSpan().spanId); + await Promise.resolve(); + seenSpanIds.push(currentSpan().spanId); + return { ok: true as const }; + }, + { arguments: [{}] } as any, + ); expect(seenSpanIds).toHaveLength(2); expect(seenSpanIds[0]).toBeTruthy(); expect(seenSpanIds[1]).toBe(seenSpanIds[0]); @@ -103,7 +118,7 @@ describe("traceAsyncChannel current span binding", () => { }); it("limits channel provenance to directly instrumented spans", async () => { - const unsubscribe = traceAsyncChannel(testChannels.asyncCall, { + traceAsyncChannel(testChannels.provenance, { name: "channel-parent", type: "function", extractInput: () => ({ input: "input", metadata: undefined }), @@ -111,32 +126,27 @@ describe("traceAsyncChannel current span binding", () => { extractMetrics: () => ({}), }); - try { - await testChannels.asyncCall.tracePromise( - async () => { - const parent = currentSpan(); - parent.startSpan({ name: "user-child" }).end(); - parent - .startSpanWithParents("user-multi-parent-child", [parent.spanId], { - name: "user-multi-parent-child", - }) - .end(); - parent - .startSpan( - withSpanInstrumentationName( - { name: "instrumentation-child" }, - INSTRUMENTATION_NAMES.OPENAI, - ), - ) - .end(); - return { ok: true as const }; - }, - { arguments: [{}] } as any, - ); - } finally { - unsubscribe(); - } - + await testChannels.provenance.tracePromise( + async () => { + const parent = currentSpan(); + parent.startSpan({ name: "user-child" }).end(); + parent + .startSpanWithParents("user-multi-parent-child", [parent.spanId], { + name: "user-multi-parent-child", + }) + .end(); + parent + .startSpan( + withSpanInstrumentationName( + { name: "instrumentation-child" }, + INSTRUMENTATION_NAMES.OPENAI, + ), + ) + .end(); + return { ok: true as const }; + }, + { arguments: [{}] } as any, + ); const spans = (await backgroundLogger.drain()) as Array< Record >; @@ -161,7 +171,7 @@ describe("traceAsyncChannel current span binding", () => { }); it("does not create a span when shouldTrace returns false", async () => { - const unsubscribe = traceAsyncChannel(testChannels.asyncCall, { + traceAsyncChannel(testChannels.skipped, { name: "channel-tracing-test", shouldTrace: ([params]) => !( @@ -181,21 +191,16 @@ describe("traceAsyncChannel current span binding", () => { const seenSpanIds: string[] = []; - try { - await testChannels.asyncCall.tracePromise( - async () => { - seenSpanIds.push(currentSpan().spanId); - await Promise.resolve(); - seenSpanIds.push(currentSpan().spanId); - - return { ok: true as const }; - }, - { arguments: [{ skip: true }] } as any, - ); - } finally { - unsubscribe(); - } + await testChannels.skipped.tracePromise( + async () => { + seenSpanIds.push(currentSpan().spanId); + await Promise.resolve(); + seenSpanIds.push(currentSpan().spanId); + return { ok: true as const }; + }, + { arguments: [{ skip: true }] } as any, + ); expect(seenSpanIds).toEqual(["", ""]); expect(currentSpan()).toBe(NOOP_SPAN); @@ -207,7 +212,7 @@ describe("traceAsyncChannel current span binding", () => { const consoleErrorSpy = vi .spyOn(console, "error") .mockImplementation(() => {}); - const unsubscribe = traceAsyncChannel(testChannels.asyncCall, { + traceAsyncChannel(testChannels.throwingPredicate, { name: "channel-tracing-test", shouldTrace: () => { throw new Error("predicate failed"); @@ -221,15 +226,10 @@ describe("traceAsyncChannel current span binding", () => { extractMetrics: () => ({}), }); - try { - await testChannels.asyncCall.tracePromise( - async () => ({ ok: true as const }), - { arguments: [{}] } as any, - ); - } finally { - unsubscribe(); - } - + await testChannels.throwingPredicate.tracePromise( + async () => ({ ok: true as const }), + { arguments: [{}] } as any, + ); expect(consoleErrorSpy).not.toHaveBeenCalled(); consoleErrorSpy.mockRestore(); @@ -238,7 +238,7 @@ describe("traceAsyncChannel current span binding", () => { }); it("skips auto instrumentation spans while suppression is active", async () => { - const unsubscribe = traceAsyncChannel(testChannels.asyncCall, { + traceAsyncChannel(testChannels.suppressed, { name: "channel-tracing-test", type: "function", extractInput: () => ({ @@ -249,23 +249,18 @@ describe("traceAsyncChannel current span binding", () => { extractMetrics: () => ({}), }); - try { - await runWithAutoInstrumentationSuppressed(() => - testChannels.asyncCall.tracePromise( - async () => { - expect(currentSpan()).toBe(NOOP_SPAN); - await Promise.resolve(); - expect(currentSpan()).toBe(NOOP_SPAN); - - return { ok: true as const }; - }, - { arguments: [{}] } as any, - ), - ); - } finally { - unsubscribe(); - } + await runWithAutoInstrumentationSuppressed(() => + testChannels.suppressed.tracePromise( + async () => { + expect(currentSpan()).toBe(NOOP_SPAN); + await Promise.resolve(); + expect(currentSpan()).toBe(NOOP_SPAN); + return { ok: true as const }; + }, + { arguments: [{}] } as any, + ), + ); const spans = await backgroundLogger.drain(); expect(spans).toHaveLength(0); }); @@ -280,7 +275,7 @@ describe("traceAsyncChannel current span binding", () => { throw new Error("logging failed"); }), } as unknown as Span; - const unsubscribe = traceStreamingChannel(testChannels.streamingCall, { + traceStreamingChannel(testChannels.streamingCleanup, { name: "streaming-channel-test", startSpan: () => child, type: "function", @@ -291,25 +286,20 @@ describe("traceAsyncChannel current span binding", () => { onError, }); - try { - await expect( - testChannels.streamingCall.tracePromise( - async () => ({ ok: true as const }), - { arguments: [{}] } as any, - ), - ).resolves.toEqual({ ok: true }); - await expect( - testChannels.streamingCall.tracePromise( - async () => { - throw new Error("call failed"); - }, - { arguments: [{}] } as any, - ), - ).rejects.toThrow("call failed"); - } finally { - unsubscribe(); - } - + await expect( + testChannels.streamingCleanup.tracePromise( + async () => ({ ok: true as const }), + { arguments: [{}] } as any, + ), + ).resolves.toEqual({ ok: true }); + await expect( + testChannels.streamingCleanup.tracePromise( + async () => { + throw new Error("call failed"); + }, + { arguments: [{}] } as any, + ), + ).rejects.toThrow("call failed"); expect(onComplete).toHaveBeenCalledTimes(1); expect(onError).toHaveBeenCalledTimes(1); expect(end).toHaveBeenCalledTimes(2); @@ -321,7 +311,7 @@ describe("traceAsyncChannel current span binding", () => { end: vi.fn(), log: vi.fn(), } as unknown as Span; - const unsubscribe = traceStreamingChannel(testChannels.streamingCall, { + traceStreamingChannel(testChannels.streamingCancellation, { name: "streaming-channel-test", startSpan: () => child, type: "function", @@ -337,17 +327,12 @@ describe("traceAsyncChannel current span binding", () => { }, }; - try { - const patched = await testChannels.streamingCall.tracePromise( - async () => stream as any, - { arguments: [{}] } as any, - ); - (patched as unknown as typeof stream).abort(); - await Promise.resolve(); - } finally { - unsubscribe(); - } - + const patched = await testChannels.streamingCancellation.tracePromise( + async () => stream as any, + { arguments: [{}] } as any, + ); + (patched as unknown as typeof stream).abort(); + await Promise.resolve(); const cancellationError = expect.objectContaining({ message: "Stream cancelled before completion", name: "AbortError", diff --git a/js/src/instrumentation/core/channel-tracing.ts b/js/src/instrumentation/core/channel-tracing.ts index a8ecc819c..1f3ff47d0 100644 --- a/js/src/instrumentation/core/channel-tracing.ts +++ b/js/src/instrumentation/core/channel-tracing.ts @@ -312,7 +312,7 @@ function bindCurrentSpanStoreToStart< }, channelName: string, instrumentationName: SpanInstrumentationName, -): (() => void) | undefined { +): void { const state = _internalGetGlobalState(); const startChannel = tracingChannel.start; const contextManager = state?.contextManager; @@ -325,7 +325,7 @@ function bindCurrentSpanStoreToStart< : undefined; if (!currentSpanStore || !startChannel) { - return undefined; + return; } startChannel.bindStore( @@ -347,10 +347,6 @@ function bindCurrentSpanStoreToStart< : currentSpanStore.getStore(); }, ); - - return () => { - startChannel.unbindStore(currentSpanStore); - }; } function logErrorAndEnd< @@ -442,13 +438,13 @@ function runStreamingErrorHook(args: { export function traceAsyncChannel( channel: TChannel, config: AsyncChannelSpanConfig, -): () => void { +): void { const tracingChannel = channel.tracingChannel() as IsoTracingChannel< ChannelMessage >; const states = new WeakMap(); const channelName = channel.channelName; - const unbindCurrentSpanStore = bindCurrentSpanStoreToStart( + bindCurrentSpanStoreToStart( tracingChannel, states, config, @@ -514,23 +510,18 @@ export function traceAsyncChannel( }; tracingChannel.subscribe(handlers); - - return () => { - unbindCurrentSpanStore?.(); - tracingChannel.unsubscribe(handlers); - }; } export function traceStreamingChannel( channel: TChannel, config: StreamingChannelSpanConfig, -): () => void { +): void { const tracingChannel = channel.tracingChannel() as IsoTracingChannel< ChannelMessage >; const states = new WeakMap(); const channelName = channel.channelName; - const unbindCurrentSpanStore = bindCurrentSpanStoreToStart( + bindCurrentSpanStoreToStart( tracingChannel, states, config, @@ -790,23 +781,18 @@ export function traceStreamingChannel( }; tracingChannel.subscribe(handlers); - - return () => { - unbindCurrentSpanStore?.(); - tracingChannel.unsubscribe(handlers); - }; } export function traceSyncStreamChannel( channel: TChannel, config: SyncStreamChannelSpanConfig, -): () => void { +): void { const tracingChannel = channel.tracingChannel() as IsoTracingChannel< ChannelMessage >; const states = new WeakMap(); const channelName = channel.channelName; - const unbindCurrentSpanStore = bindCurrentSpanStoreToStart( + bindCurrentSpanStoreToStart( tracingChannel, states, config, @@ -939,19 +925,4 @@ export function traceSyncStreamChannel( }; tracingChannel.subscribe(handlers); - - return () => { - unbindCurrentSpanStore?.(); - tracingChannel.unsubscribe(handlers); - }; -} - -export function unsubscribeAll( - unsubscribers: Array<() => void>, -): Array<() => void> { - for (const unsubscribe of unsubscribers) { - unsubscribe(); - } - - return []; } diff --git a/js/src/instrumentation/core/channel.ts b/js/src/instrumentation/core/channel.ts deleted file mode 100644 index 8cbf253d6..000000000 --- a/js/src/instrumentation/core/channel.ts +++ /dev/null @@ -1,55 +0,0 @@ -/** - * Utilities for instrumentation hook naming and management. - */ - -/** - * Channel naming convention: braintrust:{component}:{operation} - * - * Examples: - * - braintrust:openai:chat.completions.create - * - braintrust:anthropic:messages.create - * - braintrust:ai-sdk:generateText - */ - -/** - * Creates a standardized channel name. - * - * @param component - The SDK/library being instrumented (e.g., 'openai', 'anthropic') - * @param operation - The operation being traced (e.g., 'chat.completions.create') - * @returns The full channel name - */ -export function createChannelName( - component: string, - operation: string, -): string { - return `braintrust:${component}:${operation}`; -} - -/** - * Parses a channel name into its component parts. - * - * @param channelName - The full channel name - * @returns Object with component and operation, or null if invalid - */ -export function parseChannelName( - channelName: string, -): { component: string; operation: string } | null { - const match = channelName.match(/^braintrust:([^:]+):(.+)$/); - if (!match) { - return null; - } - return { - component: match[1], - operation: match[2], - }; -} - -/** - * Validates a channel name follows the expected convention. - * - * @param channelName - The channel name to validate - * @returns True if valid - */ -export function isValidChannelName(channelName: string): boolean { - return /^braintrust:[^:]+:.+$/.test(channelName); -} diff --git a/js/src/instrumentation/core/index.ts b/js/src/instrumentation/core/index.ts index 6e2d2d61c..edfe17a89 100644 --- a/js/src/instrumentation/core/index.ts +++ b/js/src/instrumentation/core/index.ts @@ -1,27 +1,3 @@ -/** - * Core utilities for building auto-instrumentation plugins. - * - * Provides BasePlugin class and channel utilities following the OpenTelemetry - * InstrumentationBase pattern - core infrastructure lives here, but individual - * instrumentations can be separate packages. - * - * Note: auto-instrumentation config types are exposed from the Braintrust - * bundler subpaths, such as `braintrust/vite`. - */ +/** Internal utilities shared by Braintrust instrumentation consumers. */ -export { BasePlugin } from "./plugin"; export { toLoggedError } from "./logging"; -export { - createChannelName, - parseChannelName, - isValidChannelName, -} from "./channel"; -export type { - BaseContext, - StartEvent, - EndEvent, - ErrorEvent, - AsyncStartEvent, - AsyncEndEvent, - ChannelHandlers, -} from "./types"; diff --git a/js/src/instrumentation/core/plugin.ts b/js/src/instrumentation/core/plugin.ts deleted file mode 100644 index ae09a43f1..000000000 --- a/js/src/instrumentation/core/plugin.ts +++ /dev/null @@ -1,505 +0,0 @@ -import iso from "../../isomorph"; -import type { IsoChannelHandlers } from "../../isomorph"; -import { isAsyncIterable, patchStreamIfNeeded } from "./stream-patcher"; -import type { StartEvent } from "./types"; -import { startSpan } from "../../logger"; -import type { Span } from "../../logger"; -import { getCurrentUnixTimestamp } from "../../util"; -import { - buildStartSpanArgs, - mergeInputMetadata, -} from "./channel-tracing-utils"; - -/** - * Base class for creating instrumentation plugins. - * - * Plugins subscribe to global instrumentation hook events and convert them - * into spans, logs, or other observability data. - */ -export abstract class BasePlugin { - protected enabled = false; - protected unsubscribers: Array<() => void> = []; - - /** - * Enables the plugin. Must be called before the plugin will receive events. - */ - enable(): void { - if (this.enabled) { - return; - } - this.enabled = true; - this.onEnable(); - } - - /** - * Disables the plugin. After this, the plugin will no longer receive events. - */ - disable(): void { - if (!this.enabled) { - return; - } - this.enabled = false; - this.onDisable(); - } - - /** - * Called when the plugin is enabled. - * Override this to set up subscriptions. - */ - protected abstract onEnable(): void; - - /** - * Called when the plugin is disabled. - * Override this to clean up subscriptions. - */ - protected abstract onDisable(): void; - - /** - * Helper to subscribe to a channel with raw handlers. - * - * @param channelName - The channel name to subscribe to - * @param handlers - Event handlers - */ - protected subscribe(channelName: string, handlers: IsoChannelHandlers): void { - const channel = iso.newTracingChannel(channelName); - channel.subscribe(handlers); - } - - /** - * Subscribe to a channel for async methods (non-streaming). - * Creates a span and logs input/output/metrics. - */ - protected subscribeToChannel( - channelName: string, - config: { - name: string; - type: string; - extractInput: (args: any[]) => { input: any; metadata: any }; - extractOutput: (result: any, endEvent?: any) => any; - extractMetadata?: (result: any, endEvent?: any) => any; - extractMetrics: ( - result: any, - startTime?: number, - endEvent?: any, - ) => Record; - }, - ): void { - const channel = iso.newTracingChannel(channelName); - - const spans = new WeakMap(); - - const handlers = { - start: (event: StartEvent) => { - const { name, spanAttributes, spanInfoMetadata } = buildStartSpanArgs( - config, - event, - ); - const span = startSpan({ - name, - spanAttributes, - }); - - const startTime = getCurrentUnixTimestamp(); - spans.set(event, { span, startTime }); - - try { - const { input, metadata } = config.extractInput(event.arguments); - span.log({ - input, - metadata: mergeInputMetadata(metadata, spanInfoMetadata), - }); - } catch (error) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error(`Error extracting input for ${channelName}:`, error); - } - }, - - asyncEnd: (event: any) => { - const spanData = spans.get(event); - if (!spanData) { - return; - } - - const { span, startTime } = spanData; - - try { - const output = config.extractOutput(event.result, event); - const metrics = config.extractMetrics(event.result, startTime, event); - const metadata = config.extractMetadata?.(event.result, event); - - span.log({ - output, - ...(metadata !== undefined ? { metadata } : {}), - metrics, - }); - } catch (error) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error(`Error extracting output for ${channelName}:`, error); - } finally { - span.end(); - spans.delete(event); - } - }, - - error: (event: any) => { - const spanData = spans.get(event); - if (!spanData) { - return; - } - - const { span } = spanData; - - span.log({ - error: event.error.message, - }); - span.end(); - spans.delete(event); - }, - }; - - channel.subscribe(handlers); - - // Store unsubscribe function - this.unsubscribers.push(() => { - channel.unsubscribe(handlers); - }); - } - - /** - * Subscribe to a channel for async methods that may return streams. - * Handles both streaming and non-streaming responses. - */ - protected subscribeToStreamingChannel( - channelName: string, - config: { - name: string; - type: string; - extractInput: (args: any[]) => { input: any; metadata: any }; - extractOutput: (result: any, endEvent?: any) => any; - extractMetadata?: (result: any, endEvent?: any) => any; - extractMetrics: ( - result: any, - startTime?: number, - endEvent?: any, - ) => Record; - aggregateChunks?: ( - chunks: any[], - result?: any, - endEvent?: any, - ) => { - output: any; - metrics: Record; - metadata?: any; - }; - }, - ): void { - const channel = iso.newTracingChannel(channelName); - - const spans = new WeakMap(); - - const handlers = { - start: (event: StartEvent) => { - const { name, spanAttributes, spanInfoMetadata } = buildStartSpanArgs( - config, - event, - ); - const span = startSpan({ - name, - spanAttributes, - }); - - const startTime = getCurrentUnixTimestamp(); - spans.set(event, { span, startTime }); - - try { - const { input, metadata } = config.extractInput(event.arguments); - span.log({ - input, - metadata: mergeInputMetadata(metadata, spanInfoMetadata), - }); - } catch (error) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error(`Error extracting input for ${channelName}:`, error); - } - }, - - asyncEnd: (event: any) => { - const spanData = spans.get(event); - if (!spanData) { - return; - } - - const { span, startTime } = spanData; - - // Check if result is a stream - if (isAsyncIterable(event.result)) { - let firstChunkTime: number | undefined; - - // Patch the stream to collect chunks - patchStreamIfNeeded(event.result, { - onChunk: () => { - if (firstChunkTime === undefined) { - firstChunkTime = getCurrentUnixTimestamp(); - } - }, - onComplete: (chunks: any[]) => { - try { - let output: any; - let metrics: Record; - let metadata: any; - - if (config.aggregateChunks) { - const aggregated = config.aggregateChunks( - chunks, - event.result, - event, - ); - output = aggregated.output; - metrics = aggregated.metrics; - metadata = aggregated.metadata; - } else { - output = config.extractOutput(chunks, event); - metrics = config.extractMetrics(chunks, startTime, event); - } - - // Add time_to_first_token if not already present - if ( - metrics.time_to_first_token === undefined && - firstChunkTime !== undefined - ) { - metrics.time_to_first_token = firstChunkTime - startTime; - } else if ( - metrics.time_to_first_token === undefined && - chunks.length > 0 - ) { - metrics.time_to_first_token = - getCurrentUnixTimestamp() - startTime; - } - - span.log({ - output, - ...(metadata !== undefined ? { metadata } : {}), - metrics, - }); - } catch (error) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error( - `Error extracting output for ${channelName}:`, - error, - ); - } finally { - span.end(); - } - }, - onError: (error: Error) => { - span.log({ - error: error.message, - }); - span.end(); - }, - }); - - // Don't delete the span from the map yet - it will be ended by the stream - } else { - // Non-streaming response - try { - const output = config.extractOutput(event.result, event); - const metadata = config.extractMetadata - ? config.extractMetadata(event.result, event) - : undefined; - const metrics = config.extractMetrics( - event.result, - startTime, - event, - ); - - span.log({ - output, - ...(metadata !== undefined ? { metadata } : {}), - metrics, - }); - } catch (error) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error(`Error extracting output for ${channelName}:`, error); - } finally { - span.end(); - spans.delete(event); - } - } - }, - - error: (event: any) => { - const spanData = spans.get(event); - if (!spanData) { - return; - } - - const { span } = spanData; - - span.log({ - error: event.error.message, - }); - span.end(); - spans.delete(event); - }, - }; - - channel.subscribe(handlers); - - // Store unsubscribe function - this.unsubscribers.push(() => { - channel.unsubscribe(handlers); - }); - } - - /** - * Subscribe to a channel for sync methods that return event-based streams. - * Used for methods like beta.chat.completions.stream() and responses.stream(). - */ - protected subscribeToSyncStreamChannel( - channelName: string, - config: { - name: string; - type: string; - extractInput: (args: any[]) => { input: any; metadata: any }; - extractFromEvent?: (event: any) => { - output?: any; - metrics?: Record; - metadata?: any; - }; - }, - ): void { - const channel = iso.newTracingChannel(channelName); - - const spans = new WeakMap(); - - const handlers = { - start: (event: StartEvent) => { - const { name, spanAttributes, spanInfoMetadata } = buildStartSpanArgs( - config, - event, - ); - const span = startSpan({ - name, - spanAttributes, - }); - - const startTime = getCurrentUnixTimestamp(); - spans.set(event, { span, startTime }); - - try { - const { input, metadata } = config.extractInput(event.arguments); - span.log({ - input, - metadata: mergeInputMetadata(metadata, spanInfoMetadata), - }); - } catch (error) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error(`Error extracting input for ${channelName}:`, error); - } - }, - - end: (event: any) => { - const spanData = spans.get(event); - if (!spanData) { - return; - } - - const { span, startTime } = spanData; - const stream = event.result; - - if (!stream || typeof stream.on !== "function") { - // Not a stream, just end the span - span.end(); - spans.delete(event); - return; - } - - let first = true; - - // Listen for stream events - stream.on("chunk", (chunk: any) => { - if (first) { - const now = getCurrentUnixTimestamp(); - span.log({ - metrics: { - time_to_first_token: now - startTime, - }, - }); - first = false; - } - }); - - stream.on("chatCompletion", (completion: any) => { - try { - span.log({ - output: completion.choices, - }); - } catch (error) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error( - `Error extracting chatCompletion for ${channelName}:`, - error, - ); - } - }); - - stream.on("event", (streamEvent: any) => { - if (config.extractFromEvent) { - try { - if (first) { - const now = getCurrentUnixTimestamp(); - span.log({ - metrics: { - time_to_first_token: now - startTime, - }, - }); - first = false; - } - - const extracted = config.extractFromEvent(streamEvent); - if (extracted && Object.keys(extracted).length > 0) { - span.log(extracted); - } - } catch (error) { - // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. - console.error( - `Error extracting event for ${channelName}:`, - error, - ); - } - } - }); - - stream.on("end", () => { - span.end(); - spans.delete(event); - }); - - // Don't delete the span from the map - it will be deleted when the stream ends - }, - - error: (event: any) => { - const spanData = spans.get(event); - if (!spanData) { - return; - } - - const { span } = spanData; - - span.log({ - error: event.error.message, - }); - span.end(); - spans.delete(event); - }, - }; - - channel.subscribe(handlers); - - // Store unsubscribe function - this.unsubscribers.push(() => { - channel.unsubscribe(handlers); - }); - } -} diff --git a/js/src/instrumentation/core/types.ts b/js/src/instrumentation/core/types.ts index dc2434dc3..0181f3196 100644 --- a/js/src/instrumentation/core/types.ts +++ b/js/src/instrumentation/core/types.ts @@ -27,7 +27,7 @@ export type SpanInfoCarrier< /** * Base context object shared across all events in a trace. */ -export interface BaseContext { +interface BaseContext { /** * Unique identifier for this trace. * Can be used to correlate start/end/error events. @@ -52,42 +52,6 @@ export interface StartEvent extends BaseContext { arguments: TInput[]; } -/** - * Event emitted after the synchronous portion completes. - * For async functions, this fires when the promise is returned (not settled). - */ -export interface EndEvent extends BaseContext { - /** - * The result of the synchronous portion. - * For async functions, this is the promise (not the resolved value). - */ - result: TResult; - - /** - * Arguments passed to the function (also available in StartEvent). - */ - arguments?: unknown[]; -} - -/** - * Event emitted when a function throws or a promise rejects. - */ -export interface ErrorEvent extends BaseContext { - /** - * The error that was thrown or the rejection reason. - */ - error: Error; - - /** - * Arguments passed to the function (also available in StartEvent). - */ - arguments?: unknown[]; -} - -/** - * Event emitted when a promise begins to settle. - * This fires after the synchronous portion and when the async continuation starts. - */ export interface TypedStartEvent< TArguments extends EventArguments = unknown[], > extends BaseContext { @@ -112,17 +76,6 @@ export interface TypedErrorEvent< arguments?: [...TArguments, ...any[]]; } -// eslint-disable-next-line @typescript-eslint/no-empty-object-type -export interface AsyncStartEvent extends StartEvent {} - -/** - * Event emitted when a promise finishes settling. - * This fires BEFORE control returns to user code after await. - * This is where you should extract output data and finalize spans. - */ -// eslint-disable-next-line @typescript-eslint/no-empty-object-type -export interface AsyncEndEvent extends EndEvent {} - export type StartEventWith< TArguments extends EventArguments = unknown[], TExtra extends object = Record, @@ -144,43 +97,3 @@ export type ErrorEventWith< TArguments extends EventArguments = unknown[], TExtra extends object = Record, > = TypedErrorEvent & TExtra; - -/** - * Subscription handlers for a tracing-compatible global hook. - * - * Common usage pattern: - * - Use start to create spans and extract input - * - Use asyncEnd to extract output and finalize spans - * - Use error to handle failures - */ -export interface ChannelHandlers { - /** - * Called before the synchronous portion of a function executes. - * Use this to create spans and extract input data. - */ - start?: (event: StartEvent) => void; - - /** - * Called after the synchronous portion completes (promise returned). - * Usually not needed for typical instrumentation. - */ - end?: (event: EndEvent) => void; - - /** - * Called when a promise begins to settle. - * Usually not needed for typical instrumentation. - */ - asyncStart?: (event: AsyncStartEvent) => void; - - /** - * Called when a promise finishes settling, before user code continues. - * Use this to extract output, patch streams, and finalize spans. - */ - asyncEnd?: (event: AsyncEndEvent) => void; - - /** - * Called when a function throws or promise rejects. - * Use this to log errors and clean up spans. - */ - error?: (event: ErrorEvent) => void; -} diff --git a/js/src/instrumentation/index.test.ts b/js/src/instrumentation/index.test.ts new file mode 100644 index 000000000..fd5442745 --- /dev/null +++ b/js/src/instrumentation/index.test.ts @@ -0,0 +1,13 @@ +import { describe, expect, it } from "vitest"; +import * as instrumentation from "./index"; + +describe("instrumentation public API", () => { + it("exposes only the curated instrumentation APIs", () => { + expect(Object.keys(instrumentation).sort()).toEqual([ + "OpenAIAgentsTraceProcessor", + "braintrustEveInstrumentation", + "braintrustFlueInstrumentation", + "configureInstrumentation", + ]); + }); +}); diff --git a/js/src/instrumentation/index.ts b/js/src/instrumentation/index.ts index 8bbbd31ec..e148f2dc5 100644 --- a/js/src/instrumentation/index.ts +++ b/js/src/instrumentation/index.ts @@ -1,42 +1,16 @@ /** * Instrumentation APIs for auto-instrumentation. * - * 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 - * can be separate packages. - * * For auto-instrumentation config types, import InstrumentationConfig from the * relevant Braintrust bundler subpath such as `braintrust/vite`. * * @module instrumentation */ -export { BasePlugin } from "./core"; -export { BraintrustPlugin } from "./braintrust-plugin"; -export type { BraintrustPluginConfig } from "./braintrust-plugin"; -export { OpenAIAgentsTraceProcessor } from "./plugins/openai-agents-trace-processor"; -export type { OpenAIAgentsTraceProcessorOptions } from "./plugins/openai-agents-trace-processor"; -export { braintrustFlueInstrumentation } from "./plugins/flue-plugin"; -export { braintrustEveInstrumentation } from "./plugins/eve-plugin"; - -// Re-export core types for external instrumentation packages -export type { - BaseContext, - StartEvent, - EndEvent, - ErrorEvent, - AsyncStartEvent, - AsyncEndEvent, - ChannelHandlers, -} from "./core"; -export { - createChannelName, - parseChannelName, - isValidChannelName, -} from "./core"; +export { OpenAIAgentsTraceProcessor } from "./providers/openai-agents-trace-processor"; +export type { OpenAIAgentsTraceProcessorOptions } from "./providers/openai-agents-trace-processor"; +export { braintrustFlueInstrumentation } from "./providers/flue-instrumentation"; +export { braintrustEveInstrumentation } from "./providers/eve-instrumentation"; // Configuration API export { configureInstrumentation } from "./registry"; diff --git a/js/src/instrumentation/instrumentation-consumers.test.ts b/js/src/instrumentation/instrumentation-consumers.test.ts new file mode 100644 index 000000000..ce4837bb5 --- /dev/null +++ b/js/src/instrumentation/instrumentation-consumers.test.ts @@ -0,0 +1,215 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { InstrumentationIntegrationsConfig } from "./config"; + +const registrations = vi.hoisted(() => ({ + aiSDK: vi.fn(), + anthropic: vi.fn(), + bedrockRuntime: vi.fn(), + claudeAgentSDK: vi.fn(), + cloudflareAgents: vi.fn(), + cloudflareAIChat: vi.fn(), + cloudflareThink: vi.fn(), + cohere: vi.fn(), + cursorSDK: vi.fn(), + genkit: vi.fn(), + gitHubCopilot: vi.fn(), + googleADK: vi.fn(), + googleGenAI: vi.fn(), + groq: vi.fn(), + huggingFace: vi.fn(), + huggingFaceTransformers: vi.fn(), + langChain: vi.fn(), + langSmith: vi.fn(), + mistral: vi.fn(), + ollama: vi.fn(), + openAI: vi.fn(), + openAIAgents: vi.fn(), + openAICodex: vi.fn(), + openRouter: vi.fn(), + openRouterAgent: vi.fn(), + piCodingAgent: vi.fn(), + strandsAgentSDK: vi.fn(), + voyageAI: vi.fn(), +})); + +vi.mock("./providers/ai-sdk-instrumentation", () => ({ + registerAISDKInstrumentation: registrations.aiSDK, +})); +vi.mock("./providers/anthropic-instrumentation", () => ({ + registerAnthropicInstrumentation: registrations.anthropic, +})); +vi.mock("./providers/bedrock-runtime-instrumentation", () => ({ + registerBedrockRuntimeInstrumentation: registrations.bedrockRuntime, +})); +vi.mock("./providers/claude-agent-sdk-instrumentation", () => ({ + registerClaudeAgentSDKInstrumentation: registrations.claudeAgentSDK, +})); +vi.mock("./providers/cloudflare-agents-instrumentation", () => ({ + registerCloudflareAgentsInstrumentation: registrations.cloudflareAgents, +})); +vi.mock("./providers/cloudflare-ai-chat-consumer", () => ({ + registerCloudflareAIChatInstrumentation: registrations.cloudflareAIChat, +})); +vi.mock("./providers/cloudflare-think-instrumentation", () => ({ + registerCloudflareThinkInstrumentation: registrations.cloudflareThink, +})); +vi.mock("./providers/cohere-instrumentation", () => ({ + registerCohereInstrumentation: registrations.cohere, +})); +vi.mock("./providers/cursor-sdk-instrumentation", () => ({ + registerCursorSDKInstrumentation: registrations.cursorSDK, +})); +vi.mock("./providers/genkit-instrumentation", () => ({ + registerGenkitInstrumentation: registrations.genkit, +})); +vi.mock("./providers/github-copilot-instrumentation", () => ({ + registerGitHubCopilotInstrumentation: registrations.gitHubCopilot, +})); +vi.mock("./providers/google-adk-instrumentation", () => ({ + registerGoogleADKInstrumentation: registrations.googleADK, +})); +vi.mock("./providers/google-genai-instrumentation", () => ({ + registerGoogleGenAIInstrumentation: registrations.googleGenAI, +})); +vi.mock("./providers/groq-instrumentation", () => ({ + registerGroqInstrumentation: registrations.groq, +})); +vi.mock("./providers/huggingface-instrumentation", () => ({ + registerHuggingFaceInstrumentation: registrations.huggingFace, +})); +vi.mock("./providers/huggingface-transformers-instrumentation", () => ({ + registerHuggingFaceTransformersInstrumentation: + registrations.huggingFaceTransformers, +})); +vi.mock("./providers/langchain-instrumentation", () => ({ + registerLangChainInstrumentation: registrations.langChain, +})); +vi.mock("./providers/langsmith-instrumentation", () => ({ + registerLangSmithInstrumentation: registrations.langSmith, +})); +vi.mock("./providers/mistral-instrumentation", () => ({ + registerMistralInstrumentation: registrations.mistral, +})); +vi.mock("./providers/ollama-instrumentation", () => ({ + registerOllamaInstrumentation: registrations.ollama, +})); +vi.mock("./providers/openai-instrumentation", () => ({ + registerOpenAIInstrumentation: registrations.openAI, +})); +vi.mock("./providers/openai-agents-instrumentation", () => ({ + registerOpenAIAgentsInstrumentation: registrations.openAIAgents, +})); +vi.mock("./providers/openai-codex-instrumentation", () => ({ + registerOpenAICodexInstrumentation: registrations.openAICodex, +})); +vi.mock("./providers/openrouter-instrumentation", () => ({ + registerOpenRouterInstrumentation: registrations.openRouter, +})); +vi.mock("./providers/openrouter-agent-instrumentation", () => ({ + registerOpenRouterAgentInstrumentation: registrations.openRouterAgent, +})); +vi.mock("./providers/pi-coding-agent-instrumentation", () => ({ + registerPiCodingAgentInstrumentation: registrations.piCodingAgent, +})); +vi.mock("./providers/strands-agent-sdk-instrumentation", () => ({ + registerStrandsAgentSDKInstrumentation: registrations.strandsAgentSDK, +})); +vi.mock("./providers/voyageai-instrumentation", () => ({ + registerVoyageAIInstrumentation: registrations.voyageAI, +})); + +import { registerInstrumentationConsumers } from "./instrumentation-consumers"; + +const integrationCases: Array< + [ + keyof InstrumentationIntegrationsConfig, + Array<(typeof registrations)[keyof typeof registrations]>, + ] +> = [ + ["aisdk", [registrations.aiSDK]], + ["anthropic", [registrations.anthropic]], + ["bedrock", [registrations.bedrockRuntime]], + ["claudeAgentSDK", [registrations.claudeAgentSDK]], + ["cloudflareAgents", [registrations.cloudflareAgents]], + ["cloudflareAIChat", [registrations.cloudflareAIChat]], + ["cloudflareThink", [registrations.cloudflareThink]], + ["cohere", [registrations.cohere]], + ["cursorSDK", [registrations.cursorSDK]], + ["genkit", [registrations.genkit]], + ["gitHubCopilot", [registrations.gitHubCopilot]], + ["googleADK", [registrations.googleADK]], + ["googleGenAI", [registrations.googleGenAI]], + ["groq", [registrations.groq]], + [ + "huggingface", + [registrations.huggingFace, registrations.huggingFaceTransformers], + ], + ["langchain", [registrations.langChain]], + ["langsmith", [registrations.langSmith]], + ["mistral", [registrations.mistral]], + ["ollama", [registrations.ollama]], + ["openai", [registrations.openAI]], + ["openAIAgents", [registrations.openAIAgents]], + ["openaiCodexSDK", [registrations.openAICodex]], + ["openrouter", [registrations.openRouter]], + ["openrouterAgent", [registrations.openRouterAgent]], + ["piCodingAgent", [registrations.piCodingAgent]], + ["strandsAgentSDK", [registrations.strandsAgentSDK]], + ["voyageai", [registrations.voyageAI]], +]; + +describe("registerInstrumentationConsumers", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + it("registers every instrumentation consumer by default", () => { + registerInstrumentationConsumers(); + + for (const register of Object.values(registrations)) { + expect(register).toHaveBeenCalledTimes(1); + } + expect(registrations.langSmith).toHaveBeenCalledWith({ + skipLangChainRuns: true, + }); + }); + + it.each(integrationCases)( + "does not register the %s integration when disabled", + (integration, disabledRegistrations) => { + registerInstrumentationConsumers({ + integrations: { [integration]: false }, + }); + + for (const register of disabledRegistrations) { + expect(register).not.toHaveBeenCalled(); + } + }, + ); + + it.each([ + ["vercel", registrations.aiSDK], + ["google", registrations.googleGenAI], + ["cursor", registrations.cursorSDK], + ["awsBedrock", registrations.bedrockRuntime], + ["awsBedrockRuntime", registrations.bedrockRuntime], + ["langgraph", registrations.langChain], + ] as const)( + "honors the legacy or related %s disable flag", + (integration, register) => { + registerInstrumentationConsumers({ + integrations: { [integration]: false }, + }); + + expect(register).not.toHaveBeenCalled(); + }, + ); + + it("tells LangSmith not to suppress LangChain runs when LangChain is disabled", () => { + registerInstrumentationConsumers({ integrations: { langchain: false } }); + + expect(registrations.langSmith).toHaveBeenCalledWith({ + skipLangChainRuns: false, + }); + }); +}); diff --git a/js/src/instrumentation/instrumentation-consumers.ts b/js/src/instrumentation/instrumentation-consumers.ts new file mode 100644 index 000000000..ed32522a9 --- /dev/null +++ b/js/src/instrumentation/instrumentation-consumers.ts @@ -0,0 +1,159 @@ +import { registerOpenAIInstrumentation } from "./providers/openai-instrumentation"; +import { registerOpenAICodexInstrumentation } from "./providers/openai-codex-instrumentation"; +import { registerAnthropicInstrumentation } from "./providers/anthropic-instrumentation"; +import { registerAISDKInstrumentation } from "./providers/ai-sdk-instrumentation"; +import { registerClaudeAgentSDKInstrumentation } from "./providers/claude-agent-sdk-instrumentation"; +import { registerCloudflareThinkInstrumentation } from "./providers/cloudflare-think-instrumentation"; +import { registerCursorSDKInstrumentation } from "./providers/cursor-sdk-instrumentation"; +import { registerOpenAIAgentsInstrumentation } from "./providers/openai-agents-instrumentation"; +import { registerGoogleGenAIInstrumentation } from "./providers/google-genai-instrumentation"; +import { registerHuggingFaceInstrumentation } from "./providers/huggingface-instrumentation"; +import { registerHuggingFaceTransformersInstrumentation } from "./providers/huggingface-transformers-instrumentation"; +import { registerOpenRouterAgentInstrumentation } from "./providers/openrouter-agent-instrumentation"; +import { registerOpenRouterInstrumentation } from "./providers/openrouter-instrumentation"; +import { registerMistralInstrumentation } from "./providers/mistral-instrumentation"; +import { registerOllamaInstrumentation } from "./providers/ollama-instrumentation"; +import { registerGoogleADKInstrumentation } from "./providers/google-adk-instrumentation"; +import { registerCohereInstrumentation } from "./providers/cohere-instrumentation"; +import { registerGroqInstrumentation } from "./providers/groq-instrumentation"; +import { registerBedrockRuntimeInstrumentation } from "./providers/bedrock-runtime-instrumentation"; +import { registerGenkitInstrumentation } from "./providers/genkit-instrumentation"; +import { registerGitHubCopilotInstrumentation } from "./providers/github-copilot-instrumentation"; +import { registerLangChainInstrumentation } from "./providers/langchain-instrumentation"; +import { registerLangSmithInstrumentation } from "./providers/langsmith-instrumentation"; +import { registerPiCodingAgentInstrumentation } from "./providers/pi-coding-agent-instrumentation"; +import { registerStrandsAgentSDKInstrumentation } from "./providers/strands-agent-sdk-instrumentation"; +import { registerVoyageAIInstrumentation } from "./providers/voyageai-instrumentation"; +import { registerCloudflareAIChatInstrumentation } from "./providers/cloudflare-ai-chat-consumer"; +import { registerCloudflareAgentsInstrumentation } from "./providers/cloudflare-agents-instrumentation"; +import type { InstrumentationConfig } from "./config"; + +/** Registers the configured process-lifetime instrumentation consumers. */ +export function registerInstrumentationConsumers( + config: InstrumentationConfig = {}, +): void { + const integrations = config.integrations ?? {}; + + if (integrations.openai !== false) { + registerOpenAIInstrumentation(); + } + + if (integrations.openaiCodexSDK !== false) { + registerOpenAICodexInstrumentation(); + } + + if (integrations.anthropic !== false) { + registerAnthropicInstrumentation(); + } + + // Support both 'aisdk' and legacy 'vercel' config keys + if (integrations.aisdk !== false && integrations.vercel !== false) { + registerAISDKInstrumentation(); + } + + if (integrations.claudeAgentSDK !== false) { + registerClaudeAgentSDKInstrumentation(); + } + + if (integrations.cloudflareThink !== false) { + registerCloudflareThinkInstrumentation(); + } + + if (integrations.cursorSDK !== false && integrations.cursor !== false) { + registerCursorSDKInstrumentation(); + } + + if (integrations.openAIAgents !== false) { + registerOpenAIAgentsInstrumentation(); + } + + // Support both 'googleGenAI' and legacy 'google' config keys + if (integrations.googleGenAI !== false && integrations.google !== false) { + registerGoogleGenAIInstrumentation(); + } + + if (integrations.huggingface !== false) { + registerHuggingFaceInstrumentation(); + registerHuggingFaceTransformersInstrumentation(); + } + + if (integrations.openrouter !== false) { + registerOpenRouterInstrumentation(); + } + + if (integrations.openrouterAgent !== false) { + registerOpenRouterAgentInstrumentation(); + } + + if (integrations.mistral !== false) { + registerMistralInstrumentation(); + } + + if (integrations.ollama !== false) { + registerOllamaInstrumentation(); + } + + if (integrations.googleADK !== false) { + registerGoogleADKInstrumentation(); + } + + if (integrations.cohere !== false) { + registerCohereInstrumentation(); + } + + if (integrations.voyageai !== false) { + registerVoyageAIInstrumentation(); + } + + if (integrations.groq !== false) { + registerGroqInstrumentation(); + } + + if ( + integrations.bedrock !== false && + integrations.awsBedrock !== false && + integrations.awsBedrockRuntime !== false + ) { + registerBedrockRuntimeInstrumentation(); + } + + if (integrations.genkit !== false) { + registerGenkitInstrumentation(); + } + + if (integrations.gitHubCopilot !== false) { + registerGitHubCopilotInstrumentation(); + } + + if (integrations.piCodingAgent !== false) { + registerPiCodingAgentInstrumentation(); + } + + if (integrations.strandsAgentSDK !== false) { + registerStrandsAgentSDKInstrumentation(); + } + + if (integrations.cloudflareAIChat !== false) { + registerCloudflareAIChatInstrumentation(); + } + + if (integrations.cloudflareAgents !== false) { + registerCloudflareAgentsInstrumentation(); + } + + if (integrations.langchain !== false && integrations.langgraph !== false) { + registerLangChainInstrumentation(); + } + + if (integrations.langsmith !== false) { + registerLangSmithInstrumentation({ + skipLangChainRuns: integrations.langchain !== false, + }); + } + + // Mastra is intentionally not wired here: `@mastra/core` ships its own + // 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 an instrumentation consumer / global hook subscription. +} diff --git a/js/src/instrumentation/plugins/cloudflare-agents-plugin.ts b/js/src/instrumentation/plugins/cloudflare-agents-plugin.ts deleted file mode 100644 index d9813f419..000000000 --- a/js/src/instrumentation/plugins/cloudflare-agents-plugin.ts +++ /dev/null @@ -1,140 +0,0 @@ -import { debugLogger } from "../../debug-logger"; -import type { IsoChannelHandlers } from "../../isomorph"; -import { _internalStartSpanWithContext } from "../../logger"; -import type { Span } from "../../logger"; -import { - INSTRUMENTATION_NAMES, - withSpanInstrumentationName, -} from "../../span-origin"; -import { SpanTypeAttribute } from "../../../util/index"; -import { BasePlugin } from "../core"; -import type { ChannelMessage } from "../core/channel-definitions"; -import { cloudflareAgentsChannels } from "./cloudflare-agents-channels"; - -const CLOUDFLARE_WORKERS_CONTEXT = { - span_origin: { - environment: { type: "server", name: "cloudflare_workers" }, - }, -}; - -export class CloudflareAgentsPlugin extends BasePlugin { - protected onEnable(): void { - const channel = cloudflareAgentsChannels.runAgentTool.tracingChannel(); - const spans = new WeakMap(); - const handlers: IsoChannelHandlers< - ChannelMessage - > = { - start: (event) => { - try { - const agentClass = event.arguments[0]; - const options = event.arguments[1]; - if (ownValue(options, "detached")) { - return; - } - - const name = ownValue(agentClass, "name"); - if (typeof name !== "string" || name.length === 0) { - debugLogger.warn( - "Skipping Cloudflare Agents runAgentTool span because the child agent class has no name.", - ); - return; - } - - const span = _internalStartSpanWithContext( - withSpanInstrumentationName( - { - name, - spanAttributes: { type: SpanTypeAttribute.TOOL }, - event: { - input: ownValue(options, "input"), - }, - }, - INSTRUMENTATION_NAMES.CLOUDFLARE_AGENTS, - ), - CLOUDFLARE_WORKERS_CONTEXT, - ); - spans.set(event, span); - } catch (error) { - logInstrumentationError("start", error); - } - }, - asyncEnd: (event) => { - const span = spans.get(event); - if (!span) { - return; - } - spans.delete(event); - - try { - const status = ownValue(event.result, "status"); - if (status === "completed") { - span.log({ output: ownValue(event.result, "output") }); - } else { - const error = ownValue(event.result, "error"); - if (typeof error === "string") { - span.log({ error }); - } - } - } catch (error) { - logInstrumentationError("completion", error); - } finally { - safelyEndSpan(span); - } - }, - error: (event) => { - const span = spans.get(event); - if (!span) { - return; - } - spans.delete(event); - - try { - span.log({ error: event.error }); - } catch (error) { - logInstrumentationError("rejection", error); - } finally { - safelyEndSpan(span); - } - }, - }; - - channel.subscribe(handlers); - this.unsubscribers.push(() => channel.unsubscribe(handlers)); - } - - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - } -} - -function ownValue(value: unknown, key: PropertyKey): unknown { - if (!isObjectLike(value)) { - return undefined; - } - const descriptor = Object.getOwnPropertyDescriptor(value, key); - return descriptor && "value" in descriptor ? descriptor.value : undefined; -} - -function isObjectLike(value: unknown): value is object { - return ( - (typeof value === "object" && value !== null) || typeof value === "function" - ); -} - -function safelyEndSpan(span: Span): void { - try { - span.end(); - } catch (error) { - logInstrumentationError("span end", error); - } -} - -function logInstrumentationError(operation: string, error: unknown): void { - debugLogger.error( - `Failed to process Cloudflare Agents ${operation} instrumentation:`, - error, - ); -} diff --git a/js/src/instrumentation/plugins/openai-plugin.ts b/js/src/instrumentation/plugins/openai-plugin.ts deleted file mode 100644 index d1ca89dcc..000000000 --- a/js/src/instrumentation/plugins/openai-plugin.ts +++ /dev/null @@ -1,613 +0,0 @@ -import { BasePlugin } from "../core"; -import { - traceAsyncChannel, - traceStreamingChannel, - traceSyncStreamChannel, - unsubscribeAll, -} from "../core/channel-tracing"; -import { Attachment } from "../../logger"; -import { SpanTypeAttribute, isObject } from "../../../util/index"; -import { getCurrentUnixTimestamp } from "../../util"; -import { processInputAttachments } from "../../wrappers/attachment-utils"; -import { openAIChannels } from "./openai-channels"; -import { - BRAINTRUST_CACHED_STREAM_METRIC, - getCachedMetricFromHeaders, - parseMetricsFromUsage, -} from "../../openai-utils"; -import type { - OpenAIChatChoice, - OpenAIChatCompletionChunk, - OpenAIChatLogprobs, - OpenAIResponseStreamEvent, -} from "../../vendor-sdk-types/openai"; - -/** - * Plugin for OpenAI SDK instrumentation. - * - * Handles instrumentation for: - * - Chat completions (streaming and non-streaming) - * - Embeddings - * - Moderations - * - Beta API (parse, stream) - * - Responses API (create, stream, parse, compact) - */ -export class OpenAIPlugin extends BasePlugin { - constructor() { - super(); - } - - protected onEnable(): void { - // Chat Completions - supports streaming - this.unsubscribers.push( - traceStreamingChannel(openAIChannels.chatCompletionsCreate, { - name: "Chat Completion", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const { messages, ...metadata } = params; - return { - input: processInputAttachments(messages), - metadata: { ...metadata, provider: "openai" }, - }; - }, - extractOutput: (result) => { - return result?.choices; - }, - extractMetrics: (result, startTime, endEvent) => { - const metrics = withCachedMetric( - parseMetricsFromUsage(result?.usage), - result, - endEvent, - ); - if (startTime) { - metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; - } - return metrics; - }, - aggregateChunks: aggregateChatCompletionChunks, - }), - ); - - // Embeddings - this.unsubscribers.push( - traceAsyncChannel(openAIChannels.embeddingsCreate, { - name: "Embedding", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const { input, ...metadata } = params; - return { - input, - metadata: { ...metadata, provider: "openai" }, - }; - }, - extractOutput: (result) => { - const embedding = result?.data?.[0]?.embedding; - return Array.isArray(embedding) - ? { embedding_length: embedding.length } - : undefined; - }, - extractMetrics: (result, _startTime, endEvent) => { - return withCachedMetric( - parseMetricsFromUsage(result?.usage), - result, - endEvent, - ); - }, - }), - ); - - // Beta Chat Completions Parse - this.unsubscribers.push( - traceStreamingChannel(openAIChannels.betaChatCompletionsParse, { - name: "Chat Completion", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const { messages, ...metadata } = params; - return { - input: processInputAttachments(messages), - metadata: { ...metadata, provider: "openai" }, - }; - }, - extractOutput: (result) => { - return result?.choices; - }, - extractMetrics: (result, startTime, endEvent) => { - const metrics = withCachedMetric( - parseMetricsFromUsage(result?.usage), - result, - endEvent, - ); - if (startTime) { - metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; - } - return metrics; - }, - aggregateChunks: aggregateChatCompletionChunks, - }), - ); - - // Beta Chat Completions Stream (sync method returning event-based stream) - this.unsubscribers.push( - traceSyncStreamChannel(openAIChannels.betaChatCompletionsStream, { - name: "Chat Completion", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const { messages, ...metadata } = params; - return { - input: processInputAttachments(messages), - metadata: { ...metadata, provider: "openai" }, - }; - }, - }), - ); - - // Moderations - this.unsubscribers.push( - traceAsyncChannel(openAIChannels.moderationsCreate, { - name: "Moderation", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const { input, ...metadata } = params; - return { - input, - metadata: { ...metadata, provider: "openai" }, - }; - }, - extractOutput: (result) => { - return result?.results; - }, - extractMetrics: (result, _startTime, endEvent) => { - return withCachedMetric( - parseMetricsFromUsage(result?.usage), - result, - endEvent, - ); - }, - }), - ); - - // Responses API - create (supports streaming via stream=true param) - this.unsubscribers.push( - traceStreamingChannel(openAIChannels.responsesCreate, { - name: "openai.responses.create", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const { input, ...metadata } = params; - return { - input: processInputAttachments(input), - metadata: { ...metadata, provider: "openai" }, - }; - }, - extractOutput: (result) => { - return processImagesInOutput(result?.output); - }, - extractMetadata: (result) => { - if (!result) { - return undefined; - } - const { output: _output, usage: _usage, ...metadata } = result; - return Object.keys(metadata).length > 0 ? metadata : undefined; - }, - extractMetrics: (result, startTime, endEvent) => { - const metrics = withCachedMetric( - parseMetricsFromUsage(result?.usage), - result, - endEvent, - ); - if (startTime) { - metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; - } - return metrics; - }, - aggregateChunks: aggregateResponseStreamEvents, - }), - ); - - // Responses API - stream (sync method returning event-based stream) - this.unsubscribers.push( - traceSyncStreamChannel(openAIChannels.responsesStream, { - name: "openai.responses.create", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const { input, ...metadata } = params; - return { - input: processInputAttachments(input), - metadata: { ...metadata, provider: "openai" }, - }; - }, - extractFromEvent: (event) => { - if (event.type !== "response.completed" || !event.response) { - return {}; - } - - const response = event.response; - const data: Record = {}; - - if (response.output !== undefined) { - data.output = processImagesInOutput(response.output); - } - - const { usage: _usage, output: _output, ...metadata } = response; - if (Object.keys(metadata).length > 0) { - data.metadata = metadata; - } - - data.metrics = parseMetricsFromUsage(response.usage); - return data; - }, - }), - ); - - // Responses API - parse - this.unsubscribers.push( - traceStreamingChannel(openAIChannels.responsesParse, { - name: "openai.responses.parse", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const { input, ...metadata } = params; - return { - input: processInputAttachments(input), - metadata: { ...metadata, provider: "openai" }, - }; - }, - extractOutput: (result) => { - return processImagesInOutput(result?.output); - }, - extractMetadata: (result) => { - if (!result) { - return undefined; - } - const { output: _output, usage: _usage, ...metadata } = result; - return Object.keys(metadata).length > 0 ? metadata : undefined; - }, - extractMetrics: (result, startTime, endEvent) => { - const metrics = withCachedMetric( - parseMetricsFromUsage(result?.usage), - result, - endEvent, - ); - if (startTime) { - metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; - } - return metrics; - }, - aggregateChunks: aggregateResponseStreamEvents, - }), - ); - - // Responses API - compact - this.unsubscribers.push( - traceAsyncChannel(openAIChannels.responsesCompact, { - name: "openai.responses.compact", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const { input, ...metadata } = params; - return { - input: processInputAttachments(input), - metadata: { ...metadata, provider: "openai" }, - }; - }, - extractOutput: (result) => { - return processImagesInOutput(result?.output); - }, - extractMetadata: (result) => { - if (!result) { - return undefined; - } - const { output: _output, usage: _usage, ...metadata } = result; - return Object.keys(metadata).length > 0 ? metadata : undefined; - }, - extractMetrics: (result, startTime, endEvent) => { - const metrics = withCachedMetric( - parseMetricsFromUsage(result?.usage), - result, - endEvent, - ); - if (startTime) { - metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; - } - return metrics; - }, - }), - ); - } - - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } -} - -function getCachedMetricFromEndEvent(endEvent: unknown): number | undefined { - if (!isObject(endEvent)) { - return undefined; - } - - const response = (endEvent as Record).response; - if (!isObject(response)) { - return undefined; - } - - const headers = (response as { headers?: unknown }).headers; - if (!headers || typeof (headers as Headers).get !== "function") { - return undefined; - } - - return getCachedMetricFromHeaders(headers as Headers); -} - -function withCachedMetric( - metrics: Record, - result: unknown, - endEvent?: unknown, -): Record { - if (metrics.cached !== undefined) { - return metrics; - } - - const cachedFromEvent = getCachedMetricFromEndEvent(endEvent); - if (cachedFromEvent !== undefined) { - return { - ...metrics, - cached: cachedFromEvent, - }; - } - - if (!isObject(result)) { - return metrics; - } - - const cached = (result as Record)[ - BRAINTRUST_CACHED_STREAM_METRIC - ]; - - if (typeof cached !== "number") { - return metrics; - } - - return { - ...metrics, - cached, - }; -} - -/** - * Process output to convert base64 images to attachments. - * Used for Responses API image generation output. - */ -export function processImagesInOutput(output: any): any { - if (Array.isArray(output)) { - return output.map(processImagesInOutput); - } - - if (isObject(output)) { - if ( - output.type === "image_generation_call" && - output.result && - typeof output.result === "string" - ) { - const fileExtension = output.output_format || "png"; - const contentType = `image/${fileExtension}`; - - const baseFilename = - output.revised_prompt && typeof output.revised_prompt === "string" - ? output.revised_prompt.slice(0, 50).replace(/[^a-zA-Z0-9]/g, "_") - : "generated_image"; - const filename = `${baseFilename}.${fileExtension}`; - - // Convert base64 string to Blob - const binaryString = atob(output.result); - const bytes = new Uint8Array(binaryString.length); - for (let i = 0; i < binaryString.length; i++) { - bytes[i] = binaryString.charCodeAt(i); - } - const blob = new Blob([bytes], { type: contentType }); - - const attachment = new Attachment({ - data: blob, - filename: filename, - contentType: contentType, - }); - - return { - ...output, - result: attachment, - }; - } - } - - return output; -} - -function mergeLogprobTokens( - existing: OpenAIChatLogprobs["content"] | OpenAIChatLogprobs["refusal"], - incoming: OpenAIChatLogprobs["content"] | OpenAIChatLogprobs["refusal"], -): OpenAIChatLogprobs["content"] | OpenAIChatLogprobs["refusal"] { - if (incoming === undefined) { - return existing; - } - - if (incoming === null) { - return existing ?? null; - } - - if (Array.isArray(existing)) { - return [...existing, ...incoming]; - } - - return [...incoming]; -} - -function aggregateChatLogprobs( - existing: OpenAIChatLogprobs | null | undefined, - incoming: OpenAIChatLogprobs | null | undefined, -): OpenAIChatLogprobs | null | undefined { - if (incoming === undefined) { - return existing; - } - - if (incoming === null) { - return existing ?? null; - } - - const aggregated: OpenAIChatLogprobs = - existing && existing !== null - ? { ...existing, ...incoming } - : { ...incoming }; - - const content = mergeLogprobTokens(existing?.content, incoming.content); - if (content !== undefined) { - aggregated.content = content; - } - - const refusal = mergeLogprobTokens(existing?.refusal, incoming.refusal); - if (refusal !== undefined) { - aggregated.refusal = refusal; - } - - return aggregated; -} - -/** - * Aggregate chat completion chunks into a single response. - * Combines role (first), content (concatenated), tool_calls (by id), - * finish_reason (last), and usage (last chunk). - */ -export function aggregateChatCompletionChunks( - chunks: OpenAIChatCompletionChunk[], - streamResult?: unknown, - endEvent?: unknown, -): { - output: OpenAIChatChoice[]; - metrics: Record; -} { - let role = undefined; - let content = undefined; - let refusal = undefined; - let tool_calls = undefined; - let logprobs: OpenAIChatLogprobs | null | undefined = undefined; - let finish_reason = undefined; - let metrics: Record = {}; - - for (const chunk of chunks) { - if (chunk.usage) { - metrics = { - ...metrics, - ...parseMetricsFromUsage(chunk.usage), - }; - } - - const choice = chunk.choices?.[0]; - if (!choice) { - continue; - } - - if (choice.finish_reason) { - finish_reason = choice.finish_reason; - } - - logprobs = aggregateChatLogprobs(logprobs, choice.logprobs); - - const delta = choice.delta; - if (!delta) { - continue; - } - - if (delta.finish_reason) { - finish_reason = delta.finish_reason; - } - - if (!role && delta.role) { - role = delta.role; - } - - if (delta.content) { - content = (content || "") + delta.content; - } - - if (delta.refusal) { - refusal = (refusal || "") + delta.refusal; - } - - if (delta.tool_calls) { - const toolDelta = delta.tool_calls[0]; - if ( - !tool_calls || - (toolDelta.id && tool_calls[tool_calls.length - 1].id !== toolDelta.id) - ) { - tool_calls = [ - ...(tool_calls || []), - { - id: toolDelta.id, - type: toolDelta.type, - function: toolDelta.function, - }, - ]; - } else { - tool_calls[tool_calls.length - 1].function.arguments += - toolDelta.function.arguments; - } - } - } - - metrics = withCachedMetric(metrics, streamResult, endEvent); - - return { - metrics, - output: [ - { - index: 0, - message: { - role, - content, - ...(refusal !== undefined ? { refusal } : {}), - tool_calls, - }, - logprobs: logprobs ?? null, - finish_reason, - }, - ], - }; -} - -function aggregateResponseStreamEvents( - chunks: OpenAIResponseStreamEvent[], - _streamResult?: unknown, - endEvent?: unknown, -): { - output: any; - metrics: Record; - metadata?: Record; -} { - let output: any = undefined; - let metrics: Record = {}; - let metadata: Record | undefined = undefined; - - for (const chunk of chunks) { - if (!chunk || !chunk.type || !chunk.response) { - continue; - } - if (chunk.type !== "response.completed") { - continue; - } - - const response = chunk.response; - if (response?.output !== undefined) { - output = processImagesInOutput(response.output); - } - - const { usage: _usage, output: _output, ...rest } = response || {}; - if (Object.keys(rest).length > 0) { - metadata = rest; - } - - metrics = parseMetricsFromUsage(response?.usage); - } - - return { - output, - metrics: withCachedMetric(metrics, undefined, endEvent), - ...(metadata !== undefined ? { metadata } : {}), - }; -} - -export { parseMetricsFromUsage }; diff --git a/js/src/instrumentation/plugins/ai-sdk-channels.ts b/js/src/instrumentation/providers/ai-sdk-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/ai-sdk-channels.ts rename to js/src/instrumentation/providers/ai-sdk-channels.ts diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts b/js/src/instrumentation/providers/ai-sdk-instrumentation.streaming.test.ts similarity index 92% rename from js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts rename to js/src/instrumentation/providers/ai-sdk-instrumentation.streaming.test.ts index 8803708c3..1921dfa02 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.streaming.test.ts +++ b/js/src/instrumentation/providers/ai-sdk-instrumentation.streaming.test.ts @@ -22,7 +22,6 @@ import { } from "../../wrappers/ai-sdk/harness-agent-context"; import { workflowAgentWrapperSpanCountForTesting } from "../../wrappers/ai-sdk/workflow-agent-context"; import { aiSDKChannels } from "./ai-sdk-channels"; -import { AISDKPlugin } from "./ai-sdk-plugin"; try { configureNode(); @@ -41,7 +40,7 @@ describe("AI SDK streaming instrumentation", () => { beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ - projectName: "ai-sdk-plugin.streaming.test.ts", + projectName: "ai-sdk-instrumentation.streaming.test.ts", projectId: "test-project-id", }); }); @@ -890,144 +889,130 @@ describe("AI SDK streaming instrumentation", () => { test("baseStream patch preserves derived stream getters", async () => { expect(await backgroundLogger.drain()).toHaveLength(0); - const plugin = new AISDKPlugin(); - plugin.enable(); - - try { - let chunkSent = false; - const result = (await aiSDKChannels.streamText.tracePromise( - async () => { - const resultRecord = { - baseStream: new ReadableStream({ - pull(controller) { - if (chunkSent) { - controller.close(); - return; - } - - chunkSent = true; - controller.enqueue({ - type: "text-delta", - id: "text-1", - delta: "fresh", - }); - }, - }), - text: Promise.resolve("fresh"), - } as any; - - Object.defineProperty(resultRecord, "textStream", { - configurable: true, - enumerable: true, - get() { - const [textBranch, baseBranch] = this.baseStream.tee(); - this.baseStream = baseBranch; - return textBranch.pipeThrough( - new TransformStream({ - transform(chunk: any, controller) { - if (chunk.type === "text-delta") { - controller.enqueue(chunk.delta); - } - }, - }), - ); - }, - }); + let chunkSent = false; + const result = (await aiSDKChannels.streamText.tracePromise( + async () => { + const resultRecord = { + baseStream: new ReadableStream({ + pull(controller) { + if (chunkSent) { + controller.close(); + return; + } - return resultRecord; - }, - { - arguments: [ - { - model: "mock-stream-model", - prompt: "Reply with fresh.", + chunkSent = true; + controller.enqueue({ + type: "text-delta", + id: "text-1", + delta: "fresh", + }); }, - ], - } as any, - )) as any; + }), + text: Promise.resolve("fresh"), + } as any; + + Object.defineProperty(resultRecord, "textStream", { + configurable: true, + enumerable: true, + get() { + const [textBranch, baseBranch] = this.baseStream.tee(); + this.baseStream = baseBranch; + return textBranch.pipeThrough( + new TransformStream({ + transform(chunk: any, controller) { + if (chunk.type === "text-delta") { + controller.enqueue(chunk.delta); + } + }, + }), + ); + }, + }); - expect( - Object.getOwnPropertyDescriptor(result, "textStream")?.get, - ).toEqual(expect.any(Function)); + return resultRecord; + }, + { + arguments: [ + { + model: "mock-stream-model", + prompt: "Reply with fresh.", + }, + ], + } as any, + )) as any; - let firstText = ""; - for await (const chunk of result.textStream) { - firstText += chunk; - } + expect(Object.getOwnPropertyDescriptor(result, "textStream")?.get).toEqual( + expect.any(Function), + ); - let secondText = ""; - for await (const chunk of result.textStream) { - secondText += chunk; - } + let firstText = ""; + for await (const chunk of result.textStream) { + firstText += chunk; + } - expect(firstText).toBe("fresh"); - expect(secondText).toBe("fresh"); - } finally { - plugin.disable(); + let secondText = ""; + for await (const chunk of result.textStream) { + secondText += chunk; } + + expect(firstText).toBe("fresh"); + expect(secondText).toBe("fresh"); }); test("async iterable stream accessors preserve ReadableStream methods", async () => { expect(await backgroundLogger.drain()).toHaveLength(0); - const plugin = new AISDKPlugin(); - plugin.enable(); - - try { - const result = (await aiSDKChannels.streamText.tracePromise( - async () => { - const resultRecord = { - stream: new ReadableStream({ - start(controller) { - controller.enqueue("v7"); - controller.close(); - }, - }), - text: Promise.resolve("v7"), - } as any; - - Object.defineProperty(resultRecord, "textStream", { - configurable: true, - enumerable: true, - get() { - return this.stream.pipeThrough( - new TransformStream({ - transform(chunk: string, controller) { - controller.enqueue(chunk.toUpperCase()); - }, - }), - ); + const result = (await aiSDKChannels.streamText.tracePromise( + async () => { + const resultRecord = { + stream: new ReadableStream({ + start(controller) { + controller.enqueue("v7"); + controller.close(); }, - }); + }), + text: Promise.resolve("v7"), + } as any; + + Object.defineProperty(resultRecord, "textStream", { + configurable: true, + enumerable: true, + get() { + return this.stream.pipeThrough( + new TransformStream({ + transform(chunk: string, controller) { + controller.enqueue(chunk.toUpperCase()); + }, + }), + ); + }, + }); - return resultRecord; - }, - { - arguments: [ - { - model: "mock-v7-stream-model", - prompt: "Reply with v7.", - }, - ], - } as any, - )) as any; + return resultRecord; + }, + { + arguments: [ + { + model: "mock-v7-stream-model", + prompt: "Reply with v7.", + }, + ], + } as any, + )) as any; - expect(result.stream.pipeThrough).toEqual(expect.any(Function)); - expect(result.stream.getReader).toEqual(expect.any(Function)); + expect(result.stream.pipeThrough).toEqual(expect.any(Function)); + expect(result.stream.getReader).toEqual(expect.any(Function)); - const textStream = result.textStream; - expect(textStream.pipeThrough).toEqual(expect.any(Function)); - expect(textStream.getReader).toEqual(expect.any(Function)); + const textStream = result.textStream; + expect(textStream.pipeThrough).toEqual(expect.any(Function)); + expect(textStream.getReader).toEqual(expect.any(Function)); - const reader = textStream.getReader(); - const first = await reader.read(); - const second = await reader.read(); + const reader = textStream.getReader(); + const first = await reader.read(); + const second = await reader.read(); - expect(first).toEqual({ done: false, value: "V7" }); - expect(second).toEqual({ done: true, value: undefined }); - } finally { - plugin.disable(); - } + expect(first).toEqual({ done: false, value: "V7" }); + expect(second).toEqual({ done: true, value: undefined }); }); test("wrapAgentClass instruments all HarnessAgent turn methods without serializing sessions", async () => { diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts b/js/src/instrumentation/providers/ai-sdk-instrumentation.test.ts similarity index 96% rename from js/src/instrumentation/plugins/ai-sdk-plugin.test.ts rename to js/src/instrumentation/providers/ai-sdk-instrumentation.test.ts index fcc538deb..60fcbc4bf 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.test.ts +++ b/js/src/instrumentation/providers/ai-sdk-instrumentation.test.ts @@ -1,4 +1,4 @@ -import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { describe, it, expect, vi, beforeEach } from "vitest"; const telemetryMocks = vi.hoisted(() => ({ braintrustAISDKTelemetry: vi.fn(), @@ -23,13 +23,13 @@ vi.mock("../../wrappers/ai-sdk/telemetry", () => ({ })); import { - AISDKPlugin, + registerAISDKInstrumentation, DEFAULT_DENY_OUTPUT_PATHS, processAISDKCallInput, processAISDKWorkflowAgentCallInput, processAISDKWorkflowAgentModelCallInput, processAISDKOutput as processAISDKOutputActual, -} from "./ai-sdk-plugin"; +} from "./ai-sdk-instrumentation"; import iso from "../../isomorph"; import { serializeAISDKToolsForLogging } from "../../wrappers/ai-sdk/tool-serialization"; import { BRAINTRUST_AI_SDK_V7_OPERATION_KEY as AI_SDK_V7_OPERATION_KEY } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; @@ -39,7 +39,6 @@ type MockTracingChannel = { handlers: any[]; hasSubscribers: boolean; subscribe: ReturnType; - unsubscribe: ReturnType; }; const mockChannels = new Map(); @@ -47,9 +46,7 @@ const mockChannels = new Map(); // Since these are private, we'll test them through the public API // But we'll also add some tests for the exported utility functions -describe("AISDKPlugin", () => { - let plugin: AISDKPlugin; - +describe("registerAISDKInstrumentation", () => { beforeEach(() => { mockChannels.clear(); telemetryMocks.telemetry = { @@ -70,39 +67,10 @@ describe("AISDKPlugin", () => { channel.handlers.push(handlers); channel.hasSubscribers = true; }), - unsubscribe: vi.fn((handlers: any) => { - channel.handlers = channel.handlers.filter( - (candidate) => candidate !== handlers, - ); - channel.hasSubscribers = channel.handlers.length > 0; - return true; - }), }; mockChannels.set(name, channel); return channel; }); - - plugin = new AISDKPlugin(); - }); - - afterEach(() => { - if (plugin) { - plugin.disable(); - } - }); - - describe("constructor", () => { - it("should create plugin with default config", () => { - const defaultPlugin = new AISDKPlugin(); - expect(defaultPlugin).toBeInstanceOf(AISDKPlugin); - }); - - it("should create plugin with custom config", () => { - const customPlugin = new AISDKPlugin({ - denyOutputPaths: ["custom.path"], - }); - expect(customPlugin).toBeInstanceOf(AISDKPlugin); - }); }); describe("WorkflowAgent input extraction", () => { @@ -192,33 +160,9 @@ describe("AISDKPlugin", () => { }); }); - describe("enable/disable", () => { - it("should enable plugin", () => { - expect(() => plugin.enable()).not.toThrow(); - }); - - it("should disable plugin", () => { - plugin.enable(); - expect(() => plugin.disable()).not.toThrow(); - }); - - it("should handle multiple enable calls", () => { - plugin.enable(); - expect(() => plugin.enable()).not.toThrow(); - }); - - it("should handle multiple disable calls", () => { - plugin.enable(); - plugin.disable(); - expect(() => plugin.disable()).not.toThrow(); - }); - - it("should unsubscribe from channels on disable", () => { - plugin.enable(); - plugin.disable(); - // Verify that unsubscribers were called - // This is tested indirectly - if it doesn't throw, unsubscribe worked - expect(true).toBe(true); + describe("registration", () => { + it("registers without throwing", () => { + expect(() => registerAISDKInstrumentation()).not.toThrow(); }); }); @@ -238,7 +182,7 @@ describe("AISDKPlugin", () => { onStepEnd: existingOnStepEnd, }; - plugin.enable(); + registerAISDKInstrumentation(); const channel = mockChannels.get( "orchestrion:ai:createTelemetryDispatcher", @@ -309,7 +253,7 @@ describe("AISDKPlugin", () => { onStart: existingOnStart, }; - plugin.enable(); + registerAISDKInstrumentation(); const channel = mockChannels.get( "orchestrion:ai:createTelemetryDispatcher", @@ -368,7 +312,7 @@ describe("AISDKPlugin", () => { onStart: vi.fn(), }; - plugin.enable(); + registerAISDKInstrumentation(); const channel = mockChannels.get( "orchestrion:ai:createTelemetryDispatcher", @@ -432,7 +376,7 @@ describe("AISDKPlugin", () => { onStart: existingOnStart, }; - plugin.enable(); + registerAISDKInstrumentation(); const channel = mockChannels.get( "orchestrion:ai:createTelemetryDispatcher", @@ -460,7 +404,7 @@ describe("AISDKPlugin", () => { onStart: vi.fn(), }; - plugin.enable(); + registerAISDKInstrumentation(); const channel = mockChannels.get( "orchestrion:ai:createTelemetryDispatcher", diff --git a/js/src/instrumentation/plugins/ai-sdk-plugin.ts b/js/src/instrumentation/providers/ai-sdk-instrumentation.ts similarity index 85% rename from js/src/instrumentation/plugins/ai-sdk-plugin.ts rename to js/src/instrumentation/providers/ai-sdk-instrumentation.ts index e6cbf92f4..8b797e987 100644 --- a/js/src/instrumentation/plugins/ai-sdk-plugin.ts +++ b/js/src/instrumentation/providers/ai-sdk-instrumentation.ts @@ -1,10 +1,9 @@ -import { BasePlugin, toLoggedError } from "../core"; +import { toLoggedError } from "../core"; import { debugLogger } from "../../debug-logger"; import { traceAsyncChannel, traceStreamingChannel, traceSyncStreamChannel, - unsubscribeAll, } from "../core/channel-tracing"; import type { ChannelMessage } from "../core/channel-definitions"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; @@ -75,7 +74,7 @@ import type { } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; import { BRAINTRUST_AI_SDK_V7_OPERATION_KEY as AI_SDK_V7_OPERATION_KEY } from "../../vendor-sdk-types/ai-sdk-v7-telemetry"; -interface AISDKPluginConfig { +interface AISDKInstrumentationConfig { /** * List of JSON paths to remove from output field. * Uses dot notation with array wildcards: "roundtrips[].request.body" @@ -150,10 +149,10 @@ const AI_SDK_V7_TELEMETRY_CALLBACKS = [ ] as const; /** - * AI SDK plugin that subscribes to instrumentation channels + * Internal AI SDK consumer that subscribes to instrumentation channels * and creates Braintrust spans. * - * This plugin handles: + * This consumer handles: * - generateText (async function) * - streamText (function returning stream) * - generateObject (async function) @@ -167,500 +166,455 @@ const AI_SDK_V7_TELEMETRY_CALLBACKS = [ * - ToolLoopAgent.stream (async method returning stream) * - WorkflowAgent.stream (async method returning stream) * - * The plugin automatically extracts: + * The consumer automatically extracts: * - Model and provider information * - Token usage metrics * - Tool calls and structured outputs * - Streaming responses with time-to-first-token */ -export class AISDKPlugin extends BasePlugin { - private config: AISDKPluginConfig; +class AISDKInstrumentationConsumer { + private config: AISDKInstrumentationConfig; - constructor(config: AISDKPluginConfig = {}) { - super(); + constructor(config: AISDKInstrumentationConfig = {}) { this.config = config; } - protected onEnable(): void { + public register(): void { this.subscribeToAISDK(); } - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } - private subscribeToAISDK(): void { const denyOutputPaths = this.config.denyOutputPaths || DEFAULT_DENY_OUTPUT_PATHS; - this.unsubscribers.push(subscribeToAISDKV7TelemetryDispatcher()); - this.unsubscribers.push(subscribeToHarnessAgentCreateSession()); - this.unsubscribers.push( - subscribeToHarnessContinuation( - harnessAgentChannels.continueGenerate, - denyOutputPaths, - ), - subscribeToHarnessContinuation( - harnessAgentChannels.continueStream, - denyOutputPaths, - ), + subscribeToAISDKV7TelemetryDispatcher(); + subscribeToHarnessAgentCreateSession(); + subscribeToHarnessContinuation( + harnessAgentChannels.continueGenerate, + denyOutputPaths, + ); + subscribeToHarnessContinuation( + harnessAgentChannels.continueStream, + denyOutputPaths, ); // generateText - async function that may return streams - this.unsubscribers.push( - traceStreamingChannel(aiSDKChannels.generateText, { - name: "generateText", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event, span) => - prepareAISDKCallInput(params, event, span, denyOutputPaths), - extractOutput: (result, endEvent) => { - finalizeAISDKChildTracing(endEvent as { [key: string]: unknown }); - return processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ); - }, - extractMetrics: (result, _startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent), - aggregateChunks: aggregateAISDKChunks, - }), - ); + traceStreamingChannel(aiSDKChannels.generateText, { + name: "generateText", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event, span) => + prepareAISDKCallInput(params, event, span, denyOutputPaths), + extractOutput: (result, endEvent) => { + finalizeAISDKChildTracing(endEvent as { [key: string]: unknown }); + return processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ); + }, + extractMetrics: (result, _startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent), + aggregateChunks: aggregateAISDKChunks, + }); // streamText - function returning stream - this.unsubscribers.push( - traceStreamingChannel(aiSDKChannels.streamText, { - name: "streamText", - type: SpanTypeAttribute.FUNCTION, - shouldTrace: () => currentCloudflareThinkSpan() === undefined, - extractInput: ([params], event, span) => - prepareAISDKCallInput(params, event, span, denyOutputPaths), - extractOutput: (result, endEvent) => - processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ), - extractMetrics: (result, startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent, startTime), - aggregateChunks: aggregateAISDKChunks, - patchResult: ({ endEvent, result, span, startTime }) => - patchAISDKStreamingResult({ - defaultDenyOutputPaths: denyOutputPaths, - endEvent, - result, - span, - startTime, - }), - }), - ); + traceStreamingChannel(aiSDKChannels.streamText, { + name: "streamText", + type: SpanTypeAttribute.FUNCTION, + shouldTrace: () => currentCloudflareThinkSpan() === undefined, + extractInput: ([params], event, span) => + prepareAISDKCallInput(params, event, span, denyOutputPaths), + extractOutput: (result, endEvent) => + processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ), + extractMetrics: (result, startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent, startTime), + aggregateChunks: aggregateAISDKChunks, + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + span, + startTime, + }), + }); // streamText - sync function returning stream (v4+, used by auto-hook) - this.unsubscribers.push( - traceSyncStreamChannel(aiSDKChannels.streamTextSync, { - name: "streamText", - type: SpanTypeAttribute.FUNCTION, - shouldTrace: () => currentCloudflareThinkSpan() === undefined, - extractInput: ([params], event, span) => - prepareAISDKCallInput(params, event, span, denyOutputPaths), - patchResult: ({ endEvent, result, span, startTime }) => - patchAISDKStreamingResult({ - defaultDenyOutputPaths: denyOutputPaths, - endEvent, - result, - span, - startTime, - }), - }), - ); + traceSyncStreamChannel(aiSDKChannels.streamTextSync, { + name: "streamText", + type: SpanTypeAttribute.FUNCTION, + shouldTrace: () => currentCloudflareThinkSpan() === undefined, + extractInput: ([params], event, span) => + prepareAISDKCallInput(params, event, span, denyOutputPaths), + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + span, + startTime, + }), + }); // generateObject - async function that may return streams - this.unsubscribers.push( - traceStreamingChannel(aiSDKChannels.generateObject, { - name: "generateObject", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event, span) => - prepareAISDKCallInput(params, event, span, denyOutputPaths), - extractOutput: (result, endEvent) => { - finalizeAISDKChildTracing(endEvent as { [key: string]: unknown }); - return processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ); - }, - extractMetrics: (result, _startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent), - aggregateChunks: aggregateAISDKChunks, - }), - ); + traceStreamingChannel(aiSDKChannels.generateObject, { + name: "generateObject", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event, span) => + prepareAISDKCallInput(params, event, span, denyOutputPaths), + extractOutput: (result, endEvent) => { + finalizeAISDKChildTracing(endEvent as { [key: string]: unknown }); + return processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ); + }, + extractMetrics: (result, _startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent), + aggregateChunks: aggregateAISDKChunks, + }); // streamObject - function returning stream - this.unsubscribers.push( - traceStreamingChannel(aiSDKChannels.streamObject, { - name: "streamObject", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event, span) => - prepareAISDKCallInput(params, event, span, denyOutputPaths), - extractOutput: (result, endEvent) => - processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ), - extractMetrics: (result, startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent, startTime), - aggregateChunks: aggregateAISDKChunks, - patchResult: ({ endEvent, result, span, startTime }) => - patchAISDKStreamingResult({ - defaultDenyOutputPaths: denyOutputPaths, - endEvent, - result, - span, - startTime, - }), - }), - ); + traceStreamingChannel(aiSDKChannels.streamObject, { + name: "streamObject", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event, span) => + prepareAISDKCallInput(params, event, span, denyOutputPaths), + extractOutput: (result, endEvent) => + processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ), + extractMetrics: (result, startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent, startTime), + aggregateChunks: aggregateAISDKChunks, + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + span, + startTime, + }), + }); // streamObject - sync function returning stream (v4+, used by auto-hook) - this.unsubscribers.push( - traceSyncStreamChannel(aiSDKChannels.streamObjectSync, { - name: "streamObject", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event, span) => - prepareAISDKCallInput(params, event, span, denyOutputPaths), - patchResult: ({ endEvent, result, span, startTime }) => - patchAISDKStreamingResult({ - defaultDenyOutputPaths: denyOutputPaths, - endEvent, - result, - span, - startTime, - }), - }), - ); + traceSyncStreamChannel(aiSDKChannels.streamObjectSync, { + name: "streamObject", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event, span) => + prepareAISDKCallInput(params, event, span, denyOutputPaths), + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + span, + startTime, + }), + }); // embed - async embedding function - this.unsubscribers.push( - traceAsyncChannel(aiSDKChannels.embed, { - name: "embed", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event) => - prepareAISDKEmbedInput(params, event.self), - extractOutput: (result, endEvent) => - processAISDKEmbeddingOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ), - extractMetrics: (result, _startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent), - }), - ); + traceAsyncChannel(aiSDKChannels.embed, { + name: "embed", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event) => + prepareAISDKEmbedInput(params, event.self), + extractOutput: (result, endEvent) => + processAISDKEmbeddingOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ), + extractMetrics: (result, _startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent), + }); // embedMany - async embedding batch function - this.unsubscribers.push( - traceAsyncChannel(aiSDKChannels.embedMany, { - name: "embedMany", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event) => - prepareAISDKEmbedInput(params, event.self), - extractOutput: (result, endEvent) => - processAISDKEmbeddingOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ), - extractMetrics: (result, _startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent), - }), - ); + traceAsyncChannel(aiSDKChannels.embedMany, { + name: "embedMany", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event) => + prepareAISDKEmbedInput(params, event.self), + extractOutput: (result, endEvent) => + processAISDKEmbeddingOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ), + extractMetrics: (result, _startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent), + }); // rerank - async reranking function - this.unsubscribers.push( - traceAsyncChannel(aiSDKChannels.rerank, { - name: "rerank", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event) => - prepareAISDKRerankInput(params, event.self), - extractOutput: (result, endEvent) => - processAISDKRerankOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ), - extractMetrics: (result, _startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent), - }), - ); + traceAsyncChannel(aiSDKChannels.rerank, { + name: "rerank", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event) => + prepareAISDKRerankInput(params, event.self), + extractOutput: (result, endEvent) => + processAISDKRerankOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ), + extractMetrics: (result, _startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent), + }); // Agent.generate - async method - this.unsubscribers.push( - traceStreamingChannel(aiSDKChannels.agentGenerate, { - name: "Agent.generate", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event, span) => - prepareAISDKCallInput(params, event, span, denyOutputPaths, { - agentOwner: true, - }), - extractOutput: (result, endEvent) => { - finalizeAISDKChildTracing(endEvent as { [key: string]: unknown }); - return processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ); - }, - extractMetrics: (result, _startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent), - aggregateChunks: aggregateAISDKChunks, - }), - ); + traceStreamingChannel(aiSDKChannels.agentGenerate, { + name: "Agent.generate", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event, span) => + prepareAISDKCallInput(params, event, span, denyOutputPaths, { + agentOwner: true, + }), + extractOutput: (result, endEvent) => { + finalizeAISDKChildTracing(endEvent as { [key: string]: unknown }); + return processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ); + }, + extractMetrics: (result, _startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent), + aggregateChunks: aggregateAISDKChunks, + }); // Agent.stream - async method returning stream (v5, used by wrapAISDK) - this.unsubscribers.push( - traceStreamingChannel(aiSDKChannels.agentStream, { - name: "Agent.stream", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event, span) => - prepareAISDKCallInput(params, event, span, denyOutputPaths, { - agentOwner: true, - }), - extractOutput: (result, endEvent) => - processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ), - extractMetrics: (result, startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent, startTime), - aggregateChunks: aggregateAISDKChunks, - patchResult: ({ endEvent, result, span, startTime }) => - patchAISDKStreamingResult({ - defaultDenyOutputPaths: denyOutputPaths, - endEvent, - result, - span, - startTime, - }), - }), - ); + traceStreamingChannel(aiSDKChannels.agentStream, { + name: "Agent.stream", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event, span) => + prepareAISDKCallInput(params, event, span, denyOutputPaths, { + agentOwner: true, + }), + extractOutput: (result, endEvent) => + processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ), + extractMetrics: (result, startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent, startTime), + aggregateChunks: aggregateAISDKChunks, + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + span, + startTime, + }), + }); // Agent.stream - sync method returning stream (v5, used by auto-hook) - this.unsubscribers.push( - traceSyncStreamChannel(aiSDKChannels.agentStreamSync, { - name: "Agent.stream", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event, span) => - prepareAISDKCallInput(params, event, span, denyOutputPaths, { - agentOwner: true, - }), - patchResult: ({ endEvent, result, span, startTime }) => - patchAISDKStreamingResult({ - defaultDenyOutputPaths: denyOutputPaths, - endEvent, - result, - span, - startTime, - }), - }), - ); + traceSyncStreamChannel(aiSDKChannels.agentStreamSync, { + name: "Agent.stream", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event, span) => + prepareAISDKCallInput(params, event, span, denyOutputPaths, { + agentOwner: true, + }), + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + span, + startTime, + }), + }); // HarnessAgent.generate - one task span per agent turn - this.unsubscribers.push( - traceStreamingChannel(harnessAgentChannels.generate, { - name: "HarnessAgent.generate", - startSpan: _internalStartSpanWithInitialMerge, - type: SpanTypeAttribute.TASK, - extractInput: ([params], event, span) => - prepareAISDKHarnessAgentInput(params, event.self, span), - extractOutput: (result, endEvent) => - processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ), - extractMetrics: (result) => extractTokenMetrics(result), - aggregateChunks: aggregateAISDKChunks, - }), - ); + traceStreamingChannel(harnessAgentChannels.generate, { + name: "HarnessAgent.generate", + startSpan: _internalStartSpanWithInitialMerge, + type: SpanTypeAttribute.TASK, + extractInput: ([params], event, span) => + prepareAISDKHarnessAgentInput(params, event.self, span), + extractOutput: (result, endEvent) => + processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ), + extractMetrics: (result) => extractTokenMetrics(result), + aggregateChunks: aggregateAISDKChunks, + }); // HarnessAgent.stream - async method returning an AI SDK stream result - this.unsubscribers.push( - traceStreamingChannel(harnessAgentChannels.stream, { - name: "HarnessAgent.stream", - startSpan: _internalStartSpanWithInitialMerge, - type: SpanTypeAttribute.TASK, - extractInput: ([params], event, span) => - prepareAISDKHarnessAgentInput(params, event.self, span), - extractOutput: (result, endEvent) => - processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ), - extractMetrics: (result, startTime) => ({ - ...extractTokenMetrics(result), - ...(startTime === undefined - ? {} - : { - time_to_first_token: getCurrentUnixTimestamp() - startTime, - }), - }), - aggregateChunks: aggregateAISDKChunks, - patchResult: ({ endEvent, result, span, startTime }) => - patchAISDKStreamingResult({ - defaultDenyOutputPaths: denyOutputPaths, - endEvent, - result, - resolvePromiseUsage: true, - span, - startTime, - }), + traceStreamingChannel(harnessAgentChannels.stream, { + name: "HarnessAgent.stream", + startSpan: _internalStartSpanWithInitialMerge, + type: SpanTypeAttribute.TASK, + extractInput: ([params], event, span) => + prepareAISDKHarnessAgentInput(params, event.self, span), + extractOutput: (result, endEvent) => + processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ), + extractMetrics: (result, startTime) => ({ + ...extractTokenMetrics(result), + ...(startTime === undefined + ? {} + : { + time_to_first_token: getCurrentUnixTimestamp() - startTime, + }), }), - ); + aggregateChunks: aggregateAISDKChunks, + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + resolvePromiseUsage: true, + span, + startTime, + }), + }); // Trace a continuation as its own task only when its original turn cannot // be recovered. Known continuations extend the original Harness task. - this.unsubscribers.push( - traceStreamingChannel(harnessAgentChannels.continueGenerate, { - name: "HarnessAgent.continueGenerate", - shouldTrace: (args) => - !harnessContinuationParent(harnessSessionFromArguments(args)), - startSpan: _internalStartSpanWithInitialMerge, - type: SpanTypeAttribute.TASK, - extractInput: ([params], event, span) => - prepareAISDKHarnessAgentInput(params, event.self, span), - extractOutput: (result, endEvent) => - processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ), - extractMetrics: (result) => extractTokenMetrics(result), - aggregateChunks: aggregateAISDKChunks, - }), - ); + traceStreamingChannel(harnessAgentChannels.continueGenerate, { + name: "HarnessAgent.continueGenerate", + shouldTrace: (args) => + !harnessContinuationParent(harnessSessionFromArguments(args)), + startSpan: _internalStartSpanWithInitialMerge, + type: SpanTypeAttribute.TASK, + extractInput: ([params], event, span) => + prepareAISDKHarnessAgentInput(params, event.self, span), + extractOutput: (result, endEvent) => + processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ), + extractMetrics: (result) => extractTokenMetrics(result), + aggregateChunks: aggregateAISDKChunks, + }); - this.unsubscribers.push( - traceStreamingChannel(harnessAgentChannels.continueStream, { - name: "HarnessAgent.continueStream", - shouldTrace: (args) => - !harnessContinuationParent(harnessSessionFromArguments(args)), - startSpan: _internalStartSpanWithInitialMerge, - type: SpanTypeAttribute.TASK, - extractInput: ([params], event, span) => - prepareAISDKHarnessAgentInput(params, event.self, span), - extractOutput: (result, endEvent) => - processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ), - extractMetrics: (result, startTime) => ({ - ...extractTokenMetrics(result), - ...(startTime === undefined - ? {} - : { - time_to_first_token: getCurrentUnixTimestamp() - startTime, - }), - }), - aggregateChunks: aggregateAISDKChunks, - patchResult: ({ endEvent, result, span, startTime }) => - patchAISDKStreamingResult({ - defaultDenyOutputPaths: denyOutputPaths, - endEvent, - result, - resolvePromiseUsage: true, - span, - startTime, - }), + traceStreamingChannel(harnessAgentChannels.continueStream, { + name: "HarnessAgent.continueStream", + shouldTrace: (args) => + !harnessContinuationParent(harnessSessionFromArguments(args)), + startSpan: _internalStartSpanWithInitialMerge, + type: SpanTypeAttribute.TASK, + extractInput: ([params], event, span) => + prepareAISDKHarnessAgentInput(params, event.self, span), + extractOutput: (result, endEvent) => + processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ), + extractMetrics: (result, startTime) => ({ + ...extractTokenMetrics(result), + ...(startTime === undefined + ? {} + : { + time_to_first_token: getCurrentUnixTimestamp() - startTime, + }), }), - ); + aggregateChunks: aggregateAISDKChunks, + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + resolvePromiseUsage: true, + span, + startTime, + }), + }); // ToolLoopAgent.generate - async method - this.unsubscribers.push( - traceStreamingChannel(aiSDKChannels.toolLoopAgentGenerate, { - name: "ToolLoopAgent.generate", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event, span) => - prepareAISDKCallInput(params, event, span, denyOutputPaths, { - agentOwner: true, - }), - extractOutput: (result, endEvent) => { - finalizeAISDKChildTracing(endEvent as { [key: string]: unknown }); - return processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ); - }, - extractMetrics: (result, _startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent), - aggregateChunks: aggregateAISDKChunks, - }), - ); + traceStreamingChannel(aiSDKChannels.toolLoopAgentGenerate, { + name: "ToolLoopAgent.generate", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event, span) => + prepareAISDKCallInput(params, event, span, denyOutputPaths, { + agentOwner: true, + }), + extractOutput: (result, endEvent) => { + finalizeAISDKChildTracing(endEvent as { [key: string]: unknown }); + return processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ); + }, + extractMetrics: (result, _startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent), + aggregateChunks: aggregateAISDKChunks, + }); // ToolLoopAgent.stream - async method returning stream - this.unsubscribers.push( - traceStreamingChannel(aiSDKChannels.toolLoopAgentStream, { - name: "ToolLoopAgent.stream", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event, span) => - prepareAISDKCallInput(params, event, span, denyOutputPaths, { - agentOwner: true, - }), - extractOutput: (result, endEvent) => - processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ), - extractMetrics: (result, startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent, startTime), - aggregateChunks: aggregateAISDKChunks, - patchResult: ({ endEvent, result, span, startTime }) => - patchAISDKStreamingResult({ - defaultDenyOutputPaths: denyOutputPaths, - endEvent, - result, - span, - startTime, - }), - }), - ); + traceStreamingChannel(aiSDKChannels.toolLoopAgentStream, { + name: "ToolLoopAgent.stream", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event, span) => + prepareAISDKCallInput(params, event, span, denyOutputPaths, { + agentOwner: true, + }), + extractOutput: (result, endEvent) => + processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ), + extractMetrics: (result, startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent, startTime), + aggregateChunks: aggregateAISDKChunks, + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + result, + span, + startTime, + }), + }); // WorkflowAgent.stream - async method returning stream - this.unsubscribers.push( - traceStreamingChannel(aiSDKChannels.workflowAgentStream, { - name: "WorkflowAgent.stream", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params], event, span) => - prepareAISDKWorkflowAgentStreamInput( - params, - event, - span, - denyOutputPaths, - ), - extractOutput: (result, endEvent) => { - finalizeAISDKChildTracing(endEvent as { [key: string]: unknown }); - return processAISDKOutput( - result, - resolveDenyOutputPaths(endEvent, denyOutputPaths), - ); - }, - extractMetrics: (result, _startTime, endEvent) => - extractTopLevelAISDKMetrics(result, endEvent), - aggregateChunks: aggregateAISDKChunks, - onComplete: ({ span }) => { - unregisterWorkflowAgentWrapperSpan(span); - }, - onError: ({ event, span }) => { - finalizeAISDKChildTracing(event as { [key: string]: unknown }); - unregisterWorkflowAgentWrapperSpan(span); - }, - patchResult: ({ endEvent, result, span, startTime }) => - patchAISDKStreamingResult({ - defaultDenyOutputPaths: denyOutputPaths, - endEvent, - onComplete: () => unregisterWorkflowAgentWrapperSpan(span), - onCancel: () => unregisterWorkflowAgentWrapperSpan(span), - onError: () => unregisterWorkflowAgentWrapperSpan(span), - result, - span, - startTime, - }), - }), - ); + traceStreamingChannel(aiSDKChannels.workflowAgentStream, { + name: "WorkflowAgent.stream", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params], event, span) => + prepareAISDKWorkflowAgentStreamInput( + params, + event, + span, + denyOutputPaths, + ), + extractOutput: (result, endEvent) => { + finalizeAISDKChildTracing(endEvent as { [key: string]: unknown }); + return processAISDKOutput( + result, + resolveDenyOutputPaths(endEvent, denyOutputPaths), + ); + }, + extractMetrics: (result, _startTime, endEvent) => + extractTopLevelAISDKMetrics(result, endEvent), + aggregateChunks: aggregateAISDKChunks, + onComplete: ({ span }) => { + unregisterWorkflowAgentWrapperSpan(span); + }, + onError: ({ event, span }) => { + finalizeAISDKChildTracing(event as { [key: string]: unknown }); + unregisterWorkflowAgentWrapperSpan(span); + }, + patchResult: ({ endEvent, result, span, startTime }) => + patchAISDKStreamingResult({ + defaultDenyOutputPaths: denyOutputPaths, + endEvent, + onComplete: () => unregisterWorkflowAgentWrapperSpan(span), + onCancel: () => unregisterWorkflowAgentWrapperSpan(span), + onError: () => unregisterWorkflowAgentWrapperSpan(span), + result, + span, + startTime, + }), + }); } } -function subscribeToHarnessAgentCreateSession(): () => void { +function subscribeToHarnessAgentCreateSession(): void { const channel = harnessAgentChannels.createSession.tracingChannel(); const parents = new WeakMap(); const handlers: IsoChannelHandlers< @@ -682,7 +636,6 @@ function subscribeToHarnessAgentCreateSession(): () => void { }; channel.subscribe(handlers); - return () => channel.unsubscribe(handlers); } type HarnessContinuationChannel = typeof harnessAgentChannels.continueGenerate; @@ -700,14 +653,11 @@ function harnessContinuationParentFromEvent( function subscribeToHarnessContinuation( continuationChannel: HarnessContinuationChannel, defaultDenyOutputPaths: string[], -): () => void { +): void { const channel = continuationChannel.tracingChannel(); const parents = new WeakMap(); const startTimes = new WeakMap(); - const unbindParentStore = bindHarnessTurnParentToStart( - channel, - harnessContinuationParentFromEvent, - ); + bindHarnessTurnParentToStart(channel, harnessContinuationParentFromEvent); const handlers: IsoChannelHandlers< ChannelMessage > = { @@ -833,13 +783,9 @@ function subscribeToHarnessContinuation( }; channel.subscribe(handlers); - return () => { - unbindParentStore(); - channel.unsubscribe(handlers); - }; } -function subscribeToAISDKV7TelemetryDispatcher(): () => void { +function subscribeToAISDKV7TelemetryDispatcher(): void { const channel = aiSDKChannels.v7CreateTelemetryDispatcher.tracingChannel(); const telemetry = braintrustAISDKTelemetry(); const handlers: IsoChannelHandlers< @@ -860,10 +806,6 @@ function subscribeToAISDKV7TelemetryDispatcher(): () => void { }; channel.subscribe(handlers); - - return () => { - channel.unsubscribe(handlers); - }; } function patchAISDKV7TelemetryDispatcher( @@ -4274,3 +4216,12 @@ function omit( return result; } + +let aiSDKInstrumentationConsumer: AISDKInstrumentationConsumer | undefined; + +export function registerAISDKInstrumentation( + config: AISDKInstrumentationConfig = {}, +): void { + aiSDKInstrumentationConsumer ??= new AISDKInstrumentationConsumer(config); + aiSDKInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/ai-sdk-v7-telemetry.test.ts b/js/src/instrumentation/providers/ai-sdk-v7-telemetry.test.ts similarity index 100% rename from js/src/instrumentation/plugins/ai-sdk-v7-telemetry.test.ts rename to js/src/instrumentation/providers/ai-sdk-v7-telemetry.test.ts diff --git a/js/src/instrumentation/plugins/anthropic-channels.ts b/js/src/instrumentation/providers/anthropic-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/anthropic-channels.ts rename to js/src/instrumentation/providers/anthropic-channels.ts diff --git a/js/src/instrumentation/plugins/anthropic-plugin.test.ts b/js/src/instrumentation/providers/anthropic-instrumentation.test.ts similarity index 99% rename from js/src/instrumentation/plugins/anthropic-plugin.test.ts rename to js/src/instrumentation/providers/anthropic-instrumentation.test.ts index ffad58aff..3e3d56708 100644 --- a/js/src/instrumentation/plugins/anthropic-plugin.test.ts +++ b/js/src/instrumentation/providers/anthropic-instrumentation.test.ts @@ -12,7 +12,7 @@ import { aggregateAnthropicStreamChunks, processAttachmentsInInput, coalesceInput, -} from "./anthropic-plugin"; +} from "./anthropic-instrumentation"; import { Attachment } from "../../logger"; const parseMetricsFromUsageForTest = (usage: unknown) => diff --git a/js/src/instrumentation/plugins/anthropic-plugin.ts b/js/src/instrumentation/providers/anthropic-instrumentation.ts similarity index 98% rename from js/src/instrumentation/plugins/anthropic-plugin.ts rename to js/src/instrumentation/providers/anthropic-instrumentation.ts index 36cad5a5c..9b6cc128e 100644 --- a/js/src/instrumentation/plugins/anthropic-plugin.ts +++ b/js/src/instrumentation/providers/anthropic-instrumentation.ts @@ -1,5 +1,5 @@ -import { BasePlugin, toLoggedError } from "../core"; -import { traceStreamingChannel, unsubscribeAll } from "../core/channel-tracing"; +import { toLoggedError } from "../core"; +import { traceStreamingChannel } from "../core/channel-tracing"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; import { Attachment, @@ -93,30 +93,26 @@ const ANTHROPIC_TOOL_RUNNER_TOOL_WRAPPED = Symbol.for( ); /** - * Auto-instrumentation plugin for the Anthropic SDK. + * Internal auto-instrumentation consumer for the Anthropic SDK. * - * This plugin subscribes to orchestrion channels for Anthropic SDK methods + * This consumer subscribes to orchestrion channels for Anthropic SDK methods * and creates Braintrust spans to track: * - messages.create (streaming and non-streaming) * - beta.messages.create (streaming and non-streaming) * - * The plugin handles: + * The consumer handles: * - Anthropic-specific token metrics (including cache tokens) * - Processing message streams * - Converting base64 attachments to Attachment objects * - Streaming and non-streaming responses */ -export class AnthropicPlugin extends BasePlugin { - protected onEnable(): void { +class AnthropicInstrumentationConsumer { + public register(): void { this.subscribeToAnthropicChannels(); this.subscribeToAnthropicToolRunner(); this.subscribeToAnthropicSessionStreams(); } - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } - private subscribeToAnthropicChannels(): void { const anthropicConfig = { name: "anthropic.messages.create", @@ -163,17 +159,13 @@ export class AnthropicPlugin extends BasePlugin { }; // Messages API - supports streaming via stream=true parameter - this.unsubscribers.push( - traceStreamingChannel(anthropicChannels.messagesCreate, anthropicConfig), - ); + traceStreamingChannel(anthropicChannels.messagesCreate, anthropicConfig); // Beta Messages API - supports streaming via stream=true parameter - this.unsubscribers.push( - traceStreamingChannel(anthropicChannels.betaMessagesCreate, { - ...anthropicConfig, - name: "anthropic.messages.create", - }), - ); + traceStreamingChannel(anthropicChannels.betaMessagesCreate, { + ...anthropicConfig, + name: "anthropic.messages.create", + }); } private subscribeToAnthropicToolRunner(): void { @@ -250,9 +242,6 @@ export class AnthropicPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - tracingChannel.unsubscribe(handlers); - }); } private subscribeToAnthropicSessionStreams(): void { @@ -305,7 +294,6 @@ export class AnthropicPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => tracingChannel.unsubscribe(handlers)); } } @@ -1686,3 +1674,12 @@ export function coalesceInput( } return input; } + +let anthropicInstrumentationConsumer: + | AnthropicInstrumentationConsumer + | undefined; + +export function registerAnthropicInstrumentation(): void { + anthropicInstrumentationConsumer ??= new AnthropicInstrumentationConsumer(); + anthropicInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/anthropic-sessions-plugin.test.ts b/js/src/instrumentation/providers/anthropic-sessions-instrumentation.test.ts similarity index 97% rename from js/src/instrumentation/plugins/anthropic-sessions-plugin.test.ts rename to js/src/instrumentation/providers/anthropic-sessions-instrumentation.test.ts index 1d4cb48ac..44bd501d6 100644 --- a/js/src/instrumentation/plugins/anthropic-sessions-plugin.test.ts +++ b/js/src/instrumentation/providers/anthropic-sessions-instrumentation.test.ts @@ -28,11 +28,11 @@ vi.mock("../../logger", async (importOriginal) => { import iso from "../../isomorph"; import { collectAnthropicSession } from "../../wrappers/anthropic-session-collector"; -import { AnthropicPlugin } from "./anthropic-plugin"; +import { registerAnthropicInstrumentation } from "./anthropic-instrumentation"; const mockNewTracingChannel = iso.newTracingChannel as ReturnType; -describe("AnthropicPlugin Sessions instrumentation", () => { +describe("registerAnthropicInstrumentation Sessions instrumentation", () => { let currentSpan: TestSpan | undefined; let handlersByName: Map; let spans: TestSpan[]; @@ -43,7 +43,6 @@ describe("AnthropicPlugin Sessions instrumentation", () => { spans = []; mockNewTracingChannel.mockImplementation((name: string) => ({ subscribe: vi.fn((handlers) => handlersByName.set(name, handlers)), - unsubscribe: vi.fn(), })); mockWithCurrent.mockImplementation( (span: TestSpan, callback: () => unknown) => { @@ -78,7 +77,7 @@ describe("AnthropicPlugin Sessions instrumentation", () => { }); it("records a session turn with model and tool children", async () => { - new AnthropicPlugin().enable(); + registerAnthropicInstrumentation(); const handlers = handlersByName.get( "orchestrion:@anthropic-ai/sdk:beta.sessions.events.stream", ); @@ -262,7 +261,7 @@ describe("AnthropicPlugin Sessions instrumentation", () => { }); it("does not record an uncollected session stream", async () => { - new AnthropicPlugin().enable(); + registerAnthropicInstrumentation(); const handlers = handlersByName.get( "orchestrion:@anthropic-ai/sdk:beta.sessions.events.stream", ); @@ -286,7 +285,7 @@ describe("AnthropicPlugin Sessions instrumentation", () => { }); it("uses a distinct task name for thread streams and records denials", async () => { - new AnthropicPlugin().enable(); + registerAnthropicInstrumentation(); const handlers = handlersByName.get( "orchestrion:@anthropic-ai/sdk:beta.sessions.threads.events.stream", ); diff --git a/js/src/instrumentation/plugins/bedrock-runtime-channels.ts b/js/src/instrumentation/providers/bedrock-runtime-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/bedrock-runtime-channels.ts rename to js/src/instrumentation/providers/bedrock-runtime-channels.ts diff --git a/js/src/instrumentation/plugins/bedrock-runtime-common.ts b/js/src/instrumentation/providers/bedrock-runtime-common.ts similarity index 100% rename from js/src/instrumentation/plugins/bedrock-runtime-common.ts rename to js/src/instrumentation/providers/bedrock-runtime-common.ts diff --git a/js/src/instrumentation/plugins/bedrock-runtime-plugin.test.ts b/js/src/instrumentation/providers/bedrock-runtime-instrumentation.test.ts similarity index 99% rename from js/src/instrumentation/plugins/bedrock-runtime-plugin.test.ts rename to js/src/instrumentation/providers/bedrock-runtime-instrumentation.test.ts index fc74a9068..b22906083 100644 --- a/js/src/instrumentation/plugins/bedrock-runtime-plugin.test.ts +++ b/js/src/instrumentation/providers/bedrock-runtime-instrumentation.test.ts @@ -9,7 +9,7 @@ import { aggregateBedrockConverseStreamChunks, aggregateInvokeModelResponseStreamChunks, parseBedrockRuntimeMetrics, -} from "./bedrock-runtime-plugin"; +} from "./bedrock-runtime-instrumentation"; try { configureNode(); @@ -29,7 +29,7 @@ class GetObjectCommand { constructor(public input: Record) {} } -describe("BedrockRuntimePlugin", () => { +describe("registerBedrockRuntimeInstrumentation", () => { let backgroundLogger: ReturnType< typeof _exportsForTestingOnly.useTestBackgroundLogger >; @@ -41,7 +41,7 @@ describe("BedrockRuntimePlugin", () => { beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ - projectName: "bedrock-runtime-plugin.test.ts", + projectName: "bedrock-runtime-instrumentation.test.ts", projectId: "test-project-id", }); }); diff --git a/js/src/instrumentation/plugins/bedrock-runtime-plugin.ts b/js/src/instrumentation/providers/bedrock-runtime-instrumentation.ts similarity index 97% rename from js/src/instrumentation/plugins/bedrock-runtime-plugin.ts rename to js/src/instrumentation/providers/bedrock-runtime-instrumentation.ts index a759d1ffe..dfd5a47be 100644 --- a/js/src/instrumentation/plugins/bedrock-runtime-plugin.ts +++ b/js/src/instrumentation/providers/bedrock-runtime-instrumentation.ts @@ -1,5 +1,4 @@ -import { BasePlugin } from "../core"; -import { traceStreamingChannel, unsubscribeAll } from "../core/channel-tracing"; +import { traceStreamingChannel } from "../core/channel-tracing"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; import { SpanTypeAttribute, isObject } from "../../../util/index"; import { getCurrentUnixTimestamp } from "../../util"; @@ -25,26 +24,18 @@ import { getBedrockRuntimeOperation, } from "./bedrock-runtime-common"; -export class BedrockRuntimePlugin extends BasePlugin { - protected onEnable(): void { - this.unsubscribers.push( - ...[ - bedrockRuntimeChannels.clientSend, - smithyCoreChannels.clientSend, - smithyClientChannels.clientSend, - ].map((channel) => traceBedrockRuntimeClientSendChannel(channel)), - ); - } - - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); +export function registerBedrockRuntimeInstrumentation(): void { + for (const channel of [ + bedrockRuntimeChannels.clientSend, + smithyCoreChannels.clientSend, + smithyClientChannels.clientSend, + ]) { + traceBedrockRuntimeClientSendChannel(channel); } } -function traceBedrockRuntimeClientSendChannel( - channel: AnyAsyncChannel, -): () => void { - return traceStreamingChannel(channel, { +function traceBedrockRuntimeClientSendChannel(channel: AnyAsyncChannel): void { + traceStreamingChannel(channel, { name: ([command]) => buildBedrockRuntimeSpanInfo(command).name, shouldTrace: ([command, optionsOrCb, cb]) => getBedrockRuntimeOperation(command) !== undefined && diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-channels.ts b/js/src/instrumentation/providers/claude-agent-sdk-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/claude-agent-sdk-channels.ts rename to js/src/instrumentation/providers/claude-agent-sdk-channels.ts diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-instrumentation-constants.ts b/js/src/instrumentation/providers/claude-agent-sdk-instrumentation-constants.ts similarity index 100% rename from js/src/instrumentation/plugins/claude-agent-sdk-instrumentation-constants.ts rename to js/src/instrumentation/providers/claude-agent-sdk-instrumentation-constants.ts diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.streaming.test.ts b/js/src/instrumentation/providers/claude-agent-sdk-instrumentation.streaming.test.ts similarity index 99% rename from js/src/instrumentation/plugins/claude-agent-sdk-plugin.streaming.test.ts rename to js/src/instrumentation/providers/claude-agent-sdk-instrumentation.streaming.test.ts index 42101e066..ddaeeab1f 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.streaming.test.ts +++ b/js/src/instrumentation/providers/claude-agent-sdk-instrumentation.streaming.test.ts @@ -120,7 +120,7 @@ describe("Claude Agent SDK streaming instrumentation", () => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ projectId: "test-project-id", - projectName: "claude-agent-sdk-plugin.streaming.test.ts", + projectName: "claude-agent-sdk-instrumentation.streaming.test.ts", }); }); diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts b/js/src/instrumentation/providers/claude-agent-sdk-instrumentation.test.ts similarity index 86% rename from js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts rename to js/src/instrumentation/providers/claude-agent-sdk-instrumentation.test.ts index f1d446bcd..93265e4a3 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.test.ts +++ b/js/src/instrumentation/providers/claude-agent-sdk-instrumentation.test.ts @@ -30,7 +30,7 @@ vi.mock("../core/stream-patcher", () => ({ }), })); -import { ClaudeAgentSDKPlugin } from "./claude-agent-sdk-plugin"; +import { registerClaudeAgentSDKInstrumentation } from "./claude-agent-sdk-instrumentation"; import iso from "../../isomorph"; import { startSpan } from "../../logger"; @@ -72,79 +72,26 @@ vi.mock("../../wrappers/attachment-utils", () => ({ // rules, so these tests run the real implementation rather than a stand-in that // could drift from it. -vi.mock("../core", async (importOriginal) => { - const actual = await importOriginal(); - - return { - ...actual, - BasePlugin: class BasePlugin { - protected enabled = false; - protected unsubscribers: Array<() => void> = []; - - enable(): void { - if (this.enabled) { - return; - } - this.enabled = true; - this.onEnable(); - } - - disable(): void { - if (!this.enabled) { - return; - } - this.enabled = false; - this.onDisable(); - } - - protected onEnable(): void { - // To be implemented by subclass - } - - protected onDisable(): void { - // To be implemented by subclass - } - }, - isAsyncIterable: vi.fn( - (val: unknown) => - val !== null && - typeof val === "object" && - Symbol.asyncIterator in val && - typeof (val as any)[Symbol.asyncIterator] === "function", - ), - patchStreamIfNeeded: vi.fn((stream, _callbacks) => { - // Return the stream unchanged for simple tests - return stream; - }), - }; -}); - -describe("ClaudeAgentSDKPlugin", () => { - let plugin: ClaudeAgentSDKPlugin; +describe("registerClaudeAgentSDKInstrumentation", () => { let mockChannel: any; - let mockUnsubscribe: any; beforeEach(() => { streamPatcherMock.options = undefined; - mockUnsubscribe = vi.fn(); mockChannel = { subscribe: vi.fn(), - unsubscribe: mockUnsubscribe, hasSubscribers: false, }; mockNewTracingChannel.mockReturnValue(mockChannel); - - plugin = new ClaudeAgentSDKPlugin(); }); afterEach(() => { vi.clearAllMocks(); }); - describe("enable", () => { - it("should enable the plugin and subscribe to channels", () => { - plugin.enable(); + describe("registration", () => { + it("subscribes to the channel", () => { + registerClaudeAgentSDKInstrumentation(); expect(mockNewTracingChannel).toHaveBeenCalledWith( "orchestrion:@anthropic-ai/claude-agent-sdk:query", @@ -158,50 +105,13 @@ describe("ClaudeAgentSDKPlugin", () => { }), ); }); - - it("should not subscribe twice if already enabled", () => { - plugin.enable(); - plugin.enable(); - - expect(mockChannel.subscribe).toHaveBeenCalledTimes(1); - }); - - it("should store unsubscribe function", () => { - plugin.enable(); - - expect((plugin as any).unsubscribers).toHaveLength(1); - expect((plugin as any).unsubscribers[0]).toBeInstanceOf(Function); - }); - }); - - describe("disable", () => { - it("should unsubscribe from all channels", () => { - plugin.enable(); - plugin.disable(); - - expect(mockUnsubscribe).toHaveBeenCalledTimes(1); - expect((plugin as any).unsubscribers).toHaveLength(0); - }); - - it("should not unsubscribe if not enabled", () => { - plugin.disable(); - - expect(mockUnsubscribe).not.toHaveBeenCalled(); - }); - - it("should clear unsubscribers array", () => { - plugin.enable(); - plugin.disable(); - - expect((plugin as any).unsubscribers).toHaveLength(0); - }); }); describe("channel subscription handlers", () => { let handlers: any; beforeEach(() => { - plugin.enable(); + registerClaudeAgentSDKInstrumentation(); handlers = mockChannel.subscribe.mock.calls[0][0]; }); @@ -745,24 +655,4 @@ describe("ClaudeAgentSDKPlugin", () => { }); }); }); - - describe("enable/disable lifecycle", () => { - it("should allow re-enabling after disable", () => { - plugin.enable(); - plugin.disable(); - plugin.enable(); - - expect(mockChannel.subscribe).toHaveBeenCalledTimes(2); - }); - - it("should properly clean up on multiple enable/disable cycles", () => { - plugin.enable(); - plugin.disable(); - plugin.enable(); - plugin.disable(); - - expect(mockUnsubscribe).toHaveBeenCalledTimes(2); - expect((plugin as any).unsubscribers).toHaveLength(0); - }); - }); }); diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts b/js/src/instrumentation/providers/claude-agent-sdk-instrumentation.ts similarity index 99% rename from js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts rename to js/src/instrumentation/providers/claude-agent-sdk-instrumentation.ts index 96f52c9b6..9944ebe60 100644 --- a/js/src/instrumentation/plugins/claude-agent-sdk-plugin.ts +++ b/js/src/instrumentation/providers/claude-agent-sdk-instrumentation.ts @@ -1,4 +1,3 @@ -import { BasePlugin } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; import type { IsoChannelHandlers } from "../../isomorph"; @@ -1578,18 +1577,11 @@ async function finalizeQuerySpan(state: QueryState): Promise { } } -export class ClaudeAgentSDKPlugin extends BasePlugin { - protected onEnable(): void { +class ClaudeAgentSDKInstrumentationConsumer { + public register(): void { this.subscribeToQuery(); } - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - } - private subscribeToQuery(): void { const channel = claudeAgentSDKChannels.query.tracingChannel(); const spans = new WeakMap(); @@ -1854,8 +1846,15 @@ export class ClaudeAgentSDKPlugin extends BasePlugin { }; channel.subscribe(handlers); - this.unsubscribers.push(() => { - channel.unsubscribe(handlers); - }); } } + +let claudeAgentSDKInstrumentationConsumer: + | ClaudeAgentSDKInstrumentationConsumer + | undefined; + +export function registerClaudeAgentSDKInstrumentation(): void { + claudeAgentSDKInstrumentationConsumer ??= + new ClaudeAgentSDKInstrumentationConsumer(); + claudeAgentSDKInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts b/js/src/instrumentation/providers/claude-agent-sdk-local-tool-context.ts similarity index 100% rename from js/src/instrumentation/plugins/claude-agent-sdk-local-tool-context.ts rename to js/src/instrumentation/providers/claude-agent-sdk-local-tool-context.ts diff --git a/js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts b/js/src/instrumentation/providers/claude-agent-sdk-local-tool-spans.ts similarity index 100% rename from js/src/instrumentation/plugins/claude-agent-sdk-local-tool-spans.ts rename to js/src/instrumentation/providers/claude-agent-sdk-local-tool-spans.ts diff --git a/js/src/instrumentation/plugins/cloudflare-agents-channels.ts b/js/src/instrumentation/providers/cloudflare-agents-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/cloudflare-agents-channels.ts rename to js/src/instrumentation/providers/cloudflare-agents-channels.ts diff --git a/js/src/instrumentation/plugins/cloudflare-agents-plugin.test.ts b/js/src/instrumentation/providers/cloudflare-agents-instrumentation.test.ts similarity index 88% rename from js/src/instrumentation/plugins/cloudflare-agents-plugin.test.ts rename to js/src/instrumentation/providers/cloudflare-agents-instrumentation.test.ts index d6d2de7de..970297b01 100644 --- a/js/src/instrumentation/plugins/cloudflare-agents-plugin.test.ts +++ b/js/src/instrumentation/providers/cloudflare-agents-instrumentation.test.ts @@ -29,14 +29,13 @@ vi.mock("../../isomorph", () => ({ })); import iso from "../../isomorph"; -import { CloudflareAgentsPlugin } from "./cloudflare-agents-plugin"; +import { registerCloudflareAgentsInstrumentation } from "./cloudflare-agents-instrumentation"; const mockNewTracingChannel = iso.newTracingChannel as ReturnType; -describe("CloudflareAgentsPlugin", () => { +describe("registerCloudflareAgentsInstrumentation", () => { let handlers: any; let subscribe: ReturnType; - let unsubscribe: ReturnType; let spans: Array<{ args: any; context: any; @@ -49,8 +48,7 @@ describe("CloudflareAgentsPlugin", () => { subscribe = vi.fn((nextHandlers) => { handlers = nextHandlers; }); - unsubscribe = vi.fn(); - mockNewTracingChannel.mockReturnValue({ subscribe, unsubscribe }); + mockNewTracingChannel.mockReturnValue({ subscribe }); mockStartSpan.mockImplementation((args: any, context: any) => { const span = { args, context, end: vi.fn(), log: vi.fn() }; spans.push(span); @@ -62,19 +60,13 @@ describe("CloudflareAgentsPlugin", () => { vi.clearAllMocks(); }); - it("subscribes idempotently to Agent.runAgentTool", () => { - const plugin = new CloudflareAgentsPlugin(); - plugin.enable(); - plugin.enable(); + it("subscribes to the process-lifetime channel", () => { + registerCloudflareAgentsInstrumentation(); expect(mockNewTracingChannel).toHaveBeenCalledWith( "orchestrion:agents:Agent.runAgentTool", ); expect(subscribe).toHaveBeenCalledTimes(1); - - plugin.disable(); - plugin.disable(); - expect(unsubscribe).toHaveBeenCalledTimes(1); }); it("keeps SDK-controlled context out of the public start-span arguments", () => { @@ -86,7 +78,7 @@ describe("CloudflareAgentsPlugin", () => { }); it("records only the child class name, input, and completed output", () => { - new CloudflareAgentsPlugin().enable(); + registerCloudflareAgentsInstrumentation(); class ResearchAgent {} const event = { arguments: [ @@ -142,7 +134,7 @@ describe("CloudflareAgentsPlugin", () => { }); it("records returned terminal error strings", () => { - new CloudflareAgentsPlugin().enable(); + registerCloudflareAgentsInstrumentation(); class FailingAgent {} const event = { arguments: [FailingAgent, { input: "fail" }] }; @@ -164,7 +156,7 @@ describe("CloudflareAgentsPlugin", () => { }); it("records the original rejection and preserves concurrent span state", () => { - new CloudflareAgentsPlugin().enable(); + registerCloudflareAgentsInstrumentation(); class FirstAgent {} class SecondAgent {} const first = { arguments: [FirstAgent, { input: 1 }] }; @@ -188,7 +180,7 @@ describe("CloudflareAgentsPlugin", () => { }); it("skips detached runs and does not invoke getters", () => { - new CloudflareAgentsPlugin().enable(); + registerCloudflareAgentsInstrumentation(); const nameGetter = vi.fn(() => "GetterAgent"); const inputGetter = vi.fn(() => "getter-input"); const AgentWithGetter = Object.defineProperty(function () {}, "name", { diff --git a/js/src/instrumentation/providers/cloudflare-agents-instrumentation.ts b/js/src/instrumentation/providers/cloudflare-agents-instrumentation.ts new file mode 100644 index 000000000..36fe3671a --- /dev/null +++ b/js/src/instrumentation/providers/cloudflare-agents-instrumentation.ts @@ -0,0 +1,129 @@ +import { debugLogger } from "../../debug-logger"; +import type { IsoChannelHandlers } from "../../isomorph"; +import { _internalStartSpanWithContext } from "../../logger"; +import type { Span } from "../../logger"; +import { + INSTRUMENTATION_NAMES, + withSpanInstrumentationName, +} from "../../span-origin"; +import { SpanTypeAttribute } from "../../../util/index"; +import type { ChannelMessage } from "../core/channel-definitions"; +import { cloudflareAgentsChannels } from "./cloudflare-agents-channels"; + +const CLOUDFLARE_WORKERS_CONTEXT = { + span_origin: { + environment: { type: "server", name: "cloudflare_workers" }, + }, +}; + +export function registerCloudflareAgentsInstrumentation(): void { + const channel = cloudflareAgentsChannels.runAgentTool.tracingChannel(); + const spans = new WeakMap(); + const handlers: IsoChannelHandlers< + ChannelMessage + > = { + start: (event) => { + try { + const agentClass = event.arguments[0]; + const options = event.arguments[1]; + if (ownValue(options, "detached")) { + return; + } + + const name = ownValue(agentClass, "name"); + if (typeof name !== "string" || name.length === 0) { + debugLogger.warn( + "Skipping Cloudflare Agents runAgentTool span because the child agent class has no name.", + ); + return; + } + + const span = _internalStartSpanWithContext( + withSpanInstrumentationName( + { + name, + spanAttributes: { type: SpanTypeAttribute.TOOL }, + event: { + input: ownValue(options, "input"), + }, + }, + INSTRUMENTATION_NAMES.CLOUDFLARE_AGENTS, + ), + CLOUDFLARE_WORKERS_CONTEXT, + ); + spans.set(event, span); + } catch (error) { + logInstrumentationError("start", error); + } + }, + asyncEnd: (event) => { + const span = spans.get(event); + if (!span) { + return; + } + spans.delete(event); + + try { + const status = ownValue(event.result, "status"); + if (status === "completed") { + span.log({ output: ownValue(event.result, "output") }); + } else { + const error = ownValue(event.result, "error"); + if (typeof error === "string") { + span.log({ error }); + } + } + } catch (error) { + logInstrumentationError("completion", error); + } finally { + safelyEndSpan(span); + } + }, + error: (event) => { + const span = spans.get(event); + if (!span) { + return; + } + spans.delete(event); + + try { + span.log({ error: event.error }); + } catch (error) { + logInstrumentationError("rejection", error); + } finally { + safelyEndSpan(span); + } + }, + }; + + channel.subscribe(handlers); +} + +function ownValue(value: unknown, key: PropertyKey): unknown { + if (!isObjectLike(value)) { + return undefined; + } + const descriptor = Object.getOwnPropertyDescriptor(value, key); + return descriptor && "value" in descriptor ? descriptor.value : undefined; +} + +function isObjectLike(value: unknown): value is object { + return ( + (typeof value === "object" && value !== null) || typeof value === "function" + ); +} + +function safelyEndSpan(span: Span): void { + try { + span.end(); + } catch (error) { + logInstrumentationError("span end", error); + } +} + +function logInstrumentationError(operation: string, error: unknown): void { + debugLogger.error( + `Failed to process Cloudflare Agents ${operation} instrumentation:`, + error, + ); +} diff --git a/js/src/instrumentation/plugins/cloudflare-ai-chat-channels.ts b/js/src/instrumentation/providers/cloudflare-ai-chat-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/cloudflare-ai-chat-channels.ts rename to js/src/instrumentation/providers/cloudflare-ai-chat-channels.ts diff --git a/js/src/instrumentation/plugins/cloudflare-ai-chat-plugin.test.ts b/js/src/instrumentation/providers/cloudflare-ai-chat-consumer.test.ts similarity index 94% rename from js/src/instrumentation/plugins/cloudflare-ai-chat-plugin.test.ts rename to js/src/instrumentation/providers/cloudflare-ai-chat-consumer.test.ts index d20631a2d..860ab0b18 100644 --- a/js/src/instrumentation/plugins/cloudflare-ai-chat-plugin.test.ts +++ b/js/src/instrumentation/providers/cloudflare-ai-chat-consumer.test.ts @@ -30,12 +30,11 @@ import { INSTRUMENTATION_NAMES, INTERNAL_SPAN_INSTRUMENTATION_NAME, } from "../../span-origin"; -import { CloudflareAIChatPlugin } from "./cloudflare-ai-chat-plugin"; +import { registerCloudflareAIChatInstrumentation } from "./cloudflare-ai-chat-consumer"; const mockNewTracingChannel = iso.newTracingChannel as ReturnType; -describe("CloudflareAIChatPlugin", () => { - let plugin: CloudflareAIChatPlugin; +describe("registerCloudflareAIChatInstrumentation", () => { let channels: Map>; beforeEach(() => { @@ -54,16 +53,14 @@ describe("CloudflareAIChatPlugin", () => { log: vi.fn(), })); mockInternalGetGlobalState.mockReturnValue(undefined); - plugin = new CloudflareAIChatPlugin(); }); afterEach(() => { - plugin.disable(); vi.clearAllMocks(); }); it("captures the full successful turn and binds queued work", async () => { - plugin.enable(); + registerCloudflareAIChatInstrumentation(); const turnHandlers = turnChannel().handlers(); const callback = vi.fn(async () => "callback-result"); const agent = { @@ -132,7 +129,7 @@ describe("CloudflareAIChatPlugin", () => { }); it("correlates response errors and preserves partial output", async () => { - plugin.enable(); + registerCloudflareAIChatInstrumentation(); const turnHandlers = turnChannel().handlers(); const responseHandlers = responseChannel().handlers(); const agent = { messages: [], onChatResponse() {} }; @@ -169,7 +166,7 @@ describe("CloudflareAIChatPlugin", () => { }); it("retains a settled turn until its queued response is observed", async () => { - plugin.enable(); + registerCloudflareAIChatInstrumentation(); const handlers = turnChannel().handlers(); const agent = { messages: [ @@ -225,7 +222,7 @@ describe("CloudflareAIChatPlugin", () => { it("drops retained turns that never produce a response", () => { vi.useFakeTimers(); try { - plugin.enable(); + registerCloudflareAIChatInstrumentation(); const handlers = turnChannel().handlers(); const agent = { messages: [], @@ -263,7 +260,7 @@ describe("CloudflareAIChatPlugin", () => { }); it("preserves pre-turn input when a continuation reuses its output id", async () => { - plugin.enable(); + registerCloudflareAIChatInstrumentation(); const handlers = turnChannel().handlers(); const agent = { messages: [ @@ -324,7 +321,7 @@ describe("CloudflareAIChatPlugin", () => { }); it("deduplicates nested manual and automatic turn events", () => { - plugin.enable(); + registerCloudflareAIChatInstrumentation(); const handlers = turnChannel().handlers(); const agent = { messages: [], onChatResponse() {} }; const outer = { @@ -348,8 +345,8 @@ describe("CloudflareAIChatPlugin", () => { expect(span.end).toHaveBeenCalledTimes(1); }); - it("logs original errors and closes outstanding spans on disable", () => { - plugin.enable(); + it("logs original errors without closing unrelated outstanding spans", () => { + registerCloudflareAIChatInstrumentation(); const handlers = turnChannel().handlers(); const failure = new Error("turn failed"); const failedEvent = { @@ -370,8 +367,7 @@ describe("CloudflareAIChatPlugin", () => { } as any; handlers.start?.(pendingEvent, "start"); const pendingSpan = mockStartSpan.mock.results[1].value; - plugin.disable(); - expect(pendingSpan.end).toHaveBeenCalledTimes(1); + expect(pendingSpan.end).not.toHaveBeenCalled(); }); function turnChannel() { @@ -394,7 +390,6 @@ function createMockChannel() { hasSubscribers: false, start: { bindStore: vi.fn(), - unbindStore: vi.fn(), }, subscribe: vi.fn((handlers) => subscribed.push(handlers)), traceSync: vi.fn((callback, event) => { @@ -410,6 +405,5 @@ function createMockChannel() { throw error; } }), - unsubscribe: vi.fn(), }; } diff --git a/js/src/instrumentation/plugins/cloudflare-ai-chat-plugin.ts b/js/src/instrumentation/providers/cloudflare-ai-chat-consumer.ts similarity index 94% rename from js/src/instrumentation/plugins/cloudflare-ai-chat-plugin.ts rename to js/src/instrumentation/providers/cloudflare-ai-chat-consumer.ts index d875824af..a75ad3ea6 100644 --- a/js/src/instrumentation/plugins/cloudflare-ai-chat-plugin.ts +++ b/js/src/instrumentation/providers/cloudflare-ai-chat-consumer.ts @@ -1,4 +1,3 @@ -import { BasePlugin } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import type { IsoChannelHandlers, IsoTracingChannel } from "../../isomorph"; import { @@ -41,7 +40,7 @@ type TurnState = { const BLOCKED_KEYS = new Set(["__proto__", "constructor", "prototype"]); const RESPONSE_STATE_RETENTION_MS = 60_000; -export class CloudflareAIChatPlugin extends BasePlugin { +class CloudflareAIChatInstrumentationConsumer { private readonly activeStates = new Set(); private readonly activeTurns = new WeakMap< object, @@ -49,30 +48,18 @@ export class CloudflareAIChatPlugin extends BasePlugin { >(); private readonly eventStates = new WeakMap(); - protected onEnable(): void { + public register(): void { this.subscribeToResponseHook(); this.subscribeToTurnRunner(); } - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - - for (const state of this.activeStates) { - this.cleanupState(state); - } - this.activeStates.clear(); - } - private subscribeToTurnRunner(): void { const tracingChannel = cloudflareAIChatChannels.runExclusiveChatTurn.tracingChannel() as IsoTracingChannel< ChannelMessage >; - const unbindCurrentSpanStore = this.bindCurrentSpanStore(tracingChannel); + this.bindCurrentSpanStore(tracingChannel); const handlers: IsoChannelHandlers> = { start: (event) => { this.ensureEventState(event); @@ -86,10 +73,6 @@ export class CloudflareAIChatPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - unbindCurrentSpanStore?.(); - tracingChannel.unsubscribe(handlers); - }); } private subscribeToResponseHook(): void { @@ -143,12 +126,11 @@ export class CloudflareAIChatPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => tracingChannel.unsubscribe(handlers)); } private bindCurrentSpanStore( tracingChannel: IsoTracingChannel>, - ): (() => void) | undefined { + ): void { const globalState = _internalGetGlobalState(); const contextManager = globalState?.contextManager; const startChannel = tracingChannel.start; @@ -161,7 +143,7 @@ export class CloudflareAIChatPlugin extends BasePlugin { : undefined; if (!startChannel || !currentSpanStore || !contextManager) { - return undefined; + return; } startChannel.bindStore(currentSpanStore, (event) => { @@ -170,8 +152,6 @@ export class CloudflareAIChatPlugin extends BasePlugin { ? contextManager.wrapSpanForStore(state.span) : currentSpanStore.getStore(); }); - - return () => startChannel.unbindStore(currentSpanStore); } private ensureEventState( @@ -481,3 +461,13 @@ function sanitizeLoggedValue( } return Object.fromEntries(entries); } + +let cloudflareAIChatInstrumentationConsumer: + | CloudflareAIChatInstrumentationConsumer + | undefined; + +export function registerCloudflareAIChatInstrumentation(): void { + cloudflareAIChatInstrumentationConsumer ??= + new CloudflareAIChatInstrumentationConsumer(); + cloudflareAIChatInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/cloudflare-ai-chat-instrumentation.ts b/js/src/instrumentation/providers/cloudflare-ai-chat-instrumentation.ts similarity index 100% rename from js/src/instrumentation/plugins/cloudflare-ai-chat-instrumentation.ts rename to js/src/instrumentation/providers/cloudflare-ai-chat-instrumentation.ts diff --git a/js/src/instrumentation/plugins/cloudflare-think-channels.ts b/js/src/instrumentation/providers/cloudflare-think-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/cloudflare-think-channels.ts rename to js/src/instrumentation/providers/cloudflare-think-channels.ts diff --git a/js/src/instrumentation/plugins/cloudflare-think-context.ts b/js/src/instrumentation/providers/cloudflare-think-context.ts similarity index 100% rename from js/src/instrumentation/plugins/cloudflare-think-context.ts rename to js/src/instrumentation/providers/cloudflare-think-context.ts diff --git a/js/src/instrumentation/plugins/cloudflare-think-plugin.test.ts b/js/src/instrumentation/providers/cloudflare-think-instrumentation.test.ts similarity index 98% rename from js/src/instrumentation/plugins/cloudflare-think-plugin.test.ts rename to js/src/instrumentation/providers/cloudflare-think-instrumentation.test.ts index 240c33df0..f3e27919e 100644 --- a/js/src/instrumentation/plugins/cloudflare-think-plugin.test.ts +++ b/js/src/instrumentation/providers/cloudflare-think-instrumentation.test.ts @@ -27,7 +27,7 @@ describe("Cloudflare Think instrumentation", () => { beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ - projectName: "cloudflare-think-plugin.test.ts", + projectName: "cloudflare-think-instrumentation.test.ts", projectId: "test-project-id", }); }); diff --git a/js/src/instrumentation/plugins/cloudflare-think-plugin.ts b/js/src/instrumentation/providers/cloudflare-think-instrumentation.ts similarity index 93% rename from js/src/instrumentation/plugins/cloudflare-think-plugin.ts rename to js/src/instrumentation/providers/cloudflare-think-instrumentation.ts index 8a29eb16c..66196becf 100644 --- a/js/src/instrumentation/plugins/cloudflare-think-plugin.ts +++ b/js/src/instrumentation/providers/cloudflare-think-instrumentation.ts @@ -1,4 +1,3 @@ -import { BasePlugin } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import type { IsoChannelHandlers } from "../../isomorph"; import { @@ -25,7 +24,7 @@ import { finalizeAISDKChildTracing, patchAISDKStreamingResult, prepareAISDKAgentCallInput, -} from "./ai-sdk-plugin"; +} from "./ai-sdk-instrumentation"; import { cloudflareThinkChannels } from "./cloudflare-think-channels"; import { registerCloudflareThinkSpan, @@ -50,27 +49,15 @@ type AISDKStreamEvent = const THINK_STATE_ID = Symbol.for("braintrust.cloudflare-think.state-id"); -export class CloudflareThinkPlugin extends BasePlugin { +class CloudflareThinkInstrumentationConsumer { private readonly statesBySpanId = new Map(); - protected onEnable(): void { + public register(): void { this.subscribeToThinkRuns(); this.subscribeToAISDKStreamTextSync(); this.subscribeToAISDKStreamTextAsync(); } - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - - for (const state of this.statesBySpanId.values()) { - this.finishState(state); - } - this.statesBySpanId.clear(); - } - private subscribeToThinkRuns(): void { const channel = cloudflareThinkChannels.runInferenceLoop.tracingChannel(); const states = new WeakMap(); @@ -137,9 +124,6 @@ export class CloudflareThinkPlugin extends BasePlugin { ? contextManager.wrapSpanForStore(runState.span) : currentSpanStore.getStore(); }); - this.unsubscribers.push(() => - channel.start?.unbindStore(currentSpanStore), - ); } const handlers: IsoChannelHandlers< @@ -166,7 +150,6 @@ export class CloudflareThinkPlugin extends BasePlugin { }; channel.subscribe(handlers); - this.unsubscribers.push(() => channel.unsubscribe(handlers)); } private subscribeToAISDKStreamTextSync(): void { @@ -186,7 +169,6 @@ export class CloudflareThinkPlugin extends BasePlugin { }; channel.subscribe(handlers); - this.unsubscribers.push(() => channel.unsubscribe(handlers)); } private subscribeToAISDKStreamTextAsync(): void { @@ -206,7 +188,6 @@ export class CloudflareThinkPlugin extends BasePlugin { }; channel.subscribe(handlers); - this.unsubscribers.push(() => channel.unsubscribe(handlers)); } private startAISDKStream(event: AISDKStreamEvent): void { @@ -382,3 +363,13 @@ function extractThinkTaskInput(input: unknown): unknown { } return input; } + +let cloudflareThinkInstrumentationConsumer: + | CloudflareThinkInstrumentationConsumer + | undefined; + +export function registerCloudflareThinkInstrumentation(): void { + cloudflareThinkInstrumentationConsumer ??= + new CloudflareThinkInstrumentationConsumer(); + cloudflareThinkInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/cohere-channels.ts b/js/src/instrumentation/providers/cohere-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/cohere-channels.ts rename to js/src/instrumentation/providers/cohere-channels.ts diff --git a/js/src/instrumentation/plugins/cohere-plugin.test.ts b/js/src/instrumentation/providers/cohere-instrumentation.test.ts similarity index 99% rename from js/src/instrumentation/plugins/cohere-plugin.test.ts rename to js/src/instrumentation/providers/cohere-instrumentation.test.ts index c2ea4fbf1..d7bd3603c 100644 --- a/js/src/instrumentation/plugins/cohere-plugin.test.ts +++ b/js/src/instrumentation/providers/cohere-instrumentation.test.ts @@ -3,7 +3,7 @@ import { aggregateCohereChatStreamChunks, extractCohereResponseMetadata, parseCohereMetricsFromUsage, -} from "./cohere-plugin"; +} from "./cohere-instrumentation"; describe("parseCohereMetricsFromUsage", () => { it("parses usage and meta token counters", () => { diff --git a/js/src/instrumentation/plugins/cohere-plugin.ts b/js/src/instrumentation/providers/cohere-instrumentation.ts similarity index 89% rename from js/src/instrumentation/plugins/cohere-plugin.ts rename to js/src/instrumentation/providers/cohere-instrumentation.ts index afaea628c..14e21179c 100644 --- a/js/src/instrumentation/plugins/cohere-plugin.ts +++ b/js/src/instrumentation/providers/cohere-instrumentation.ts @@ -1,8 +1,6 @@ -import { BasePlugin } from "../core"; import { traceAsyncChannel, traceStreamingChannel, - unsubscribeAll, } from "../core/channel-tracing"; import { SpanTypeAttribute, isObject } from "../../../util/index"; import { processInputAttachments } from "../../wrappers/attachment-utils"; @@ -16,79 +14,67 @@ import type { CohereUsageLike, } from "../../vendor-sdk-types/cohere"; -export class CoherePlugin extends BasePlugin { - protected onEnable(): void { +class CohereInstrumentationConsumer { + public register(): void { this.subscribeToCohereChannels(); } - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } - private subscribeToCohereChannels(): void { - this.unsubscribers.push( - traceStreamingChannel(cohereChannels.chat, { - name: "cohere.chat", - type: SpanTypeAttribute.LLM, - extractInput: extractChatInputWithMetadata, - extractOutput: (result) => extractCohereChatOutput(result), - extractMetadata: (result) => extractCohereResponseMetadata(result), - extractMetrics: (result, startTime) => { - const metrics = parseCohereMetricsFromUsage(result); - if (startTime) { - metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; - } - return metrics; - }, - }), - ); + traceStreamingChannel(cohereChannels.chat, { + name: "cohere.chat", + type: SpanTypeAttribute.LLM, + extractInput: extractChatInputWithMetadata, + extractOutput: (result) => extractCohereChatOutput(result), + extractMetadata: (result) => extractCohereResponseMetadata(result), + extractMetrics: (result, startTime) => { + const metrics = parseCohereMetricsFromUsage(result); + if (startTime) { + metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; + } + return metrics; + }, + }); - this.unsubscribers.push( - traceStreamingChannel(cohereChannels.chatStream, { - name: "cohere.chatStream", - type: SpanTypeAttribute.LLM, - extractInput: extractChatInputWithMetadata, - extractOutput: () => undefined, - extractMetadata: () => undefined, - extractMetrics: () => ({}), - aggregateChunks: aggregateCohereChatStreamChunks, - }), - ); + traceStreamingChannel(cohereChannels.chatStream, { + name: "cohere.chatStream", + type: SpanTypeAttribute.LLM, + extractInput: extractChatInputWithMetadata, + extractOutput: () => undefined, + extractMetadata: () => undefined, + extractMetrics: () => ({}), + aggregateChunks: aggregateCohereChatStreamChunks, + }); - this.unsubscribers.push( - traceAsyncChannel(cohereChannels.embed, { - name: "cohere.embed", - type: SpanTypeAttribute.LLM, - extractInput: extractEmbedInputWithMetadata, - extractOutput: extractCohereEmbeddingOutput, - extractMetadata: (result) => extractCohereResponseMetadata(result), - extractMetrics: (result) => parseCohereMetricsFromUsage(result), - }), - ); + traceAsyncChannel(cohereChannels.embed, { + name: "cohere.embed", + type: SpanTypeAttribute.LLM, + extractInput: extractEmbedInputWithMetadata, + extractOutput: extractCohereEmbeddingOutput, + extractMetadata: (result) => extractCohereResponseMetadata(result), + extractMetrics: (result) => parseCohereMetricsFromUsage(result), + }); - this.unsubscribers.push( - traceAsyncChannel(cohereChannels.rerank, { - name: "cohere.rerank", - type: SpanTypeAttribute.LLM, - extractInput: extractRerankInputWithMetadata, - extractOutput: (result) => { - if (!isObject(result) || !Array.isArray(result.results)) { - return undefined; - } - - return result.results.slice(0, 100).map((item) => ({ - index: isObject(item) ? item.index : undefined, - relevance_score: isObject(item) - ? ((typeof item.relevanceScore === "number" - ? item.relevanceScore - : item.relevance_score) ?? null) - : null, - })); - }, - extractMetadata: (result) => extractCohereResponseMetadata(result), - extractMetrics: (result) => parseCohereMetricsFromUsage(result), - }), - ); + traceAsyncChannel(cohereChannels.rerank, { + name: "cohere.rerank", + type: SpanTypeAttribute.LLM, + extractInput: extractRerankInputWithMetadata, + extractOutput: (result) => { + if (!isObject(result) || !Array.isArray(result.results)) { + return undefined; + } + + return result.results.slice(0, 100).map((item) => ({ + index: isObject(item) ? item.index : undefined, + relevance_score: isObject(item) + ? ((typeof item.relevanceScore === "number" + ? item.relevanceScore + : item.relevance_score) ?? null) + : null, + })); + }, + extractMetadata: (result) => extractCohereResponseMetadata(result), + extractMetrics: (result) => parseCohereMetricsFromUsage(result), + }); } } @@ -949,3 +935,10 @@ export function aggregateCohereChatStreamChunks( output, }; } + +let cohereInstrumentationConsumer: CohereInstrumentationConsumer | undefined; + +export function registerCohereInstrumentation(): void { + cohereInstrumentationConsumer ??= new CohereInstrumentationConsumer(); + cohereInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/cursor-sdk-channels.ts b/js/src/instrumentation/providers/cursor-sdk-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/cursor-sdk-channels.ts rename to js/src/instrumentation/providers/cursor-sdk-channels.ts diff --git a/js/src/instrumentation/plugins/cursor-sdk-plugin.test.ts b/js/src/instrumentation/providers/cursor-sdk-instrumentation.test.ts similarity index 94% rename from js/src/instrumentation/plugins/cursor-sdk-plugin.test.ts rename to js/src/instrumentation/providers/cursor-sdk-instrumentation.test.ts index b22052c8c..c30169c35 100644 --- a/js/src/instrumentation/plugins/cursor-sdk-plugin.test.ts +++ b/js/src/instrumentation/providers/cursor-sdk-instrumentation.test.ts @@ -15,11 +15,11 @@ vi.mock("../../logger", () => ({ })); import iso from "../../isomorph"; -import { CursorSDKPlugin } from "./cursor-sdk-plugin"; +import { registerCursorSDKInstrumentation } from "./cursor-sdk-instrumentation"; const mockNewTracingChannel = iso.newTracingChannel as ReturnType; -describe("CursorSDKPlugin", () => { +describe("registerCursorSDKInstrumentation", () => { let handlersByName: Map; let spans: Array<{ end: ReturnType; @@ -34,7 +34,6 @@ describe("CursorSDKPlugin", () => { mockNewTracingChannel.mockImplementation((name: string) => ({ subscribe: vi.fn((handlers) => handlersByName.set(name, handlers)), tracePromise: vi.fn((fn) => fn()), - unsubscribe: vi.fn(), })); mockStartSpan.mockImplementation((args: any) => { const span = { @@ -56,9 +55,7 @@ describe("CursorSDKPlugin", () => { }); it("subscribes to Cursor SDK channels", () => { - const plugin = new CursorSDKPlugin(); - - plugin.enable(); + registerCursorSDKInstrumentation(); expect(handlersByName.has("orchestrion:@cursor/sdk:Agent.create")).toBe( true, @@ -73,8 +70,7 @@ describe("CursorSDKPlugin", () => { }); it("patches agents returned by Agent.create and traces send/wait", async () => { - const plugin = new CursorSDKPlugin(); - plugin.enable(); + registerCursorSDKInstrumentation(); const createHandlers = handlersByName.get( "orchestrion:@cursor/sdk:Agent.create", @@ -131,8 +127,7 @@ describe("CursorSDKPlugin", () => { }); it("captures stream tool calls and usage", async () => { - const plugin = new CursorSDKPlugin(); - plugin.enable(); + registerCursorSDKInstrumentation(); const sendHandlers = handlersByName.get( "orchestrion:@cursor/sdk:agent.send", @@ -214,8 +209,7 @@ describe("CursorSDKPlugin", () => { }); it("traces Agent.prompt without a nested send span", () => { - const plugin = new CursorSDKPlugin(); - plugin.enable(); + registerCursorSDKInstrumentation(); const promptHandlers = handlersByName.get( "orchestrion:@cursor/sdk:Agent.prompt", diff --git a/js/src/instrumentation/plugins/cursor-sdk-plugin.ts b/js/src/instrumentation/providers/cursor-sdk-instrumentation.ts similarity index 98% rename from js/src/instrumentation/plugins/cursor-sdk-plugin.ts rename to js/src/instrumentation/providers/cursor-sdk-instrumentation.ts index f01db5987..41118be17 100644 --- a/js/src/instrumentation/plugins/cursor-sdk-plugin.ts +++ b/js/src/instrumentation/providers/cursor-sdk-instrumentation.ts @@ -1,4 +1,4 @@ -import { BasePlugin, toLoggedError } from "../core"; +import { toLoggedError } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import type { IsoChannelHandlers } from "../../isomorph"; import { debugLogger } from "../../debug-logger"; @@ -64,23 +64,15 @@ type PromptState = { startTime: number; }; -export class CursorSDKPlugin extends BasePlugin { +class CursorSDKInstrumentationConsumer { private promptDepth = 0; - protected onEnable(): void { + public register(): void { this.subscribeToAgentFactories(); this.subscribeToPrompt(); this.subscribeToSend(); } - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - this.promptDepth = 0; - } - private subscribeToAgentFactories(): void { this.subscribeToAgentFactory(cursorSDKChannels.create); this.subscribeToAgentFactory(cursorSDKChannels.resume); @@ -98,9 +90,6 @@ export class CursorSDKPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - tracingChannel.unsubscribe(handlers); - }); } private subscribeToPrompt(): void { @@ -171,9 +160,6 @@ export class CursorSDKPlugin extends BasePlugin { }; channel.subscribe(handlers); - this.unsubscribers.push(() => { - channel.unsubscribe(handlers); - }); } private subscribeToSend(): void { @@ -267,9 +253,6 @@ export class CursorSDKPlugin extends BasePlugin { }; channel.subscribe(handlers); - this.unsubscribers.push(() => { - channel.unsubscribe(handlers); - }); } } @@ -1185,3 +1168,12 @@ function cleanMetrics(metrics: Record): Record { } return cleaned; } + +let cursorSDKInstrumentationConsumer: + | CursorSDKInstrumentationConsumer + | undefined; + +export function registerCursorSDKInstrumentation(): void { + cursorSDKInstrumentationConsumer ??= new CursorSDKInstrumentationConsumer(); + cursorSDKInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/eve-plugin.test.ts b/js/src/instrumentation/providers/eve-instrumentation.test.ts similarity index 99% rename from js/src/instrumentation/plugins/eve-plugin.test.ts rename to js/src/instrumentation/providers/eve-instrumentation.test.ts index 38e15bef0..9e1a6bc28 100644 --- a/js/src/instrumentation/plugins/eve-plugin.test.ts +++ b/js/src/instrumentation/providers/eve-instrumentation.test.ts @@ -11,7 +11,7 @@ import * as braintrustExports from "../../exports"; import { configureNode } from "../../node/config"; import { Attachment, _exportsForTestingOnly, initLogger } from "../../logger"; import * as instrumentationExports from "../index"; -import { braintrustEveInstrumentation } from "./eve-plugin"; +import { braintrustEveInstrumentation } from "./eve-instrumentation"; import type { EveInstrumentationAttemptScope, EveInstrumentationHandlerContext, @@ -60,7 +60,7 @@ describe("braintrustEveInstrumentation", () => { beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ - projectName: "eve-plugin.test.ts", + projectName: "eve-instrumentation.test.ts", projectId: "test-project-id", }); }); @@ -107,7 +107,7 @@ describe("braintrustEveInstrumentation", () => { it("records a turn with LLM and action spans", async () => { const provider = braintrustEveInstrumentation({ metadata: { - scenario: "eve-plugin-unit", + scenario: "eve-instrumentation-unit", testRunId: "test-run-flat-tree", }, }); @@ -272,7 +272,7 @@ describe("braintrustEveInstrumentation", () => { input: [{ content: "Search then read", role: "user" }], metadata: { "eve.session_id": "session-root", - scenario: "eve-plugin-unit", + scenario: "eve-instrumentation-unit", testRunId: "test-run-flat-tree", }, output: "Here is the Eve instrumentation guide.", @@ -300,7 +300,7 @@ describe("braintrustEveInstrumentation", () => { "eve.session_id": "session-root", model: "qwen/qwen3-30b-a3b", provider: "openrouter", - scenario: "eve-plugin-unit", + scenario: "eve-instrumentation-unit", testRunId: "test-run-flat-tree", }, span_attributes: { @@ -345,7 +345,7 @@ describe("braintrustEveInstrumentation", () => { input: { query: "Eve instrumentation" }, metadata: { "eve.session_id": "session-root", - scenario: "eve-plugin-unit", + scenario: "eve-instrumentation-unit", testRunId: "test-run-flat-tree", }, output: { hits: ["eve.dev/docs"] }, diff --git a/js/src/instrumentation/plugins/eve-plugin.ts b/js/src/instrumentation/providers/eve-instrumentation.ts similarity index 100% rename from js/src/instrumentation/plugins/eve-plugin.ts rename to js/src/instrumentation/providers/eve-instrumentation.ts diff --git a/js/src/instrumentation/plugins/flue-plugin.test.ts b/js/src/instrumentation/providers/flue-instrumentation.test.ts similarity index 99% rename from js/src/instrumentation/plugins/flue-plugin.test.ts rename to js/src/instrumentation/providers/flue-instrumentation.test.ts index 04901f8cc..ae6fbec8b 100644 --- a/js/src/instrumentation/plugins/flue-plugin.test.ts +++ b/js/src/instrumentation/providers/flue-instrumentation.test.ts @@ -64,7 +64,7 @@ vi.mock("../../logger", () => ({ }, })); -import { braintrustFlueInstrumentation } from "./flue-plugin"; +import { braintrustFlueInstrumentation } from "./flue-instrumentation"; describe("Flue instrumentation", () => { let spans: Array<{ diff --git a/js/src/instrumentation/plugins/flue-plugin.ts b/js/src/instrumentation/providers/flue-instrumentation.ts similarity index 100% rename from js/src/instrumentation/plugins/flue-plugin.ts rename to js/src/instrumentation/providers/flue-instrumentation.ts diff --git a/js/src/instrumentation/plugins/genkit-channels.ts b/js/src/instrumentation/providers/genkit-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/genkit-channels.ts rename to js/src/instrumentation/providers/genkit-channels.ts diff --git a/js/src/instrumentation/plugins/genkit-plugin.test.ts b/js/src/instrumentation/providers/genkit-instrumentation.test.ts similarity index 91% rename from js/src/instrumentation/plugins/genkit-plugin.test.ts rename to js/src/instrumentation/providers/genkit-instrumentation.test.ts index dcf41c0f7..11f172e5b 100644 --- a/js/src/instrumentation/plugins/genkit-plugin.test.ts +++ b/js/src/instrumentation/providers/genkit-instrumentation.test.ts @@ -1,6 +1,6 @@ import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; import { _exportsForTestingOnly, initLogger } from "../../logger"; -import { GenkitPlugin } from "./genkit-plugin"; +import { registerGenkitInstrumentation } from "./genkit-instrumentation"; import { genkitChannels } from "./genkit-channels"; function singleQueueStream( @@ -35,9 +35,7 @@ async function collectAsync(stream: AsyncIterable): Promise { return chunks; } -describe("GenkitPlugin stream patching", () => { - const plugin = new GenkitPlugin(); - +describe("registerGenkitInstrumentation stream patching", () => { beforeAll(async () => { await _exportsForTestingOnly.simulateLoginForTests(); }); @@ -45,18 +43,17 @@ describe("GenkitPlugin stream patching", () => { beforeEach(() => { _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ - projectName: "genkit-plugin.test.ts", + projectName: "genkit-instrumentation.test.ts", projectId: "test-project-id", }); }); afterEach(() => { - plugin.disable(); _exportsForTestingOnly.clearTestBackgroundLogger(); }); it("does not consume generateStream chunks before user code reads them", async () => { - plugin.enable(); + registerGenkitInstrumentation(); const stream = singleQueueStream([{ text: "hello" }, { text: " world" }]); const result = genkitChannels.generateStream.traceSync( @@ -85,7 +82,7 @@ describe("GenkitPlugin stream patching", () => { }); it("does not consume action.stream chunks before user code reads them", async () => { - plugin.enable(); + registerGenkitInstrumentation(); const stream = singleQueueStream(["first", "second"]); const action = Object.assign(() => Promise.resolve(), { __action: { diff --git a/js/src/instrumentation/plugins/genkit-plugin.ts b/js/src/instrumentation/providers/genkit-instrumentation.ts similarity index 85% rename from js/src/instrumentation/plugins/genkit-plugin.ts rename to js/src/instrumentation/providers/genkit-instrumentation.ts index 90edb66e4..e485295ea 100644 --- a/js/src/instrumentation/plugins/genkit-plugin.ts +++ b/js/src/instrumentation/providers/genkit-instrumentation.ts @@ -1,8 +1,7 @@ -import { BasePlugin, toLoggedError } from "../core"; +import { toLoggedError } from "../core"; import { traceAsyncChannel, traceSyncStreamChannel, - unsubscribeAll, } from "../core/channel-tracing"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; import type { ChannelMessage } from "../core/channel-definitions"; @@ -38,61 +37,49 @@ type SpanState = { startTime: number; }; -export class GenkitPlugin extends BasePlugin { - protected onEnable(): void { +class GenkitInstrumentationConsumer { + public register(): void { this.subscribeToGenkitChannels(); } - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } - private subscribeToGenkitChannels(): void { - this.unsubscribers.push( - traceAsyncChannel(genkitChannels.generate, { - name: "genkit.generate", - type: SpanTypeAttribute.LLM, - extractInput: ([input]) => extractGenerateInput(input), - extractOutput: extractGenerateOutput, - extractMetadata: (result, event) => - extractGenerateResponseMetadata(result, event?.arguments?.[0]), - extractMetrics: (result) => parseGenkitUsageMetrics(result?.usage), - }), - ); + traceAsyncChannel(genkitChannels.generate, { + name: "genkit.generate", + type: SpanTypeAttribute.LLM, + extractInput: ([input]) => extractGenerateInput(input), + extractOutput: extractGenerateOutput, + extractMetadata: (result, event) => + extractGenerateResponseMetadata(result, event?.arguments?.[0]), + extractMetrics: (result) => parseGenkitUsageMetrics(result?.usage), + }); - this.unsubscribers.push( - traceSyncStreamChannel(genkitChannels.generateStream, { - name: "genkit.generateStream", - type: SpanTypeAttribute.LLM, - extractInput: ([input]) => extractGenerateInput(input), - patchResult: ({ result, span, startTime }) => - patchGenerateStreamResult(result, span, startTime), - }), - ); + traceSyncStreamChannel(genkitChannels.generateStream, { + name: "genkit.generateStream", + type: SpanTypeAttribute.LLM, + extractInput: ([input]) => extractGenerateInput(input), + patchResult: ({ result, span, startTime }) => + patchGenerateStreamResult(result, span, startTime), + }); - this.unsubscribers.push( - traceAsyncChannel(genkitChannels.embed, { - name: "genkit.embed", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params]) => extractEmbedInput(params), - extractOutput: (result) => summarizeEmbeddingResult(result), - extractMetadata: (_result, event) => - extractEmbedMetadata(event?.arguments?.[0]), - extractMetrics: () => ({}), - }), - ); + traceAsyncChannel(genkitChannels.embed, { + name: "genkit.embed", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params]) => extractEmbedInput(params), + extractOutput: (result) => summarizeEmbeddingResult(result), + extractMetadata: (_result, event) => + extractEmbedMetadata(event?.arguments?.[0]), + extractMetrics: () => ({}), + }); - this.unsubscribers.push( - traceAsyncChannel(genkitChannels.embedMany, { - name: "genkit.embedMany", - type: SpanTypeAttribute.FUNCTION, - extractInput: ([params]) => extractEmbedManyInput(params), - extractOutput: summarizeEmbeddingResult, - extractMetadata: (_result, event) => - extractEmbedMetadata(event?.arguments?.[0]), - extractMetrics: () => ({}), - }), - ); + traceAsyncChannel(genkitChannels.embedMany, { + name: "genkit.embedMany", + type: SpanTypeAttribute.FUNCTION, + extractInput: ([params]) => extractEmbedManyInput(params), + extractOutput: summarizeEmbeddingResult, + extractMetadata: (_result, event) => + extractEmbedMetadata(event?.arguments?.[0]), + extractMetrics: () => ({}), + }); this.subscribeToActionRun(); this.subscribeToActionSpan(); @@ -105,10 +92,8 @@ export class GenkitPlugin extends BasePlugin { ChannelMessage >; const states = new WeakMap(); - const unbindCurrentSpanStore = bindActionCurrentSpanStoreToStart( - tracingChannel, - states, - (event) => startActionRunSpan(event), + bindActionCurrentSpanStoreToStart(tracingChannel, states, (event) => + startActionRunSpan(event), ); const handlers: IsoChannelHandlers< @@ -147,10 +132,6 @@ export class GenkitPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - unbindCurrentSpanStore?.(); - tracingChannel.unsubscribe(handlers); - }); } private subscribeToActionSpan(): void { @@ -159,10 +140,8 @@ export class GenkitPlugin extends BasePlugin { ChannelMessage >; const states = new WeakMap(); - const unbindCurrentSpanStore = bindActionCurrentSpanStoreToStart( - tracingChannel, - states, - (event) => startActionSpan(event), + bindActionCurrentSpanStoreToStart(tracingChannel, states, (event) => + startActionSpan(event), ); const handlers: IsoChannelHandlers< @@ -202,25 +181,19 @@ export class GenkitPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - unbindCurrentSpanStore?.(); - tracingChannel.unsubscribe(handlers); - }); } private subscribeToActionStream(): void { - this.unsubscribers.push( - traceSyncStreamChannel(genkitChannels.actionStream, { - name: "genkit.action.stream", - type: SpanTypeAttribute.TASK, - extractInput: ([input], event) => ({ - input, - metadata: actionMetadataForLog(extractActionMetadata(event.self)), - }), - patchResult: ({ result, span, startTime }) => - patchActionStreamResult(result, span, startTime), + traceSyncStreamChannel(genkitChannels.actionStream, { + name: "genkit.action.stream", + type: SpanTypeAttribute.TASK, + extractInput: ([input], event) => ({ + input, + metadata: actionMetadataForLog(extractActionMetadata(event.self)), }), - ); + patchResult: ({ result, span, startTime }) => + patchActionStreamResult(result, span, startTime), + }); } } @@ -307,7 +280,7 @@ function bindActionCurrentSpanStoreToStart< tracingChannel: IsoTracingChannel>, states: WeakMap, create: (event: ChannelMessage) => SpanState | undefined, -): (() => void) | undefined { +): void { const state = _internalGetGlobalState(); const contextManager = state?.contextManager; const startChannel = tracingChannel.start as @@ -316,7 +289,6 @@ function bindActionCurrentSpanStoreToStart< store: CurrentSpanStore, callback: (event: ChannelMessage) => unknown, ) => void; - unbindStore?: (store: CurrentSpanStore) => void; } & object) | undefined; const currentSpanStore = contextManager @@ -328,7 +300,7 @@ function bindActionCurrentSpanStoreToStart< : undefined; if (!startChannel?.bindStore || !currentSpanStore) { - return undefined; + return; } startChannel.bindStore(currentSpanStore, (event) => { @@ -339,10 +311,6 @@ function bindActionCurrentSpanStoreToStart< ? contextManager!.wrapSpanForStore(state.span) : currentSpanStore.getStore(); }); - - return () => { - startChannel.unbindStore?.(currentSpanStore); - }; } function normalizeInput(input: GenkitGenerateInput): GenkitGenerateInput { @@ -771,3 +739,10 @@ function pickNumberMetrics( function stringValue(value: unknown): string | undefined { return typeof value === "string" ? value : undefined; } + +let genkitInstrumentationConsumer: GenkitInstrumentationConsumer | undefined; + +export function registerGenkitInstrumentation(): void { + genkitInstrumentationConsumer ??= new GenkitInstrumentationConsumer(); + genkitInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/github-copilot-channels.ts b/js/src/instrumentation/providers/github-copilot-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/github-copilot-channels.ts rename to js/src/instrumentation/providers/github-copilot-channels.ts diff --git a/js/src/instrumentation/plugins/github-copilot-plugin.test.ts b/js/src/instrumentation/providers/github-copilot-instrumentation.test.ts similarity index 79% rename from js/src/instrumentation/plugins/github-copilot-plugin.test.ts rename to js/src/instrumentation/providers/github-copilot-instrumentation.test.ts index 17ecb38d8..5e88b3b51 100644 --- a/js/src/instrumentation/plugins/github-copilot-plugin.test.ts +++ b/js/src/instrumentation/providers/github-copilot-instrumentation.test.ts @@ -1,8 +1,8 @@ import { describe, expect, it } from "vitest"; import { - GitHubCopilotPlugin, + registerGitHubCopilotInstrumentation, extractMetricsFromUsage, -} from "./github-copilot-plugin"; +} from "./github-copilot-instrumentation"; describe("extractMetricsFromUsage", () => { it("maps input/output tokens to standard metric keys", () => { @@ -119,38 +119,8 @@ describe("extractMetricsFromUsage", () => { }); }); -describe("GitHubCopilotPlugin lifecycle", () => { - it("enables without throwing", () => { - const plugin = new GitHubCopilotPlugin(); - expect(() => plugin.enable()).not.toThrow(); - plugin.disable(); - }); - - it("disables without throwing", () => { - const plugin = new GitHubCopilotPlugin(); - plugin.enable(); - expect(() => plugin.disable()).not.toThrow(); - }); - - it("enable is idempotent", () => { - const plugin = new GitHubCopilotPlugin(); - plugin.enable(); - expect(() => plugin.enable()).not.toThrow(); - plugin.disable(); - }); - - it("disable is idempotent", () => { - const plugin = new GitHubCopilotPlugin(); - plugin.enable(); - plugin.disable(); - expect(() => plugin.disable()).not.toThrow(); - }); - - it("can be re-enabled after disable", () => { - const plugin = new GitHubCopilotPlugin(); - plugin.enable(); - plugin.disable(); - expect(() => plugin.enable()).not.toThrow(); - plugin.disable(); +describe("registerGitHubCopilotInstrumentation", () => { + it("registers without throwing", () => { + expect(() => registerGitHubCopilotInstrumentation()).not.toThrow(); }); }); diff --git a/js/src/instrumentation/plugins/github-copilot-plugin.ts b/js/src/instrumentation/providers/github-copilot-instrumentation.ts similarity index 97% rename from js/src/instrumentation/plugins/github-copilot-plugin.ts rename to js/src/instrumentation/providers/github-copilot-instrumentation.ts index 8fa156cfc..b5f14976f 100644 --- a/js/src/instrumentation/plugins/github-copilot-plugin.ts +++ b/js/src/instrumentation/providers/github-copilot-instrumentation.ts @@ -1,4 +1,3 @@ -import { BasePlugin } from "../core"; import type { IsoChannelHandlers } from "../../isomorph"; import { startSpan as startBaseSpan } from "../../logger"; import type { Span } from "../../logger"; @@ -672,7 +671,7 @@ function isGitHubCopilotSession(value: unknown): value is GitHubCopilotSession { } // --------------------------------------------------------------------------- -// Plugin and handler factory +// Consumer and handler factory // --------------------------------------------------------------------------- // eslint-disable-next-line @typescript-eslint/no-explicit-any @@ -756,18 +755,11 @@ function makeSessionHandlers( }; } -export class GitHubCopilotPlugin extends BasePlugin { - protected onEnable(): void { +class GitHubCopilotInstrumentationConsumer { + public register(): void { this.subscribeToSessionChannels(); } - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - } - private subscribeToSessionChannels(): void { const createChannel = gitHubCopilotChannels.createSession.tracingChannel(); const resumeChannel = gitHubCopilotChannels.resumeSession.tracingChannel(); @@ -787,10 +779,15 @@ export class GitHubCopilotPlugin extends BasePlugin { createChannel.subscribe(createHandlers); resumeChannel.subscribe(resumeHandlers); - - this.unsubscribers.push( - () => createChannel.unsubscribe(createHandlers), - () => resumeChannel.unsubscribe(resumeHandlers), - ); } } + +let gitHubCopilotInstrumentationConsumer: + | GitHubCopilotInstrumentationConsumer + | undefined; + +export function registerGitHubCopilotInstrumentation(): void { + gitHubCopilotInstrumentationConsumer ??= + new GitHubCopilotInstrumentationConsumer(); + gitHubCopilotInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/google-adk-channels.ts b/js/src/instrumentation/providers/google-adk-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/google-adk-channels.ts rename to js/src/instrumentation/providers/google-adk-channels.ts diff --git a/js/src/instrumentation/plugins/google-adk-plugin.test.ts b/js/src/instrumentation/providers/google-adk-instrumentation.test.ts similarity index 91% rename from js/src/instrumentation/plugins/google-adk-plugin.test.ts rename to js/src/instrumentation/providers/google-adk-instrumentation.test.ts index 073af75f4..f4efcfc7d 100644 --- a/js/src/instrumentation/plugins/google-adk-plugin.test.ts +++ b/js/src/instrumentation/providers/google-adk-instrumentation.test.ts @@ -15,7 +15,7 @@ vi.mock("../../isomorph", () => ({ }, })); -import { GoogleADKPlugin } from "./google-adk-plugin"; +import { registerGoogleADKInstrumentation } from "./google-adk-instrumentation"; import iso from "../../isomorph"; const mockNewTracingChannel = iso.newTracingChannel as ReturnType; @@ -44,26 +44,19 @@ vi.mock("../../logger", () => ({ }, })); -describe("GoogleADKPlugin", () => { - let plugin: GoogleADKPlugin; +describe("registerGoogleADKInstrumentation", () => { let mockChannel: any; let subscribeSpy: any; - let unsubscribeSpy: any; let bindStoreSpy: any; - let unbindStoreSpy: any; beforeEach(() => { subscribeSpy = vi.fn(); - unsubscribeSpy = vi.fn(); bindStoreSpy = vi.fn(); - unbindStoreSpy = vi.fn(); mockChannel = { subscribe: subscribeSpy, - unsubscribe: unsubscribeSpy, hasSubscribers: false, start: { bindStore: bindStoreSpy, - unbindStore: unbindStoreSpy, }, }; @@ -71,16 +64,15 @@ describe("GoogleADKPlugin", () => { mockStartSpan.mockClear(); mockInternalGetGlobalState.mockReset(); mockInternalGetGlobalState.mockReturnValue(undefined); - plugin = new GoogleADKPlugin(); }); afterEach(() => { vi.clearAllMocks(); }); - describe("enable/disable lifecycle", () => { - it("should subscribe to channels when enabled", () => { - plugin.enable(); + describe("registration", () => { + it("subscribes to channels", () => { + registerGoogleADKInstrumentation(); // Should subscribe to 3 channels: runner.runAsync, agent.runAsync, tool.runAsync expect(mockNewTracingChannel).toHaveBeenCalledWith( @@ -94,43 +86,11 @@ describe("GoogleADKPlugin", () => { ); expect(subscribeSpy).toHaveBeenCalledTimes(3); }); - - it("should not subscribe multiple times if enabled twice", () => { - plugin.enable(); - const firstCallCount = subscribeSpy.mock.calls.length; - - plugin.enable(); - const secondCallCount = subscribeSpy.mock.calls.length; - - expect(firstCallCount).toBe(secondCallCount); - }); - - it("should unsubscribe from channels when disabled", () => { - plugin.enable(); - plugin.disable(); - - expect(unsubscribeSpy).toHaveBeenCalled(); - }); - - it("should clear unsubscribers array after disable", () => { - plugin.enable(); - plugin.disable(); - - // Enable again should re-subscribe - subscribeSpy.mockClear(); - plugin.enable(); - - expect(subscribeSpy).toHaveBeenCalledTimes(3); - }); - - it("should not crash when disabled without being enabled", () => { - expect(() => plugin.disable()).not.toThrow(); - }); }); describe("runner.runAsync channel", () => { it("should create a TASK span with runner metadata on start", () => { - plugin.enable(); + registerGoogleADKInstrumentation(); // Find the first subscribe call (runner channel) const handlers = subscribeSpy.mock.calls[0][0]; @@ -163,7 +123,7 @@ describe("GoogleADKPlugin", () => { }); it("should handle stream end with async iterable result", () => { - plugin.enable(); + registerGoogleADKInstrumentation(); const handlers = subscribeSpy.mock.calls[0][0]; const event: any = { @@ -196,7 +156,7 @@ describe("GoogleADKPlugin", () => { }); it("should handle error events", () => { - plugin.enable(); + registerGoogleADKInstrumentation(); const handlers = subscribeSpy.mock.calls[0][0]; const event: any = { @@ -228,7 +188,7 @@ describe("GoogleADKPlugin", () => { }, } as any); - plugin.enable(); + registerGoogleADKInstrumentation(); expect(bindStoreSpy).toHaveBeenNthCalledWith( 1, @@ -367,7 +327,7 @@ describe("GoogleADKPlugin", () => { ])( "normalizes $name", async ({ usageMetadata, expectedMetrics, absentMetrics }) => { - plugin.enable(); + registerGoogleADKInstrumentation(); const handlers = subscribeSpy.mock.calls[0][0]; const event: any = { @@ -398,7 +358,7 @@ describe("GoogleADKPlugin", () => { ); it("preserves explicitly reported zero usage", async () => { - plugin.enable(); + registerGoogleADKInstrumentation(); const handlers = subscribeSpy.mock.calls[0][0]; const event: any = { @@ -447,7 +407,7 @@ describe("GoogleADKPlugin", () => { describe("agent.runAsync channel", () => { it("should create a TASK span with agent metadata on start", () => { - plugin.enable(); + registerGoogleADKInstrumentation(); // Agent channel is the second subscribe call const handlers = subscribeSpy.mock.calls[1][0]; @@ -474,7 +434,7 @@ describe("GoogleADKPlugin", () => { }); it("uses the invoked agent instance for names when parent context still points at a parent agent", () => { - plugin.enable(); + registerGoogleADKInstrumentation(); const handlers = subscribeSpy.mock.calls[1][0]; const event = { @@ -511,7 +471,7 @@ describe("GoogleADKPlugin", () => { }); it("should handle agent without a name gracefully", () => { - plugin.enable(); + registerGoogleADKInstrumentation(); const handlers = subscribeSpy.mock.calls[1][0]; const event = { @@ -537,7 +497,7 @@ describe("GoogleADKPlugin", () => { }, } as any); - plugin.enable(); + registerGoogleADKInstrumentation(); expect(bindStoreSpy).toHaveBeenNthCalledWith( 2, @@ -588,7 +548,7 @@ describe("GoogleADKPlugin", () => { .mockImplementationOnce(() => runnerSpan) .mockImplementationOnce(() => agentSpan); - plugin.enable(); + registerGoogleADKInstrumentation(); const runnerHandlers = subscribeSpy.mock.calls[0][0]; const agentHandlers = subscribeSpy.mock.calls[1][0]; @@ -638,7 +598,7 @@ describe("GoogleADKPlugin", () => { describe("tool.runAsync channel", () => { it("should create a TOOL span on start", () => { - plugin.enable(); + registerGoogleADKInstrumentation(); // Tool channel is the third subscribe call const handlers = subscribeSpy.mock.calls[2][0]; @@ -675,7 +635,7 @@ describe("GoogleADKPlugin", () => { }); it("should log output and metrics on asyncEnd", () => { - plugin.enable(); + registerGoogleADKInstrumentation(); const handlers = subscribeSpy.mock.calls[2][0]; const event: any = { @@ -706,7 +666,7 @@ describe("GoogleADKPlugin", () => { }); it("should handle tool execution errors", () => { - plugin.enable(); + registerGoogleADKInstrumentation(); const handlers = subscribeSpy.mock.calls[2][0]; const event: any = { diff --git a/js/src/instrumentation/plugins/google-adk-plugin.ts b/js/src/instrumentation/providers/google-adk-instrumentation.ts similarity index 95% rename from js/src/instrumentation/plugins/google-adk-plugin.ts rename to js/src/instrumentation/providers/google-adk-instrumentation.ts index 623264016..994779305 100644 --- a/js/src/instrumentation/plugins/google-adk-plugin.ts +++ b/js/src/instrumentation/providers/google-adk-instrumentation.ts @@ -1,4 +1,3 @@ -import { BasePlugin } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; import type { IsoChannelHandlers, IsoTracingChannel } from "../../isomorph"; @@ -51,9 +50,9 @@ type GoogleADKStreamChannel = | typeof googleADKChannels.agentRunAsync; /** - * Auto-instrumentation plugin for the Google ADK. + * Internal auto-instrumentation consumer for the Google ADK. * - * This plugin subscribes to orchestrion channels for Google ADK methods + * This consumer subscribes to orchestrion channels for Google ADK methods * and creates Braintrust spans to track: * - Runner.runAsync — top-level agent execution (TASK span) * - BaseAgent.runAsync — individual agent invocations (TASK span) @@ -62,25 +61,16 @@ type GoogleADKStreamChannel = * LLM calls made through ADK are automatically captured by the existing * @google/genai instrumentation since ADK uses GenAI internally. */ -export class GoogleADKPlugin extends BasePlugin { +class GoogleADKInstrumentationConsumer { private activeRunnerSpans = new Map(); private activeAgentSpans = new Map(); - protected onEnable(): void { + public register(): void { this.subscribeToRunnerRunAsync(); this.subscribeToAgentRunAsync(); this.subscribeToToolRunAsync(); } - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - this.activeRunnerSpans.clear(); - this.activeAgentSpans.clear(); - } - private subscribeToRunnerRunAsync(): void { const tracingChannel = googleADKChannels.runnerRunAsync.tracingChannel() as IsoTracingChannel< @@ -126,11 +116,7 @@ export class GoogleADKPlugin extends BasePlugin { return { span, startTime, events: [], contextKey }; }; - const unbindCurrentSpanStore = bindCurrentSpanStoreToStart( - tracingChannel, - states, - createState, - ); + bindCurrentSpanStoreToStart(tracingChannel, states, createState); const handlers: IsoChannelHandlers< ChannelMessage @@ -189,10 +175,6 @@ export class GoogleADKPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - unbindCurrentSpanStore?.(); - tracingChannel.unsubscribe(handlers); - }); } private subscribeToAgentRunAsync(): void { @@ -262,11 +244,7 @@ export class GoogleADKPlugin extends BasePlugin { return { span, startTime, events: [], contextKey, name: agentName }; }; - const unbindCurrentSpanStore = bindCurrentSpanStoreToStart( - tracingChannel, - states, - createState, - ); + bindCurrentSpanStoreToStart(tracingChannel, states, createState); const handlers: IsoChannelHandlers< ChannelMessage @@ -324,10 +302,6 @@ export class GoogleADKPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - unbindCurrentSpanStore?.(); - tracingChannel.unsubscribe(handlers); - }); } private subscribeToToolRunAsync(): void { @@ -413,9 +387,6 @@ export class GoogleADKPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - tracingChannel.unsubscribe(handlers); - }); } } @@ -536,7 +507,7 @@ function bindCurrentSpanStoreToStart< tracingChannel: IsoTracingChannel>, states: WeakMap, create: (event: ChannelMessage) => TState, -): (() => void) | undefined { +): void { const state = _internalGetGlobalState(); const contextManager = state?.contextManager; const startChannel = tracingChannel.start as @@ -545,7 +516,6 @@ function bindCurrentSpanStoreToStart< store: CurrentSpanStore, callback: (event: ChannelMessage) => unknown, ) => void; - unbindStore?: (store: CurrentSpanStore) => void; } & object) | undefined; const currentSpanStore = contextManager @@ -557,7 +527,7 @@ function bindCurrentSpanStoreToStart< : undefined; if (!startChannel?.bindStore || !currentSpanStore) { - return undefined; + return; } startChannel.bindStore(currentSpanStore, (event) => { @@ -566,10 +536,6 @@ function bindCurrentSpanStoreToStart< ).span; return contextManager.wrapSpanForStore(span); }); - - return () => { - startChannel.unbindStore?.(currentSpanStore); - }; } // ---- Helper functions ---- @@ -1026,3 +992,12 @@ function cleanMetrics(metrics: Record): Record { } return cleaned; } + +let googleADKInstrumentationConsumer: + | GoogleADKInstrumentationConsumer + | undefined; + +export function registerGoogleADKInstrumentation(): void { + googleADKInstrumentationConsumer ??= new GoogleADKInstrumentationConsumer(); + googleADKInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/google-genai-channels.ts b/js/src/instrumentation/providers/google-genai-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/google-genai-channels.ts rename to js/src/instrumentation/providers/google-genai-channels.ts diff --git a/js/src/instrumentation/plugins/google-genai-plugin.test.ts b/js/src/instrumentation/providers/google-genai-instrumentation.test.ts similarity index 95% rename from js/src/instrumentation/plugins/google-genai-plugin.test.ts rename to js/src/instrumentation/providers/google-genai-instrumentation.test.ts index f14690e73..175cfb309 100644 --- a/js/src/instrumentation/plugins/google-genai-plugin.test.ts +++ b/js/src/instrumentation/providers/google-genai-instrumentation.test.ts @@ -25,7 +25,7 @@ vi.mock("../../isomorph", () => ({ }, })); -import { GoogleGenAIPlugin } from "./google-genai-plugin"; +import { registerGoogleGenAIInstrumentation } from "./google-genai-instrumentation"; import { startSpan } from "../../logger"; import iso from "../../isomorph"; @@ -50,66 +50,27 @@ vi.mock("../../logger", () => ({ }, })); -describe("GoogleGenAIPlugin", () => { - let plugin: GoogleGenAIPlugin; +describe("registerGoogleGenAIInstrumentation", () => { let mockChannel: any; let subscribeSpy: any; - let unsubscribeSpy: any; beforeEach(() => { subscribeSpy = vi.fn(); - unsubscribeSpy = vi.fn(); mockChannel = { subscribe: subscribeSpy, - unsubscribe: unsubscribeSpy, hasSubscribers: false, }; mockNewTracingChannel.mockReturnValue(mockChannel); - plugin = new GoogleGenAIPlugin(); }); afterEach(() => { vi.clearAllMocks(); }); - describe("enable/disable lifecycle", () => { - it("should not subscribe multiple times if enabled twice", () => { - plugin.enable(); - const firstCallCount = subscribeSpy.mock.calls.length; - - plugin.enable(); - const secondCallCount = subscribeSpy.mock.calls.length; - - expect(firstCallCount).toBe(secondCallCount); - }); - - it("should unsubscribe from channels when disabled", () => { - plugin.enable(); - plugin.disable(); - - expect(unsubscribeSpy).toHaveBeenCalled(); - }); - - it("should clear unsubscribers array after disable", () => { - plugin.enable(); - plugin.disable(); - - // Enable again should re-subscribe - subscribeSpy.mockClear(); - plugin.enable(); - - expect(subscribeSpy).toHaveBeenCalled(); - }); - - it("should not crash when disabled without being enabled", () => { - expect(() => plugin.disable()).not.toThrow(); - }); - }); - describe("generateContent channel subscription", () => { it("should extract input correctly", () => { - plugin.enable(); + registerGoogleGenAIInstrumentation(); const subscribeCall = subscribeSpy.mock.calls.find( (call: any) => @@ -199,7 +160,7 @@ describe("GoogleGenAIPlugin", () => { ])( "normalizes $name", ({ usageMetadata, expectedMetrics, absentMetrics }) => { - plugin.enable(); + registerGoogleGenAIInstrumentation(); const handlers = subscribeSpy.mock.calls[0][0]; const event: any = { @@ -227,7 +188,7 @@ describe("GoogleGenAIPlugin", () => { ); it("preserves explicitly reported zero usage", () => { - plugin.enable(); + registerGoogleGenAIInstrumentation(); const handlers = subscribeSpy.mock.calls[0][0]; const event: any = { @@ -275,7 +236,7 @@ describe("GoogleGenAIPlugin", () => { describe("interactions.create channel subscription", () => { it("subscribes to the interactions.create channel", () => { - plugin.enable(); + registerGoogleGenAIInstrumentation(); expect(mockNewTracingChannel).toHaveBeenCalledWith( "orchestrion:@google/genai:interactions.create", @@ -284,7 +245,7 @@ describe("GoogleGenAIPlugin", () => { }); it("logs non-streaming interaction output and metrics", () => { - plugin.enable(); + registerGoogleGenAIInstrumentation(); const handlers = subscribeSpy.mock.calls[3][0]; const scheduledAt = new Date("2026-01-02T03:04:05.000Z"); @@ -412,7 +373,7 @@ describe("GoogleGenAIPlugin", () => { }); it("preserves zero and missing interaction usage values", () => { - plugin.enable(); + registerGoogleGenAIInstrumentation(); const handlers = subscribeSpy.mock.calls[3][0]; const event: any = { @@ -488,7 +449,7 @@ describe("GoogleGenAIPlugin", () => { }); it("does not trace background interaction tasks", () => { - plugin.enable(); + registerGoogleGenAIInstrumentation(); const handlers = subscribeSpy.mock.calls[3][0]; const event: any = { @@ -512,7 +473,7 @@ describe("GoogleGenAIPlugin", () => { }); it("aggregates streaming interaction events when consumed", async () => { - plugin.enable(); + registerGoogleGenAIInstrumentation(); async function* stream() { yield { @@ -599,7 +560,7 @@ describe("GoogleGenAIPlugin", () => { }); it("ends the interaction span when a stream errors", async () => { - plugin.enable(); + registerGoogleGenAIInstrumentation(); const streamError = new Error("stream failed"); async function* stream() { @@ -651,7 +612,7 @@ describe("Google GenAI serialization functions", () => { contents: "Hello world", }; - // Since these are private functions, we'll test them through the plugin behavior + // Since these are private functions, test them through consumer behavior. // For now, we'll verify the structure by examining what gets logged expect(params.model).toBe("gemini-pro"); expect(params.contents).toBe("Hello world"); diff --git a/js/src/instrumentation/plugins/google-genai-plugin.ts b/js/src/instrumentation/providers/google-genai-instrumentation.ts similarity index 92% rename from js/src/instrumentation/plugins/google-genai-plugin.ts rename to js/src/instrumentation/providers/google-genai-instrumentation.ts index fc70e112c..85e5ae049 100644 --- a/js/src/instrumentation/plugins/google-genai-plugin.ts +++ b/js/src/instrumentation/providers/google-genai-instrumentation.ts @@ -1,5 +1,4 @@ -import { BasePlugin } from "../core"; -import { traceStreamingChannel, unsubscribeAll } from "../core/channel-tracing"; +import { traceStreamingChannel } from "../core/channel-tracing"; import type { ChannelMessage, ErrorOf, @@ -74,29 +73,25 @@ function createWrapperParityEvent(args: { } /** - * Auto-instrumentation plugin for the Google GenAI SDK. + * Internal auto-instrumentation consumer for the Google GenAI SDK. * - * This plugin subscribes to orchestrion channels for Google GenAI SDK methods + * This consumer subscribes to orchestrion channels for Google GenAI SDK methods * and creates Braintrust spans to track: * - models.generateContent (non-streaming) * - models.generateContentStream (streaming) * - models.embedContent (embeddings) * - * The plugin handles: + * The consumer handles: * - Google-specific token metrics (promptTokenCount, candidatesTokenCount, cachedContentTokenCount) * - Processing streaming responses * - Converting inline data (images) to Attachment objects * - Tool calls (functionCall, functionResponse) and executable code results */ -export class GoogleGenAIPlugin extends BasePlugin { - protected onEnable(): void { +class GoogleGenAIInstrumentationConsumer { + public register(): void { this.subscribeToGoogleGenAIChannels(); } - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } - private subscribeToGoogleGenAIChannels(): void { this.subscribeToGenerateContentChannel(); this.subscribeToGenerateContentStreamChannel(); @@ -110,32 +105,28 @@ export class GoogleGenAIPlugin extends BasePlugin { ChannelMessage >; const states = new WeakMap(); - const unbindCurrentSpanStore = bindCurrentSpanStoreToStart( - tracingChannel, - states, - (event) => { - const params = event.arguments[0]; - const input = serializeGenerateContentInput(params); - const metadata = extractGenerateContentMetadata(params); - const span = startBaseSpan( - withSpanInstrumentationName( - { - name: "generate_content", - spanAttributes: { - type: SpanTypeAttribute.LLM, - }, - event: createWrapperParityEvent({ input, metadata }), + bindCurrentSpanStoreToStart(tracingChannel, states, (event) => { + const params = event.arguments[0]; + const input = serializeGenerateContentInput(params); + const metadata = extractGenerateContentMetadata(params); + const span = startBaseSpan( + withSpanInstrumentationName( + { + name: "generate_content", + spanAttributes: { + type: SpanTypeAttribute.LLM, }, - INSTRUMENTATION_NAMES.GOOGLE_GENAI, - ), - ); - - return { - span, - startTime: getCurrentUnixTimestamp(), - }; - }, - ); + event: createWrapperParityEvent({ input, metadata }), + }, + INSTRUMENTATION_NAMES.GOOGLE_GENAI, + ), + ); + + return { + span, + startTime: getCurrentUnixTimestamp(), + }; + }); const handlers: IsoChannelHandlers> = { @@ -192,10 +183,6 @@ export class GoogleGenAIPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - unbindCurrentSpanStore?.(); - tracingChannel.unsubscribe(handlers); - }); } private subscribeToGenerateContentStreamChannel(): void { @@ -228,9 +215,6 @@ export class GoogleGenAIPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - tracingChannel.unsubscribe(handlers); - }); } private subscribeToEmbedContentChannel(): void { @@ -239,32 +223,28 @@ export class GoogleGenAIPlugin extends BasePlugin { ChannelMessage >; const states = new WeakMap(); - const unbindCurrentSpanStore = bindCurrentSpanStoreToStart( - tracingChannel, - states, - (event) => { - const params = event.arguments[0]; - const input = serializeEmbedContentInput(params); - const metadata = extractEmbedContentMetadata(params); - const span = startBaseSpan( - withSpanInstrumentationName( - { - name: "embed_content", - spanAttributes: { - type: SpanTypeAttribute.LLM, - }, - event: createWrapperParityEvent({ input, metadata }), + bindCurrentSpanStoreToStart(tracingChannel, states, (event) => { + const params = event.arguments[0]; + const input = serializeEmbedContentInput(params); + const metadata = extractEmbedContentMetadata(params); + const span = startBaseSpan( + withSpanInstrumentationName( + { + name: "embed_content", + spanAttributes: { + type: SpanTypeAttribute.LLM, }, - INSTRUMENTATION_NAMES.GOOGLE_GENAI, - ), - ); - - return { - span, - startTime: getCurrentUnixTimestamp(), - }; - }, - ); + event: createWrapperParityEvent({ input, metadata }), + }, + INSTRUMENTATION_NAMES.GOOGLE_GENAI, + ), + ); + + return { + span, + startTime: getCurrentUnixTimestamp(), + }; + }); const handlers: IsoChannelHandlers> = { start: (event) => { @@ -316,33 +296,26 @@ export class GoogleGenAIPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - unbindCurrentSpanStore?.(); - tracingChannel.unsubscribe(handlers); - }); } private subscribeToInteractionsCreateChannel(): void { - this.unsubscribers.push( - traceStreamingChannel( - googleGenAIChannels.interactionsCreate as InteractionsCreateChannel, - { - name: "create_interaction", - shouldTrace: ([params]) => !isBackgroundInteractionCreate(params), - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => ({ - input: serializeInteractionInput(params), - metadata: extractInteractionMetadata(params), - }), - extractOutput: (result) => serializeInteractionValue(result), - extractMetadata: (result) => - extractInteractionResponseMetadata(result), - extractMetrics: (result, startTime) => - cleanMetrics(extractInteractionMetrics(result, startTime)), - aggregateChunks: (chunks, _result, _event, startTime) => - aggregateInteractionEvents(chunks, startTime), - }, - ), + traceStreamingChannel( + googleGenAIChannels.interactionsCreate as InteractionsCreateChannel, + { + name: "create_interaction", + shouldTrace: ([params]) => !isBackgroundInteractionCreate(params), + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => ({ + input: serializeInteractionInput(params), + metadata: extractInteractionMetadata(params), + }), + extractOutput: (result) => serializeInteractionValue(result), + extractMetadata: (result) => extractInteractionResponseMetadata(result), + extractMetrics: (result, startTime) => + cleanMetrics(extractInteractionMetrics(result, startTime)), + aggregateChunks: (chunks, _result, _event, startTime) => + aggregateInteractionEvents(chunks, startTime), + }, ); } } @@ -372,7 +345,7 @@ function bindCurrentSpanStoreToStart< tracingChannel: IsoTracingChannel>, states: WeakMap, create: (event: StartOf) => SpanState, -): (() => void) | undefined { +): void { const state = _internalGetGlobalState(); const contextManager = state?.contextManager; const startChannel = tracingChannel.start as @@ -381,7 +354,6 @@ function bindCurrentSpanStoreToStart< store: CurrentSpanStore, callback: (event: ChannelMessage) => unknown, ) => void; - unbindStore?: (store: CurrentSpanStore) => void; } & object) | undefined; const currentSpanStore = contextManager @@ -393,7 +365,7 @@ function bindCurrentSpanStoreToStart< : undefined; if (!startChannel?.bindStore || !currentSpanStore) { - return undefined; + return; } startChannel.bindStore(currentSpanStore, (event) => { @@ -402,10 +374,6 @@ function bindCurrentSpanStoreToStart< ).span; return contextManager!.wrapSpanForStore(span); }); - - return () => { - startChannel.unbindStore?.(currentSpanStore); - }; } function logErrorAndEndSpan( @@ -1626,3 +1594,13 @@ function tryToDict(obj: unknown): Record | null { return null; } + +let googleGenAIInstrumentationConsumer: + | GoogleGenAIInstrumentationConsumer + | undefined; + +export function registerGoogleGenAIInstrumentation(): void { + googleGenAIInstrumentationConsumer ??= + new GoogleGenAIInstrumentationConsumer(); + googleGenAIInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/groq-channels.ts b/js/src/instrumentation/providers/groq-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/groq-channels.ts rename to js/src/instrumentation/providers/groq-channels.ts diff --git a/js/src/instrumentation/plugins/groq-plugin.test.ts b/js/src/instrumentation/providers/groq-instrumentation.test.ts similarity index 98% rename from js/src/instrumentation/plugins/groq-plugin.test.ts rename to js/src/instrumentation/providers/groq-instrumentation.test.ts index a140deb80..3b10736d4 100644 --- a/js/src/instrumentation/plugins/groq-plugin.test.ts +++ b/js/src/instrumentation/providers/groq-instrumentation.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it } from "vitest"; import { aggregateGroqChatCompletionChunks, parseGroqMetrics, -} from "./groq-plugin"; +} from "./groq-instrumentation"; describe("parseGroqMetrics", () => { it("merges OpenAI-compatible usage metrics with Groq cache metrics", () => { diff --git a/js/src/instrumentation/plugins/groq-plugin.ts b/js/src/instrumentation/providers/groq-instrumentation.ts similarity index 58% rename from js/src/instrumentation/plugins/groq-plugin.ts rename to js/src/instrumentation/providers/groq-instrumentation.ts index 10b6d2c0a..6f2b85d5f 100644 --- a/js/src/instrumentation/plugins/groq-plugin.ts +++ b/js/src/instrumentation/providers/groq-instrumentation.ts @@ -1,8 +1,6 @@ -import { BasePlugin } from "../core"; import { traceAsyncChannel, traceStreamingChannel, - unsubscribeAll, } from "../core/channel-tracing"; import { SpanTypeAttribute } from "../../../util/index"; import { processInputAttachments } from "../../wrappers/attachment-utils"; @@ -10,63 +8,53 @@ import { getCurrentUnixTimestamp } from "../../util"; import { aggregateChatCompletionChunks, parseMetricsFromUsage, -} from "./openai-plugin"; +} from "./openai-instrumentation"; import { groqChannels } from "./groq-channels"; import type { GroqChatCompletion, GroqChatCompletionChunk, } from "../../vendor-sdk-types/groq"; -export class GroqPlugin extends BasePlugin { - protected onEnable(): void { - this.unsubscribers.push( - traceStreamingChannel(groqChannels.chatCompletionsCreate, { - name: "groq.chat.completions.create", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const { messages, ...metadata } = params; - return { - input: processInputAttachments(messages), - metadata: { ...metadata, provider: "groq" }, - }; - }, - extractOutput: (result) => result?.choices, - extractMetrics: (result, startTime) => { - const metrics = parseGroqMetrics(result); - if (startTime) { - metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; - } - return metrics; - }, - aggregateChunks: aggregateGroqChatCompletionChunks, - }), - ); +export function registerGroqInstrumentation(): void { + traceStreamingChannel(groqChannels.chatCompletionsCreate, { + name: "groq.chat.completions.create", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { messages, ...metadata } = params; + return { + input: processInputAttachments(messages), + metadata: { ...metadata, provider: "groq" }, + }; + }, + extractOutput: (result) => result?.choices, + extractMetrics: (result, startTime) => { + const metrics = parseGroqMetrics(result); + if (startTime) { + metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; + } + return metrics; + }, + aggregateChunks: aggregateGroqChatCompletionChunks, + }); - this.unsubscribers.push( - traceAsyncChannel(groqChannels.embeddingsCreate, { - name: "groq.embeddings.create", - type: SpanTypeAttribute.LLM, - extractInput: ([params]) => { - const { input, ...metadata } = params; - return { - input, - metadata: { ...metadata, provider: "groq" }, - }; - }, - extractOutput: (result) => { - const embedding = result?.data?.[0]?.embedding; - return Array.isArray(embedding) - ? { embedding_length: embedding.length } - : undefined; - }, - extractMetrics: (result) => parseGroqMetrics(result), - }), - ); - } - - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } + traceAsyncChannel(groqChannels.embeddingsCreate, { + name: "groq.embeddings.create", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { input, ...metadata } = params; + return { + input, + metadata: { ...metadata, provider: "groq" }, + }; + }, + extractOutput: (result) => { + const embedding = result?.data?.[0]?.embedding; + return Array.isArray(embedding) + ? { embedding_length: embedding.length } + : undefined; + }, + extractMetrics: (result) => parseGroqMetrics(result), + }); } export function parseGroqMetrics( diff --git a/js/src/instrumentation/plugins/harness-agent-context.test.ts b/js/src/instrumentation/providers/harness-agent-context.test.ts similarity index 100% rename from js/src/instrumentation/plugins/harness-agent-context.test.ts rename to js/src/instrumentation/providers/harness-agent-context.test.ts diff --git a/js/src/instrumentation/plugins/huggingface-channels.ts b/js/src/instrumentation/providers/huggingface-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/huggingface-channels.ts rename to js/src/instrumentation/providers/huggingface-channels.ts diff --git a/js/src/instrumentation/plugins/huggingface-plugin.test.ts b/js/src/instrumentation/providers/huggingface-instrumentation.test.ts similarity index 98% rename from js/src/instrumentation/plugins/huggingface-plugin.test.ts rename to js/src/instrumentation/providers/huggingface-instrumentation.test.ts index 8b2b80286..6c55e8bbe 100644 --- a/js/src/instrumentation/plugins/huggingface-plugin.test.ts +++ b/js/src/instrumentation/providers/huggingface-instrumentation.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it } from "vitest"; import { aggregateChatCompletionChunks, extractResponseMetadata, -} from "./huggingface-plugin"; +} from "./huggingface-instrumentation"; describe("extractResponseMetadata", () => { it("keeps allowlisted response fields", () => { diff --git a/js/src/instrumentation/plugins/huggingface-plugin.ts b/js/src/instrumentation/providers/huggingface-instrumentation.ts similarity index 87% rename from js/src/instrumentation/plugins/huggingface-plugin.ts rename to js/src/instrumentation/providers/huggingface-instrumentation.ts index 797c3bf82..0407adffa 100644 --- a/js/src/instrumentation/plugins/huggingface-plugin.ts +++ b/js/src/instrumentation/providers/huggingface-instrumentation.ts @@ -1,10 +1,8 @@ import { traceAsyncChannel, traceSyncStreamChannel, - unsubscribeAll, } from "../core/channel-tracing"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; -import { BasePlugin } from "../core"; import { SpanTypeAttribute, isObject } from "../../../util/index"; import { getCurrentUnixTimestamp } from "../../util"; import { parseMetricsFromUsage } from "../../openai-utils"; @@ -41,62 +39,54 @@ const RESPONSE_METADATA_ALLOWLIST = new Set([ "object", ]); -export class HuggingFacePlugin extends BasePlugin { - protected onEnable(): void { - this.unsubscribers.push( - traceAsyncChannel(huggingFaceChannels.chatCompletion, { - name: "huggingface.chat_completion", - type: SpanTypeAttribute.LLM, - extractInput: extractChatInputWithMetadata, - extractOutput: (result) => result?.choices, - extractMetadata: (result) => extractResponseMetadata(result), - extractMetrics: (result) => parseMetricsFromUsage(result?.usage), - }), - traceSyncStreamChannel(huggingFaceChannels.chatCompletionStream, { - name: "huggingface.chat_completion_stream", - type: SpanTypeAttribute.LLM, - extractInput: extractChatInputWithMetadata, - patchResult: ({ result, span, startTime }) => - patchChatCompletionStream({ - result, - span, - startTime, - }), - }), - traceAsyncChannel(huggingFaceChannels.textGeneration, { - name: "huggingface.text_generation", - type: SpanTypeAttribute.LLM, - extractInput: extractTextGenerationInputWithMetadata, - extractOutput: (result) => - isObject(result) ? { generated_text: result.generated_text } : result, - extractMetadata: extractTextGenerationMetadata, - extractMetrics: (result) => - extractTextGenerationMetrics(result?.details ?? null), - }), - traceSyncStreamChannel(huggingFaceChannels.textGenerationStream, { - name: "huggingface.text_generation_stream", - type: SpanTypeAttribute.LLM, - extractInput: extractTextGenerationInputWithMetadata, - patchResult: ({ result, span, startTime }) => - patchTextGenerationStream({ - result, - span, - startTime, - }), +export function registerHuggingFaceInstrumentation(): void { + traceAsyncChannel(huggingFaceChannels.chatCompletion, { + name: "huggingface.chat_completion", + type: SpanTypeAttribute.LLM, + extractInput: extractChatInputWithMetadata, + extractOutput: (result) => result?.choices, + extractMetadata: (result) => extractResponseMetadata(result), + extractMetrics: (result) => parseMetricsFromUsage(result?.usage), + }); + traceSyncStreamChannel(huggingFaceChannels.chatCompletionStream, { + name: "huggingface.chat_completion_stream", + type: SpanTypeAttribute.LLM, + extractInput: extractChatInputWithMetadata, + patchResult: ({ result, span, startTime }) => + patchChatCompletionStream({ + result, + span, + startTime, }), - traceAsyncChannel(huggingFaceChannels.featureExtraction, { - name: "huggingface.feature_extraction", - type: SpanTypeAttribute.LLM, - extractInput: extractFeatureExtractionInputWithMetadata, - extractOutput: summarizeFeatureExtractionOutput, - extractMetrics: () => ({}), + }); + traceAsyncChannel(huggingFaceChannels.textGeneration, { + name: "huggingface.text_generation", + type: SpanTypeAttribute.LLM, + extractInput: extractTextGenerationInputWithMetadata, + extractOutput: (result) => + isObject(result) ? { generated_text: result.generated_text } : result, + extractMetadata: extractTextGenerationMetadata, + extractMetrics: (result) => + extractTextGenerationMetrics(result?.details ?? null), + }); + traceSyncStreamChannel(huggingFaceChannels.textGenerationStream, { + name: "huggingface.text_generation_stream", + type: SpanTypeAttribute.LLM, + extractInput: extractTextGenerationInputWithMetadata, + patchResult: ({ result, span, startTime }) => + patchTextGenerationStream({ + result, + span, + startTime, }), - ); - } - - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } + }); + traceAsyncChannel(huggingFaceChannels.featureExtraction, { + name: "huggingface.feature_extraction", + type: SpanTypeAttribute.LLM, + extractInput: extractFeatureExtractionInputWithMetadata, + extractOutput: summarizeFeatureExtractionOutput, + extractMetrics: () => ({}), + }); } function addProviderMetadata( diff --git a/js/src/instrumentation/plugins/huggingface-transformers-channels.ts b/js/src/instrumentation/providers/huggingface-transformers-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/huggingface-transformers-channels.ts rename to js/src/instrumentation/providers/huggingface-transformers-channels.ts diff --git a/js/src/instrumentation/plugins/huggingface-transformers-plugin.test.ts b/js/src/instrumentation/providers/huggingface-transformers-instrumentation.test.ts similarity index 97% rename from js/src/instrumentation/plugins/huggingface-transformers-plugin.test.ts rename to js/src/instrumentation/providers/huggingface-transformers-instrumentation.test.ts index 667fc19c7..b32e93579 100644 --- a/js/src/instrumentation/plugins/huggingface-transformers-plugin.test.ts +++ b/js/src/instrumentation/providers/huggingface-transformers-instrumentation.test.ts @@ -3,13 +3,13 @@ import { registerHuggingFaceTransformersPipeline, type HuggingFaceTransformersEventContext, } from "./huggingface-transformers-channels"; -import { _exportsForTestingOnly } from "./huggingface-transformers-plugin"; +import { _exportsForTestingOnly } from "./huggingface-transformers-instrumentation"; import type { HuggingFaceTransformersPipeline } from "../../vendor-sdk-types/huggingface-transformers"; const { extractInput, extractMetadata, extractOutput, isSupportedTask } = _exportsForTestingOnly; -describe("HuggingFaceTransformersPlugin extraction", () => { +describe("registerHuggingFaceTransformersInstrumentation extraction", () => { it("normalizes generation and chat payloads", () => { expect(extractInput("text-generation", ["Hello"])).toEqual([ { role: "user", content: "Hello" }, diff --git a/js/src/instrumentation/plugins/huggingface-transformers-plugin.ts b/js/src/instrumentation/providers/huggingface-transformers-instrumentation.ts similarity index 82% rename from js/src/instrumentation/plugins/huggingface-transformers-plugin.ts rename to js/src/instrumentation/providers/huggingface-transformers-instrumentation.ts index 8d41b535a..09e817685 100644 --- a/js/src/instrumentation/plugins/huggingface-transformers-plugin.ts +++ b/js/src/instrumentation/providers/huggingface-transformers-instrumentation.ts @@ -1,5 +1,4 @@ -import { BasePlugin } from "../core"; -import { traceAsyncChannel, unsubscribeAll } from "../core/channel-tracing"; +import { traceAsyncChannel } from "../core/channel-tracing"; import type { ChannelMessage } from "../core/channel-definitions"; import type { IsoChannelHandlers, IsoTracingChannel } from "../../isomorph"; import { SpanTypeAttribute, isObject } from "../../../util"; @@ -19,43 +18,37 @@ const REQUEST_METADATA_KEYS = [ "stop", ] as const; -export class HuggingFaceTransformersPlugin extends BasePlugin { - protected onEnable(): void { +class HuggingFaceTransformersInstrumentationConsumer { + public register(): void { this.subscribeToPipelineFactory(); - this.unsubscribers.push( - traceAsyncChannel(huggingFaceTransformersChannels.pipelineCall, { - name: (_args, event) => { - const task = getTask(event as HuggingFaceTransformersEventContext); - const operation = task?.replaceAll("-", "_") ?? "unknown"; - return `huggingface.transformers.${operation}`; - }, - type: SpanTypeAttribute.LLM, - shouldTrace: (_args, event) => - isSupportedHuggingFaceTransformersTask( - getTask(event as HuggingFaceTransformersEventContext), - ), - extractInput: (args, event) => ({ - input: extractInput( - getTask(event as HuggingFaceTransformersEventContext), - args, - ), - metadata: extractMetadata( - event as HuggingFaceTransformersEventContext, - args, - ), - }), - extractOutput: (result, event) => - extractOutput( - getTask(event as HuggingFaceTransformersEventContext), - result, - ), - extractMetrics: () => ({}), + traceAsyncChannel(huggingFaceTransformersChannels.pipelineCall, { + name: (_args, event) => { + const task = getTask(event as HuggingFaceTransformersEventContext); + const operation = task?.replaceAll("-", "_") ?? "unknown"; + return `huggingface.transformers.${operation}`; + }, + type: SpanTypeAttribute.LLM, + shouldTrace: (_args, event) => + isSupportedHuggingFaceTransformersTask( + getTask(event as HuggingFaceTransformersEventContext), + ), + extractInput: (args, event) => ({ + input: extractInput( + getTask(event as HuggingFaceTransformersEventContext), + args, + ), + metadata: extractMetadata( + event as HuggingFaceTransformersEventContext, + args, + ), }), - ); - } - - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); + extractOutput: (result, event) => + extractOutput( + getTask(event as HuggingFaceTransformersEventContext), + result, + ), + extractMetrics: () => ({}), + }); } private subscribeToPipelineFactory(): void { @@ -79,7 +72,6 @@ export class HuggingFaceTransformersPlugin extends BasePlugin { }; channel.subscribe(handlers); - this.unsubscribers.push(() => channel.unsubscribe(handlers)); } } @@ -321,3 +313,13 @@ export const _exportsForTestingOnly = { extractOutput, isSupportedTask: isSupportedHuggingFaceTransformersTask, }; + +let huggingFaceTransformersInstrumentationConsumer: + | HuggingFaceTransformersInstrumentationConsumer + | undefined; + +export function registerHuggingFaceTransformersInstrumentation(): void { + huggingFaceTransformersInstrumentationConsumer ??= + new HuggingFaceTransformersInstrumentationConsumer(); + huggingFaceTransformersInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/instrumentation-names.test.ts b/js/src/instrumentation/providers/instrumentation-names.test.ts similarity index 100% rename from js/src/instrumentation/plugins/instrumentation-names.test.ts rename to js/src/instrumentation/providers/instrumentation-names.test.ts diff --git a/js/src/instrumentation/plugins/langchain-channels.ts b/js/src/instrumentation/providers/langchain-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/langchain-channels.ts rename to js/src/instrumentation/providers/langchain-channels.ts diff --git a/js/src/instrumentation/plugins/langchain-plugin.test.ts b/js/src/instrumentation/providers/langchain-instrumentation.test.ts similarity index 79% rename from js/src/instrumentation/plugins/langchain-plugin.test.ts rename to js/src/instrumentation/providers/langchain-instrumentation.test.ts index 29ace20ab..570972a9d 100644 --- a/js/src/instrumentation/plugins/langchain-plugin.test.ts +++ b/js/src/instrumentation/providers/langchain-instrumentation.test.ts @@ -1,5 +1,5 @@ -import { describe, expect, it } from "vitest"; -import { LangChainPlugin } from "./langchain-plugin"; +import { beforeAll, describe, expect, it } from "vitest"; +import { registerLangChainInstrumentation } from "./langchain-instrumentation"; import { langChainChannels } from "./langchain-channels"; function createManager(handlers: unknown[] = []) { @@ -35,14 +35,15 @@ function createArgumentsObject(...args: unknown[]): IArguments { })(...args); } -describe("LangChainPlugin", () => { +describe("registerLangChainInstrumentation", () => { + beforeAll(() => { + registerLangChainInstrumentation(); + }); + it("injects a Braintrust callback handler into empty CallbackManager.configure() arguments", () => { - const plugin = new LangChainPlugin(); const args: unknown[] = []; - plugin.enable(); traceConfigureArguments(args); - plugin.disable(); expect(args[0]).toEqual([ expect.objectContaining({ @@ -52,12 +53,9 @@ describe("LangChainPlugin", () => { }); it("injects a Braintrust callback handler into real arguments objects", () => { - const plugin = new LangChainPlugin(); const args = createArgumentsObject(); - plugin.enable(); traceConfigureArgumentsObject(args); - plugin.disable(); expect(args[0]).toEqual([ expect.objectContaining({ @@ -67,12 +65,9 @@ describe("LangChainPlugin", () => { }); it("injects a Braintrust callback handler into CallbackManager.configure() results", () => { - const plugin = new LangChainPlugin(); const manager = createManager(); - plugin.enable(); traceConfigureResult(manager); - plugin.disable(); expect(manager.handlers).toHaveLength(1); expect(manager.handlers[0]).toMatchObject({ @@ -81,37 +76,25 @@ describe("LangChainPlugin", () => { }); it("does not inject duplicate handlers into the same manager", () => { - const plugin = new LangChainPlugin(); const manager = createManager(); - plugin.enable(); traceConfigureResult(manager); traceConfigureResult(manager); - plugin.disable(); expect(manager.handlers).toHaveLength(1); }); it("does not inject when a Braintrust callback handler is already present", () => { - const plugin = new LangChainPlugin(); const existingHandler = { name: "BraintrustCallbackHandler" }; const manager = createManager([existingHandler]); - plugin.enable(); traceConfigureResult(manager); - plugin.disable(); expect(manager.handlers).toEqual([existingHandler]); }); it("gracefully ignores undefined and non-manager results", () => { - const plugin = new LangChainPlugin(); - - plugin.enable(); - expect(() => traceConfigureResult(undefined)).not.toThrow(); expect(() => traceConfigureResult({ handlers: [] })).not.toThrow(); - - plugin.disable(); }); }); diff --git a/js/src/instrumentation/plugins/langchain-plugin.ts b/js/src/instrumentation/providers/langchain-instrumentation.ts similarity index 89% rename from js/src/instrumentation/plugins/langchain-plugin.ts rename to js/src/instrumentation/providers/langchain-instrumentation.ts index 85ef66083..42a5f7fd5 100644 --- a/js/src/instrumentation/plugins/langchain-plugin.ts +++ b/js/src/instrumentation/providers/langchain-instrumentation.ts @@ -1,4 +1,3 @@ -import { BasePlugin } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import type { IsoChannelHandlers, IsoTracingChannel } from "../../isomorph"; import type { LangChainCallbackManager } from "../../vendor-sdk-types/langchain"; @@ -12,22 +11,14 @@ type LangChainConfigureChannel = | typeof langChainChannels.configure | typeof langChainChannels.configureSync; -export class LangChainPlugin extends BasePlugin { +class LangChainInstrumentationConsumer { private injectedManagers = new WeakSet(); - protected onEnable(): void { + public register(): void { this.subscribeToConfigure(langChainChannels.configure); this.subscribeToConfigure(langChainChannels.configureSync); } - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - this.injectedManagers = new WeakSet(); - } - private subscribeToConfigure(channel: LangChainConfigureChannel): void { const tracingChannel: IsoTracingChannel< ChannelMessage @@ -45,9 +36,6 @@ export class LangChainPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - tracingChannel.unsubscribe(handlers); - }); } private injectHandler(result: unknown): void { @@ -126,3 +114,12 @@ function isBraintrustHandler(handler: unknown): boolean { Reflect.get(handler, "name") === BRAINTRUST_LANGCHAIN_CALLBACK_HANDLER_NAME ); } + +let langChainInstrumentationConsumer: + | LangChainInstrumentationConsumer + | undefined; + +export function registerLangChainInstrumentation(): void { + langChainInstrumentationConsumer ??= new LangChainInstrumentationConsumer(); + langChainInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/langsmith-channels.ts b/js/src/instrumentation/providers/langsmith-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/langsmith-channels.ts rename to js/src/instrumentation/providers/langsmith-channels.ts diff --git a/js/src/instrumentation/plugins/langsmith-plugin.test.ts b/js/src/instrumentation/providers/langsmith-instrumentation.test.ts similarity index 88% rename from js/src/instrumentation/plugins/langsmith-plugin.test.ts rename to js/src/instrumentation/providers/langsmith-instrumentation.test.ts index 8e9b3aed4..55a986a74 100644 --- a/js/src/instrumentation/plugins/langsmith-plugin.test.ts +++ b/js/src/instrumentation/providers/langsmith-instrumentation.test.ts @@ -1,37 +1,36 @@ import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; import { configureNode } from "../../node/config"; import { _exportsForTestingOnly, initLogger } from "../../logger"; +import { configureInstrumentation } from "../registry"; import { langSmithChannels } from "./langsmith-channels"; -import { LangSmithPlugin } from "./langsmith-plugin"; +import { registerLangSmithInstrumentation } from "./langsmith-instrumentation"; try { + configureInstrumentation({ integrations: { langsmith: false } }); configureNode(); } catch { // Best-effort initialization for test environments. } -describe("LangSmithPlugin", () => { +describe("registerLangSmithInstrumentation", () => { let backgroundLogger: ReturnType< typeof _exportsForTestingOnly.useTestBackgroundLogger >; - let plugin: LangSmithPlugin; beforeAll(async () => { await _exportsForTestingOnly.simulateLoginForTests(); + registerLangSmithInstrumentation(); }); beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ - projectName: "langsmith-plugin.test.ts", + projectName: "langsmith-instrumentation.test.ts", projectId: "test-project-id", }); - plugin = new LangSmithPlugin(); - plugin.enable(); }); afterEach(() => { - plugin.disable(); _exportsForTestingOnly.clearTestBackgroundLogger(); }); @@ -269,7 +268,7 @@ describe("LangSmithPlugin", () => { expect(await backgroundLogger.drain()).toEqual([]); }); - it("suppresses LangChain-serialized runs only when configured", async () => { + it("suppresses LangChain-serialized runs by default", async () => { const run = { id: "44444444-4444-4444-8444-444444444444", trace_id: "44444444-4444-4444-8444-444444444444", @@ -282,31 +281,5 @@ describe("LangSmithPlugin", () => { arguments: [run.id, run], }); expect(await backgroundLogger.drain()).toEqual([]); - - plugin.disable(); - plugin = new LangSmithPlugin({ skipLangChainRuns: false }); - plugin.enable(); - await langSmithChannels.updateRun.tracePromise(async () => undefined, { - arguments: [run.id, run], - }); - expect(await backgroundLogger.drain()).toHaveLength(1); - }); - - it("bounds the completed-run deduplication cache", () => { - const completedRuns = ( - plugin as unknown as { - completedRuns: { - get(key: string): true | undefined; - set(key: string, value: true): void; - }; - } - ).completedRuns; - - for (let index = 0; index <= 10_000; index++) { - completedRuns.set(`run-${index}`, true); - } - - expect(completedRuns.get("run-0")).toBeUndefined(); - expect(completedRuns.get("run-10000")).toBe(true); }); }); diff --git a/js/src/instrumentation/plugins/langsmith-plugin.ts b/js/src/instrumentation/providers/langsmith-instrumentation.ts similarity index 95% rename from js/src/instrumentation/plugins/langsmith-plugin.ts rename to js/src/instrumentation/providers/langsmith-instrumentation.ts index 5027fa694..1716b88cb 100644 --- a/js/src/instrumentation/plugins/langsmith-plugin.ts +++ b/js/src/instrumentation/providers/langsmith-instrumentation.ts @@ -11,8 +11,6 @@ import type { LangSmithBatchIngestRuns, LangSmithRun, } from "../../vendor-sdk-types/langsmith"; -import { BasePlugin } from "../core"; -import { unsubscribeAll } from "../core/channel-tracing"; import type { ChannelMessage } from "../core/channel-definitions"; import { langSmithChannels } from "./langsmith-channels"; @@ -21,7 +19,7 @@ type ActiveRun = { span: Span; }; -type LangSmithPluginOptions = { +type LangSmithInstrumentationOptions = { skipLangChainRuns?: boolean; }; @@ -43,19 +41,18 @@ const LLM_SETTING_KEYS = [ "response_format", ] as const; -export class LangSmithPlugin extends BasePlugin { +class LangSmithInstrumentationConsumer { private readonly activeRuns = new Map(); private readonly completedRuns = new LRUCache({ max: MAX_COMPLETED_RUNS, }); private readonly skipLangChainRuns: boolean; - constructor(options: LangSmithPluginOptions = {}) { - super(); + constructor(options: LangSmithInstrumentationOptions = {}) { this.skipLangChainRuns = options.skipLangChainRuns ?? true; } - protected onEnable(): void { + public register(): void { const createChannel = langSmithChannels.createRun.tracingChannel(); const createHandlers = { start: ( @@ -67,7 +64,6 @@ export class LangSmithPlugin extends BasePlugin { }, }; createChannel.subscribe(createHandlers); - this.unsubscribers.push(() => createChannel.unsubscribe(createHandlers)); const updateChannel = langSmithChannels.updateRun.tracingChannel(); const updateHandlers = { @@ -80,7 +76,6 @@ export class LangSmithPlugin extends BasePlugin { }, }; updateChannel.subscribe(updateHandlers); - this.unsubscribers.push(() => updateChannel.unsubscribe(updateHandlers)); const batchChannel = langSmithChannels.batchIngestRuns.tracingChannel(); const batchHandlers = { @@ -93,16 +88,6 @@ export class LangSmithPlugin extends BasePlugin { }, }; batchChannel.subscribe(batchHandlers); - this.unsubscribers.push(() => batchChannel.unsubscribe(batchHandlers)); - } - - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - for (const { span } of this.activeRuns.values()) { - span.end(); - } - this.activeRuns.clear(); - this.completedRuns.clear(); } private processBatch(batch: LangSmithBatchIngestRuns): void { @@ -337,6 +322,19 @@ function mergeRuns( return Object.fromEntries(entries); } +let langSmithInstrumentationConsumer: + | LangSmithInstrumentationConsumer + | undefined; + +export function registerLangSmithInstrumentation( + options: LangSmithInstrumentationOptions = {}, +): void { + langSmithInstrumentationConsumer ??= new LangSmithInstrumentationConsumer( + options, + ); + langSmithInstrumentationConsumer.register(); +} + function isRecord(value: unknown): value is Record { return typeof value === "object" && value !== null && !Array.isArray(value); } diff --git a/js/src/instrumentation/plugins/mistral-channels.ts b/js/src/instrumentation/providers/mistral-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/mistral-channels.ts rename to js/src/instrumentation/providers/mistral-channels.ts diff --git a/js/src/instrumentation/plugins/mistral-plugin.test.ts b/js/src/instrumentation/providers/mistral-instrumentation.test.ts similarity index 99% rename from js/src/instrumentation/plugins/mistral-plugin.test.ts rename to js/src/instrumentation/providers/mistral-instrumentation.test.ts index 96d410440..adab7b87b 100644 --- a/js/src/instrumentation/plugins/mistral-plugin.test.ts +++ b/js/src/instrumentation/providers/mistral-instrumentation.test.ts @@ -4,7 +4,7 @@ import { extractMistralRequestMetadata, extractMistralResponseMetadata, parseMistralMetricsFromUsage, -} from "./mistral-plugin"; +} from "./mistral-instrumentation"; describe("extractMistralRequestMetadata", () => { it("keeps only allowlisted request metadata", () => { diff --git a/js/src/instrumentation/plugins/mistral-plugin.ts b/js/src/instrumentation/providers/mistral-instrumentation.ts similarity index 77% rename from js/src/instrumentation/plugins/mistral-plugin.ts rename to js/src/instrumentation/providers/mistral-instrumentation.ts index f33f82554..7179b569d 100644 --- a/js/src/instrumentation/plugins/mistral-plugin.ts +++ b/js/src/instrumentation/providers/mistral-instrumentation.ts @@ -1,8 +1,6 @@ -import { BasePlugin } from "../core"; import { traceAsyncChannel, traceStreamingChannel, - unsubscribeAll, } from "../core/channel-tracing"; import { SpanTypeAttribute, isObject } from "../../../util/index"; import { processInputAttachments } from "../../wrappers/attachment-utils"; @@ -19,156 +17,130 @@ import type { MistralToolCallDelta, } from "../../vendor-sdk-types/mistral"; -export class MistralPlugin extends BasePlugin { - protected onEnable(): void { +class MistralInstrumentationConsumer { + public register(): void { this.subscribeToMistralChannels(); } - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } - private subscribeToMistralChannels(): void { - this.unsubscribers.push( - traceStreamingChannel(mistralChannels.chatComplete, { - name: "mistral.chat.complete", - type: SpanTypeAttribute.LLM, - extractInput: extractMessagesInputWithMetadata, - extractOutput: (result) => { - return result?.choices; - }, - extractMetadata: (result) => extractMistralResponseMetadata(result), - extractMetrics: (result, startTime) => - extractMistralMetrics(result?.usage, startTime), - }), - ); - - this.unsubscribers.push( - traceStreamingChannel(mistralChannels.chatStream, { - name: "mistral.chat.stream", - type: SpanTypeAttribute.LLM, - extractInput: extractMessagesInputWithMetadata, - extractOutput: extractMistralStreamOutput, - extractMetadata: (result) => extractMistralResponseMetadata(result), - extractMetrics: (result, startTime) => - extractMistralStreamingMetrics(result, startTime), - aggregateChunks: aggregateMistralStreamChunks, - }), - ); - - this.unsubscribers.push( - traceAsyncChannel(mistralChannels.embeddingsCreate, { - name: "mistral.embeddings.create", - type: SpanTypeAttribute.LLM, - extractInput: extractEmbeddingInputWithMetadata, - extractOutput: (result) => { - const embedding = result?.data?.[0]?.embedding; - return Array.isArray(embedding) - ? { embedding_length: embedding.length } - : undefined; - }, - extractMetadata: (result) => extractMistralResponseMetadata(result), - extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), - }), - ); - - this.unsubscribers.push( - traceAsyncChannel(mistralChannels.classifiersModerate, { - name: "mistral.classifiers.moderate", - type: SpanTypeAttribute.LLM, - extractInput: extractClassifierInputWithMetadata, - extractOutput: extractClassifierOutput, - extractMetadata: (result) => extractMistralResponseMetadata(result), - extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), - }), - ); - - this.unsubscribers.push( - traceAsyncChannel(mistralChannels.classifiersModerateChat, { - name: "mistral.classifiers.moderateChat", - type: SpanTypeAttribute.LLM, - extractInput: extractClassifierInputWithMetadata, - extractOutput: extractClassifierOutput, - extractMetadata: (result) => extractMistralResponseMetadata(result), - extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), - }), - ); - - this.unsubscribers.push( - traceAsyncChannel(mistralChannels.classifiersClassify, { - name: "mistral.classifiers.classify", - type: SpanTypeAttribute.LLM, - extractInput: extractClassifierInputWithMetadata, - extractOutput: extractClassifierOutput, - extractMetadata: (result) => extractMistralResponseMetadata(result), - extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), - }), - ); - - this.unsubscribers.push( - traceAsyncChannel(mistralChannels.classifiersClassifyChat, { - name: "mistral.classifiers.classifyChat", - type: SpanTypeAttribute.LLM, - extractInput: extractClassifierInputWithMetadata, - extractOutput: extractClassifierOutput, - extractMetadata: (result) => extractMistralResponseMetadata(result), - extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), - }), - ); - - this.unsubscribers.push( - traceStreamingChannel(mistralChannels.fimComplete, { - name: "mistral.fim.complete", - type: SpanTypeAttribute.LLM, - extractInput: extractPromptInputWithMetadata, - extractOutput: (result) => { - return result?.choices; - }, - extractMetadata: (result) => extractMistralResponseMetadata(result), - extractMetrics: (result, startTime) => - extractMistralMetrics(result?.usage, startTime), - }), - ); - - this.unsubscribers.push( - traceStreamingChannel(mistralChannels.fimStream, { - name: "mistral.fim.stream", - type: SpanTypeAttribute.LLM, - extractInput: extractPromptInputWithMetadata, - extractOutput: extractMistralStreamOutput, - extractMetadata: (result) => extractMistralResponseMetadata(result), - extractMetrics: (result, startTime) => - extractMistralStreamingMetrics(result, startTime), - aggregateChunks: aggregateMistralStreamChunks, - }), - ); - - this.unsubscribers.push( - traceStreamingChannel(mistralChannels.agentsComplete, { - name: "mistral.agents.complete", - type: SpanTypeAttribute.LLM, - extractInput: extractMessagesInputWithMetadata, - extractOutput: (result) => { - return result?.choices; - }, - extractMetadata: (result) => extractMistralResponseMetadata(result), - extractMetrics: (result, startTime) => - extractMistralMetrics(result?.usage, startTime), - }), - ); - - this.unsubscribers.push( - traceStreamingChannel(mistralChannels.agentsStream, { - name: "mistral.agents.stream", - type: SpanTypeAttribute.LLM, - extractInput: extractMessagesInputWithMetadata, - extractOutput: extractMistralStreamOutput, - extractMetadata: (result) => extractMistralResponseMetadata(result), - extractMetrics: (result, startTime) => - extractMistralStreamingMetrics(result, startTime), - aggregateChunks: aggregateMistralStreamChunks, - }), - ); + traceStreamingChannel(mistralChannels.chatComplete, { + name: "mistral.chat.complete", + type: SpanTypeAttribute.LLM, + extractInput: extractMessagesInputWithMetadata, + extractOutput: (result) => { + return result?.choices; + }, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result, startTime) => + extractMistralMetrics(result?.usage, startTime), + }); + + traceStreamingChannel(mistralChannels.chatStream, { + name: "mistral.chat.stream", + type: SpanTypeAttribute.LLM, + extractInput: extractMessagesInputWithMetadata, + extractOutput: extractMistralStreamOutput, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result, startTime) => + extractMistralStreamingMetrics(result, startTime), + aggregateChunks: aggregateMistralStreamChunks, + }); + + traceAsyncChannel(mistralChannels.embeddingsCreate, { + name: "mistral.embeddings.create", + type: SpanTypeAttribute.LLM, + extractInput: extractEmbeddingInputWithMetadata, + extractOutput: (result) => { + const embedding = result?.data?.[0]?.embedding; + return Array.isArray(embedding) + ? { embedding_length: embedding.length } + : undefined; + }, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), + }); + + traceAsyncChannel(mistralChannels.classifiersModerate, { + name: "mistral.classifiers.moderate", + type: SpanTypeAttribute.LLM, + extractInput: extractClassifierInputWithMetadata, + extractOutput: extractClassifierOutput, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), + }); + + traceAsyncChannel(mistralChannels.classifiersModerateChat, { + name: "mistral.classifiers.moderateChat", + type: SpanTypeAttribute.LLM, + extractInput: extractClassifierInputWithMetadata, + extractOutput: extractClassifierOutput, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), + }); + + traceAsyncChannel(mistralChannels.classifiersClassify, { + name: "mistral.classifiers.classify", + type: SpanTypeAttribute.LLM, + extractInput: extractClassifierInputWithMetadata, + extractOutput: extractClassifierOutput, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), + }); + + traceAsyncChannel(mistralChannels.classifiersClassifyChat, { + name: "mistral.classifiers.classifyChat", + type: SpanTypeAttribute.LLM, + extractInput: extractClassifierInputWithMetadata, + extractOutput: extractClassifierOutput, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result) => parseMistralMetricsFromUsage(result?.usage), + }); + + traceStreamingChannel(mistralChannels.fimComplete, { + name: "mistral.fim.complete", + type: SpanTypeAttribute.LLM, + extractInput: extractPromptInputWithMetadata, + extractOutput: (result) => { + return result?.choices; + }, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result, startTime) => + extractMistralMetrics(result?.usage, startTime), + }); + + traceStreamingChannel(mistralChannels.fimStream, { + name: "mistral.fim.stream", + type: SpanTypeAttribute.LLM, + extractInput: extractPromptInputWithMetadata, + extractOutput: extractMistralStreamOutput, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result, startTime) => + extractMistralStreamingMetrics(result, startTime), + aggregateChunks: aggregateMistralStreamChunks, + }); + + traceStreamingChannel(mistralChannels.agentsComplete, { + name: "mistral.agents.complete", + type: SpanTypeAttribute.LLM, + extractInput: extractMessagesInputWithMetadata, + extractOutput: (result) => { + return result?.choices; + }, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result, startTime) => + extractMistralMetrics(result?.usage, startTime), + }); + + traceStreamingChannel(mistralChannels.agentsStream, { + name: "mistral.agents.stream", + type: SpanTypeAttribute.LLM, + extractInput: extractMessagesInputWithMetadata, + extractOutput: extractMistralStreamOutput, + extractMetadata: (result) => extractMistralResponseMetadata(result), + extractMetrics: (result, startTime) => + extractMistralStreamingMetrics(result, startTime), + aggregateChunks: aggregateMistralStreamChunks, + }); } } @@ -909,3 +881,10 @@ export function aggregateMistralStreamChunks( ...(metadata ? { metadata } : {}), }; } + +let mistralInstrumentationConsumer: MistralInstrumentationConsumer | undefined; + +export function registerMistralInstrumentation(): void { + mistralInstrumentationConsumer ??= new MistralInstrumentationConsumer(); + mistralInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/ollama-channels.ts b/js/src/instrumentation/providers/ollama-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/ollama-channels.ts rename to js/src/instrumentation/providers/ollama-channels.ts diff --git a/js/src/instrumentation/plugins/ollama-plugin.test.ts b/js/src/instrumentation/providers/ollama-instrumentation.test.ts similarity index 99% rename from js/src/instrumentation/plugins/ollama-plugin.test.ts rename to js/src/instrumentation/providers/ollama-instrumentation.test.ts index 0006efea8..4503d4ddf 100644 --- a/js/src/instrumentation/plugins/ollama-plugin.test.ts +++ b/js/src/instrumentation/providers/ollama-instrumentation.test.ts @@ -9,7 +9,7 @@ import { extractOllamaEmbedOutput, extractOllamaGenerateInput, extractOllamaMetrics, -} from "./ollama-plugin"; +} from "./ollama-instrumentation"; configureNode(); diff --git a/js/src/instrumentation/plugins/ollama-plugin.ts b/js/src/instrumentation/providers/ollama-instrumentation.ts similarity index 91% rename from js/src/instrumentation/plugins/ollama-plugin.ts rename to js/src/instrumentation/providers/ollama-instrumentation.ts index f48040348..405b3dff9 100644 --- a/js/src/instrumentation/plugins/ollama-plugin.ts +++ b/js/src/instrumentation/providers/ollama-instrumentation.ts @@ -14,54 +14,44 @@ import type { OllamaToolCall, OllamaUsageResponse, } from "../../vendor-sdk-types/ollama"; -import { BasePlugin } from "../core"; import type { AsyncEndOf } from "../core/channel-definitions"; import { traceAsyncChannel, traceStreamingChannel, - unsubscribeAll, } from "../core/channel-tracing"; import { ollamaChannels } from "./ollama-channels"; -export class OllamaPlugin extends BasePlugin { - protected onEnable(): void { - this.unsubscribers.push( - traceStreamingChannel(ollamaChannels.chat, { - name: "ollama.chat", - type: SpanTypeAttribute.LLM, - extractInput: extractOllamaChatInput, - extractOutput: (result, event) => - extractOllamaChatOutput( - result, - countOllamaToolCalls(event?.arguments?.[0]?.messages), - ), - extractMetadata: extractOllamaResponseMetadata, - extractMetrics: extractOllamaMetrics, - aggregateChunks: aggregateOllamaChatChunks, - }), - traceStreamingChannel(ollamaChannels.generate, { - name: "ollama.generate", - type: SpanTypeAttribute.LLM, - extractInput: extractOllamaGenerateInput, - extractOutput: extractOllamaGenerateOutput, - extractMetadata: extractOllamaResponseMetadata, - extractMetrics: extractOllamaMetrics, - aggregateChunks: aggregateOllamaGenerateChunks, - }), - traceAsyncChannel(ollamaChannels.embed, { - name: "ollama.embed", - type: SpanTypeAttribute.LLM, - extractInput: extractOllamaEmbedInput, - extractOutput: extractOllamaEmbedOutput, - extractMetadata: extractOllamaResponseMetadata, - extractMetrics: extractOllamaMetrics, - }), - ); - } - - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } +export function registerOllamaInstrumentation(): void { + traceStreamingChannel(ollamaChannels.chat, { + name: "ollama.chat", + type: SpanTypeAttribute.LLM, + extractInput: extractOllamaChatInput, + extractOutput: (result, event) => + extractOllamaChatOutput( + result, + countOllamaToolCalls(event?.arguments?.[0]?.messages), + ), + extractMetadata: extractOllamaResponseMetadata, + extractMetrics: extractOllamaMetrics, + aggregateChunks: aggregateOllamaChatChunks, + }); + traceStreamingChannel(ollamaChannels.generate, { + name: "ollama.generate", + type: SpanTypeAttribute.LLM, + extractInput: extractOllamaGenerateInput, + extractOutput: extractOllamaGenerateOutput, + extractMetadata: extractOllamaResponseMetadata, + extractMetrics: extractOllamaMetrics, + aggregateChunks: aggregateOllamaGenerateChunks, + }); + traceAsyncChannel(ollamaChannels.embed, { + name: "ollama.embed", + type: SpanTypeAttribute.LLM, + extractInput: extractOllamaEmbedInput, + extractOutput: extractOllamaEmbedOutput, + extractMetadata: extractOllamaResponseMetadata, + extractMetrics: extractOllamaMetrics, + }); } function isNonNegativeNumber(value: unknown): value is number { diff --git a/js/src/instrumentation/plugins/openai-agents-channels.ts b/js/src/instrumentation/providers/openai-agents-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/openai-agents-channels.ts rename to js/src/instrumentation/providers/openai-agents-channels.ts diff --git a/js/src/instrumentation/plugins/openai-agents-plugin.test.ts b/js/src/instrumentation/providers/openai-agents-instrumentation.test.ts similarity index 88% rename from js/src/instrumentation/plugins/openai-agents-plugin.test.ts rename to js/src/instrumentation/providers/openai-agents-instrumentation.test.ts index 9f2a604bd..fb6f62fba 100644 --- a/js/src/instrumentation/plugins/openai-agents-plugin.test.ts +++ b/js/src/instrumentation/providers/openai-agents-instrumentation.test.ts @@ -1,37 +1,36 @@ import { afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest"; import { configureNode } from "../../node/config"; import { _exportsForTestingOnly, initLogger } from "../../logger"; +import { configureInstrumentation } from "../registry"; import { openAIAgentsCoreChannels } from "./openai-agents-channels"; -import { OpenAIAgentsPlugin } from "./openai-agents-plugin"; +import { registerOpenAIAgentsInstrumentation } from "./openai-agents-instrumentation"; try { + configureInstrumentation({ integrations: { openAIAgents: false } }); configureNode(); } catch { // Best-effort initialization for test environments. } -describe("OpenAIAgentsPlugin", () => { +describe("registerOpenAIAgentsInstrumentation", () => { let backgroundLogger: ReturnType< typeof _exportsForTestingOnly.useTestBackgroundLogger >; - let plugin: OpenAIAgentsPlugin; beforeAll(async () => { await _exportsForTestingOnly.simulateLoginForTests(); + registerOpenAIAgentsInstrumentation(); }); beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ - projectName: "openai-agents-plugin.test.ts", + projectName: "openai-agents-instrumentation.test.ts", projectId: "test-project-id", }); - plugin = new OpenAIAgentsPlugin(); - plugin.enable(); }); afterEach(() => { - plugin.disable(); _exportsForTestingOnly.clearTestBackgroundLogger(); }); diff --git a/js/src/instrumentation/plugins/openai-agents-plugin.ts b/js/src/instrumentation/providers/openai-agents-instrumentation.ts similarity index 80% rename from js/src/instrumentation/plugins/openai-agents-plugin.ts rename to js/src/instrumentation/providers/openai-agents-instrumentation.ts index 4cf8af1a0..9db90ac4c 100644 --- a/js/src/instrumentation/plugins/openai-agents-plugin.ts +++ b/js/src/instrumentation/providers/openai-agents-instrumentation.ts @@ -1,5 +1,3 @@ -import { BasePlugin } from "../core"; -import { unsubscribeAll } from "../core/channel-tracing"; import { isObject } from "../../../util/index"; import { openAIAgentsCoreChannels } from "./openai-agents-channels"; import { OpenAIAgentsTraceProcessor } from "./openai-agents-trace-processor"; @@ -41,18 +39,13 @@ function isOpenAIAgentsSpan(value: unknown): value is OpenAIAgentsSpan { ); } -export class OpenAIAgentsPlugin extends BasePlugin { +class OpenAIAgentsInstrumentationConsumer { private processor = new OpenAIAgentsTraceProcessor(); - protected onEnable(): void { + public register(): void { this.subscribeToTraceLifecycle(); } - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - void this.processor.shutdown(); - } - private subscribeToTraceLifecycle(): void { const traceStartChannel = openAIAgentsCoreChannels.onTraceStart.tracingChannel(); @@ -65,9 +58,6 @@ export class OpenAIAgentsPlugin extends BasePlugin { }, }; traceStartChannel.subscribe(traceStartHandlers); - this.unsubscribers.push(() => - traceStartChannel.unsubscribe(traceStartHandlers), - ); const traceEndChannel = openAIAgentsCoreChannels.onTraceEnd.tracingChannel(); @@ -80,9 +70,6 @@ export class OpenAIAgentsPlugin extends BasePlugin { }, }; traceEndChannel.subscribe(traceEndHandlers); - this.unsubscribers.push(() => - traceEndChannel.unsubscribe(traceEndHandlers), - ); const spanStartChannel = openAIAgentsCoreChannels.onSpanStart.tracingChannel(); @@ -95,9 +82,6 @@ export class OpenAIAgentsPlugin extends BasePlugin { }, }; spanStartChannel.subscribe(spanStartHandlers); - this.unsubscribers.push(() => - spanStartChannel.unsubscribe(spanStartHandlers), - ); const spanEndChannel = openAIAgentsCoreChannels.onSpanEnd.tracingChannel(); const spanEndHandlers = { @@ -109,6 +93,15 @@ export class OpenAIAgentsPlugin extends BasePlugin { }, }; spanEndChannel.subscribe(spanEndHandlers); - this.unsubscribers.push(() => spanEndChannel.unsubscribe(spanEndHandlers)); } } + +let openAIAgentsInstrumentationConsumer: + | OpenAIAgentsInstrumentationConsumer + | undefined; + +export function registerOpenAIAgentsInstrumentation(): void { + openAIAgentsInstrumentationConsumer ??= + new OpenAIAgentsInstrumentationConsumer(); + openAIAgentsInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/openai-agents-trace-processor.ts b/js/src/instrumentation/providers/openai-agents-trace-processor.ts similarity index 100% rename from js/src/instrumentation/plugins/openai-agents-trace-processor.ts rename to js/src/instrumentation/providers/openai-agents-trace-processor.ts diff --git a/js/src/instrumentation/plugins/openai-channels.ts b/js/src/instrumentation/providers/openai-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/openai-channels.ts rename to js/src/instrumentation/providers/openai-channels.ts diff --git a/js/src/instrumentation/plugins/openai-codex-channels.ts b/js/src/instrumentation/providers/openai-codex-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/openai-codex-channels.ts rename to js/src/instrumentation/providers/openai-codex-channels.ts diff --git a/js/src/instrumentation/plugins/openai-codex-plugin.test.ts b/js/src/instrumentation/providers/openai-codex-instrumentation.test.ts similarity index 92% rename from js/src/instrumentation/plugins/openai-codex-plugin.test.ts rename to js/src/instrumentation/providers/openai-codex-instrumentation.test.ts index 932abe0ba..f7250bb3b 100644 --- a/js/src/instrumentation/plugins/openai-codex-plugin.test.ts +++ b/js/src/instrumentation/providers/openai-codex-instrumentation.test.ts @@ -15,11 +15,11 @@ vi.mock("../../logger", () => ({ })); import iso from "../../isomorph"; -import { OpenAICodexPlugin } from "./openai-codex-plugin"; +import { registerOpenAICodexInstrumentation } from "./openai-codex-instrumentation"; const mockNewTracingChannel = iso.newTracingChannel as ReturnType; -describe("OpenAICodexPlugin", () => { +describe("registerOpenAICodexInstrumentation", () => { let handlersByName: Map; let spans: Array<{ end: ReturnType; @@ -33,7 +33,6 @@ describe("OpenAICodexPlugin", () => { spans = []; mockNewTracingChannel.mockImplementation((name: string) => ({ subscribe: vi.fn((handlers) => handlersByName.set(name, handlers)), - unsubscribe: vi.fn(), })); mockStartSpan.mockImplementation((args: any) => { const span = { @@ -55,8 +54,7 @@ describe("OpenAICodexPlugin", () => { }); it("uses explicit token totals by precedence instead of summing breakdowns", async () => { - const plugin = new OpenAICodexPlugin(); - plugin.enable(); + registerOpenAICodexInstrumentation(); const runHandlers = handlersByName.get( "orchestrion:@openai/codex-sdk:Thread.run", diff --git a/js/src/instrumentation/plugins/openai-codex-plugin.ts b/js/src/instrumentation/providers/openai-codex-instrumentation.ts similarity index 98% rename from js/src/instrumentation/plugins/openai-codex-plugin.ts rename to js/src/instrumentation/providers/openai-codex-instrumentation.ts index 4005ecb86..abf442120 100644 --- a/js/src/instrumentation/plugins/openai-codex-plugin.ts +++ b/js/src/instrumentation/providers/openai-codex-instrumentation.ts @@ -1,4 +1,4 @@ -import { BasePlugin, toLoggedError } from "../core"; +import { toLoggedError } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import type { IsoChannelHandlers } from "../../isomorph"; import { debugLogger } from "../../debug-logger"; @@ -55,19 +55,12 @@ const PATCHED_STREAMED_TURN = Symbol.for( "braintrust.openai-codex.patched-streamed-turn", ); -export class OpenAICodexPlugin extends BasePlugin { - protected onEnable(): void { +class OpenAICodexInstrumentationConsumer { + public register(): void { this.subscribeToRun(); this.subscribeToRunStreamed(); } - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - } - private subscribeToRun(): void { const channel = openAICodexChannels.run.tracingChannel(); const states = new WeakMap(); @@ -97,9 +90,6 @@ export class OpenAICodexPlugin extends BasePlugin { }; channel.subscribe(handlers); - this.unsubscribers.push(() => { - channel.unsubscribe(handlers); - }); } private subscribeToRunStreamed(): void { @@ -131,9 +121,6 @@ export class OpenAICodexPlugin extends BasePlugin { }; channel.subscribe(handlers); - this.unsubscribers.push(() => { - channel.unsubscribe(handlers); - }); } } @@ -941,3 +928,13 @@ function safeLog(span: Span, event: Parameters[0]): void { function logInstrumentationError(context: string, error: unknown): void { debugLogger.error(`Error processing ${context}:`, error); } + +let openAICodexInstrumentationConsumer: + | OpenAICodexInstrumentationConsumer + | undefined; + +export function registerOpenAICodexInstrumentation(): void { + openAICodexInstrumentationConsumer ??= + new OpenAICodexInstrumentationConsumer(); + openAICodexInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/openai-plugin.test.ts b/js/src/instrumentation/providers/openai-instrumentation.test.ts similarity index 99% rename from js/src/instrumentation/plugins/openai-plugin.test.ts rename to js/src/instrumentation/providers/openai-instrumentation.test.ts index 5f2b237fa..d51301e85 100644 --- a/js/src/instrumentation/plugins/openai-plugin.test.ts +++ b/js/src/instrumentation/providers/openai-instrumentation.test.ts @@ -3,7 +3,7 @@ import { parseMetricsFromUsage, processImagesInOutput, aggregateChatCompletionChunks, -} from "./openai-plugin"; +} from "./openai-instrumentation"; import { Attachment } from "../../logger"; describe("parseMetricsFromUsage", () => { diff --git a/js/src/instrumentation/providers/openai-instrumentation.ts b/js/src/instrumentation/providers/openai-instrumentation.ts new file mode 100644 index 000000000..30cddb5c2 --- /dev/null +++ b/js/src/instrumentation/providers/openai-instrumentation.ts @@ -0,0 +1,573 @@ +import { + traceAsyncChannel, + traceStreamingChannel, + traceSyncStreamChannel, +} from "../core/channel-tracing"; +import { Attachment } from "../../logger"; +import { SpanTypeAttribute, isObject } from "../../../util/index"; +import { getCurrentUnixTimestamp } from "../../util"; +import { processInputAttachments } from "../../wrappers/attachment-utils"; +import { openAIChannels } from "./openai-channels"; +import { + BRAINTRUST_CACHED_STREAM_METRIC, + getCachedMetricFromHeaders, + parseMetricsFromUsage, +} from "../../openai-utils"; +import type { + OpenAIChatChoice, + OpenAIChatCompletionChunk, + OpenAIChatLogprobs, + OpenAIResponseStreamEvent, +} from "../../vendor-sdk-types/openai"; + +export function registerOpenAIInstrumentation(): void { + // Chat Completions - supports streaming + traceStreamingChannel(openAIChannels.chatCompletionsCreate, { + name: "Chat Completion", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { messages, ...metadata } = params; + return { + input: processInputAttachments(messages), + metadata: { ...metadata, provider: "openai" }, + }; + }, + extractOutput: (result) => { + return result?.choices; + }, + extractMetrics: (result, startTime, endEvent) => { + const metrics = withCachedMetric( + parseMetricsFromUsage(result?.usage), + result, + endEvent, + ); + if (startTime) { + metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; + } + return metrics; + }, + aggregateChunks: aggregateChatCompletionChunks, + }); + + // Embeddings + traceAsyncChannel(openAIChannels.embeddingsCreate, { + name: "Embedding", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { input, ...metadata } = params; + return { + input, + metadata: { ...metadata, provider: "openai" }, + }; + }, + extractOutput: (result) => { + const embedding = result?.data?.[0]?.embedding; + return Array.isArray(embedding) + ? { embedding_length: embedding.length } + : undefined; + }, + extractMetrics: (result, _startTime, endEvent) => { + return withCachedMetric( + parseMetricsFromUsage(result?.usage), + result, + endEvent, + ); + }, + }); + + // Beta Chat Completions Parse + traceStreamingChannel(openAIChannels.betaChatCompletionsParse, { + name: "Chat Completion", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { messages, ...metadata } = params; + return { + input: processInputAttachments(messages), + metadata: { ...metadata, provider: "openai" }, + }; + }, + extractOutput: (result) => { + return result?.choices; + }, + extractMetrics: (result, startTime, endEvent) => { + const metrics = withCachedMetric( + parseMetricsFromUsage(result?.usage), + result, + endEvent, + ); + if (startTime) { + metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; + } + return metrics; + }, + aggregateChunks: aggregateChatCompletionChunks, + }); + + // Beta Chat Completions Stream (sync method returning event-based stream) + traceSyncStreamChannel(openAIChannels.betaChatCompletionsStream, { + name: "Chat Completion", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { messages, ...metadata } = params; + return { + input: processInputAttachments(messages), + metadata: { ...metadata, provider: "openai" }, + }; + }, + }); + + // Moderations + traceAsyncChannel(openAIChannels.moderationsCreate, { + name: "Moderation", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { input, ...metadata } = params; + return { + input, + metadata: { ...metadata, provider: "openai" }, + }; + }, + extractOutput: (result) => { + return result?.results; + }, + extractMetrics: (result, _startTime, endEvent) => { + return withCachedMetric( + parseMetricsFromUsage(result?.usage), + result, + endEvent, + ); + }, + }); + + // Responses API - create (supports streaming via stream=true param) + traceStreamingChannel(openAIChannels.responsesCreate, { + name: "openai.responses.create", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { input, ...metadata } = params; + return { + input: processInputAttachments(input), + metadata: { ...metadata, provider: "openai" }, + }; + }, + extractOutput: (result) => { + return processImagesInOutput(result?.output); + }, + extractMetadata: (result) => { + if (!result) { + return undefined; + } + const { output: _output, usage: _usage, ...metadata } = result; + return Object.keys(metadata).length > 0 ? metadata : undefined; + }, + extractMetrics: (result, startTime, endEvent) => { + const metrics = withCachedMetric( + parseMetricsFromUsage(result?.usage), + result, + endEvent, + ); + if (startTime) { + metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; + } + return metrics; + }, + aggregateChunks: aggregateResponseStreamEvents, + }); + + // Responses API - stream (sync method returning event-based stream) + traceSyncStreamChannel(openAIChannels.responsesStream, { + name: "openai.responses.create", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { input, ...metadata } = params; + return { + input: processInputAttachments(input), + metadata: { ...metadata, provider: "openai" }, + }; + }, + extractFromEvent: (event) => { + if (event.type !== "response.completed" || !event.response) { + return {}; + } + + const response = event.response; + const data: Record = {}; + + if (response.output !== undefined) { + data.output = processImagesInOutput(response.output); + } + + const { usage: _usage, output: _output, ...metadata } = response; + if (Object.keys(metadata).length > 0) { + data.metadata = metadata; + } + + data.metrics = parseMetricsFromUsage(response.usage); + return data; + }, + }); + + // Responses API - parse + traceStreamingChannel(openAIChannels.responsesParse, { + name: "openai.responses.parse", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { input, ...metadata } = params; + return { + input: processInputAttachments(input), + metadata: { ...metadata, provider: "openai" }, + }; + }, + extractOutput: (result) => { + return processImagesInOutput(result?.output); + }, + extractMetadata: (result) => { + if (!result) { + return undefined; + } + const { output: _output, usage: _usage, ...metadata } = result; + return Object.keys(metadata).length > 0 ? metadata : undefined; + }, + extractMetrics: (result, startTime, endEvent) => { + const metrics = withCachedMetric( + parseMetricsFromUsage(result?.usage), + result, + endEvent, + ); + if (startTime) { + metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; + } + return metrics; + }, + aggregateChunks: aggregateResponseStreamEvents, + }); + + // Responses API - compact + traceAsyncChannel(openAIChannels.responsesCompact, { + name: "openai.responses.compact", + type: SpanTypeAttribute.LLM, + extractInput: ([params]) => { + const { input, ...metadata } = params; + return { + input: processInputAttachments(input), + metadata: { ...metadata, provider: "openai" }, + }; + }, + extractOutput: (result) => { + return processImagesInOutput(result?.output); + }, + extractMetadata: (result) => { + if (!result) { + return undefined; + } + const { output: _output, usage: _usage, ...metadata } = result; + return Object.keys(metadata).length > 0 ? metadata : undefined; + }, + extractMetrics: (result, startTime, endEvent) => { + const metrics = withCachedMetric( + parseMetricsFromUsage(result?.usage), + result, + endEvent, + ); + if (startTime) { + metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; + } + return metrics; + }, + }); +} + +function getCachedMetricFromEndEvent(endEvent: unknown): number | undefined { + if (!isObject(endEvent)) { + return undefined; + } + + const response = (endEvent as Record).response; + if (!isObject(response)) { + return undefined; + } + + const headers = (response as { headers?: unknown }).headers; + if (!headers || typeof (headers as Headers).get !== "function") { + return undefined; + } + + return getCachedMetricFromHeaders(headers as Headers); +} + +function withCachedMetric( + metrics: Record, + result: unknown, + endEvent?: unknown, +): Record { + if (metrics.cached !== undefined) { + return metrics; + } + + const cachedFromEvent = getCachedMetricFromEndEvent(endEvent); + if (cachedFromEvent !== undefined) { + return { + ...metrics, + cached: cachedFromEvent, + }; + } + + if (!isObject(result)) { + return metrics; + } + + const cached = (result as Record)[ + BRAINTRUST_CACHED_STREAM_METRIC + ]; + + if (typeof cached !== "number") { + return metrics; + } + + return { + ...metrics, + cached, + }; +} + +/** + * Process output to convert base64 images to attachments. + * Used for Responses API image generation output. + */ +export function processImagesInOutput(output: any): any { + if (Array.isArray(output)) { + return output.map(processImagesInOutput); + } + + if (isObject(output)) { + if ( + output.type === "image_generation_call" && + output.result && + typeof output.result === "string" + ) { + const fileExtension = output.output_format || "png"; + const contentType = `image/${fileExtension}`; + + const baseFilename = + output.revised_prompt && typeof output.revised_prompt === "string" + ? output.revised_prompt.slice(0, 50).replace(/[^a-zA-Z0-9]/g, "_") + : "generated_image"; + const filename = `${baseFilename}.${fileExtension}`; + + // Convert base64 string to Blob + const binaryString = atob(output.result); + const bytes = new Uint8Array(binaryString.length); + for (let i = 0; i < binaryString.length; i++) { + bytes[i] = binaryString.charCodeAt(i); + } + const blob = new Blob([bytes], { type: contentType }); + + const attachment = new Attachment({ + data: blob, + filename: filename, + contentType: contentType, + }); + + return { + ...output, + result: attachment, + }; + } + } + + return output; +} + +function mergeLogprobTokens( + existing: OpenAIChatLogprobs["content"] | OpenAIChatLogprobs["refusal"], + incoming: OpenAIChatLogprobs["content"] | OpenAIChatLogprobs["refusal"], +): OpenAIChatLogprobs["content"] | OpenAIChatLogprobs["refusal"] { + if (incoming === undefined) { + return existing; + } + + if (incoming === null) { + return existing ?? null; + } + + if (Array.isArray(existing)) { + return [...existing, ...incoming]; + } + + return [...incoming]; +} + +function aggregateChatLogprobs( + existing: OpenAIChatLogprobs | null | undefined, + incoming: OpenAIChatLogprobs | null | undefined, +): OpenAIChatLogprobs | null | undefined { + if (incoming === undefined) { + return existing; + } + + if (incoming === null) { + return existing ?? null; + } + + const aggregated: OpenAIChatLogprobs = + existing && existing !== null + ? { ...existing, ...incoming } + : { ...incoming }; + + const content = mergeLogprobTokens(existing?.content, incoming.content); + if (content !== undefined) { + aggregated.content = content; + } + + const refusal = mergeLogprobTokens(existing?.refusal, incoming.refusal); + if (refusal !== undefined) { + aggregated.refusal = refusal; + } + + return aggregated; +} + +/** + * Aggregate chat completion chunks into a single response. + * Combines role (first), content (concatenated), tool_calls (by id), + * finish_reason (last), and usage (last chunk). + */ +export function aggregateChatCompletionChunks( + chunks: OpenAIChatCompletionChunk[], + streamResult?: unknown, + endEvent?: unknown, +): { + output: OpenAIChatChoice[]; + metrics: Record; +} { + let role = undefined; + let content = undefined; + let refusal = undefined; + let tool_calls = undefined; + let logprobs: OpenAIChatLogprobs | null | undefined = undefined; + let finish_reason = undefined; + let metrics: Record = {}; + + for (const chunk of chunks) { + if (chunk.usage) { + metrics = { + ...metrics, + ...parseMetricsFromUsage(chunk.usage), + }; + } + + const choice = chunk.choices?.[0]; + if (!choice) { + continue; + } + + if (choice.finish_reason) { + finish_reason = choice.finish_reason; + } + + logprobs = aggregateChatLogprobs(logprobs, choice.logprobs); + + const delta = choice.delta; + if (!delta) { + continue; + } + + if (delta.finish_reason) { + finish_reason = delta.finish_reason; + } + + if (!role && delta.role) { + role = delta.role; + } + + if (delta.content) { + content = (content || "") + delta.content; + } + + if (delta.refusal) { + refusal = (refusal || "") + delta.refusal; + } + + if (delta.tool_calls) { + const toolDelta = delta.tool_calls[0]; + if ( + !tool_calls || + (toolDelta.id && tool_calls[tool_calls.length - 1].id !== toolDelta.id) + ) { + tool_calls = [ + ...(tool_calls || []), + { + id: toolDelta.id, + type: toolDelta.type, + function: toolDelta.function, + }, + ]; + } else { + tool_calls[tool_calls.length - 1].function.arguments += + toolDelta.function.arguments; + } + } + } + + metrics = withCachedMetric(metrics, streamResult, endEvent); + + return { + metrics, + output: [ + { + index: 0, + message: { + role, + content, + ...(refusal !== undefined ? { refusal } : {}), + tool_calls, + }, + logprobs: logprobs ?? null, + finish_reason, + }, + ], + }; +} + +function aggregateResponseStreamEvents( + chunks: OpenAIResponseStreamEvent[], + _streamResult?: unknown, + endEvent?: unknown, +): { + output: any; + metrics: Record; + metadata?: Record; +} { + let output: any = undefined; + let metrics: Record = {}; + let metadata: Record | undefined = undefined; + + for (const chunk of chunks) { + if (!chunk || !chunk.type || !chunk.response) { + continue; + } + if (chunk.type !== "response.completed") { + continue; + } + + const response = chunk.response; + if (response?.output !== undefined) { + output = processImagesInOutput(response.output); + } + + const { usage: _usage, output: _output, ...rest } = response || {}; + if (Object.keys(rest).length > 0) { + metadata = rest; + } + + metrics = parseMetricsFromUsage(response?.usage); + } + + return { + output, + metrics: withCachedMetric(metrics, undefined, endEvent), + ...(metadata !== undefined ? { metadata } : {}), + }; +} + +export { parseMetricsFromUsage }; diff --git a/js/src/instrumentation/plugins/openrouter-agent-channels.ts b/js/src/instrumentation/providers/openrouter-agent-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/openrouter-agent-channels.ts rename to js/src/instrumentation/providers/openrouter-agent-channels.ts diff --git a/js/src/instrumentation/plugins/openrouter-agent-plugin.test.ts b/js/src/instrumentation/providers/openrouter-agent-instrumentation.test.ts similarity index 99% rename from js/src/instrumentation/plugins/openrouter-agent-plugin.test.ts rename to js/src/instrumentation/providers/openrouter-agent-instrumentation.test.ts index 85d66d51c..27fa78ca2 100644 --- a/js/src/instrumentation/plugins/openrouter-agent-plugin.test.ts +++ b/js/src/instrumentation/providers/openrouter-agent-instrumentation.test.ts @@ -14,7 +14,7 @@ import { aggregateOpenRouterChatChunks, aggregateOpenRouterResponseStreamEvents, parseOpenRouterMetricsFromUsage, -} from "./openrouter-agent-plugin"; +} from "./openrouter-agent-instrumentation"; const TEST_PROVIDER = "openai"; const TEST_MODEL = "gpt-4.1-mini"; @@ -25,7 +25,7 @@ try { // Best-effort initialization for test environments. } -describe("OpenRouter Agent Plugin", () => { +describe("OpenRouter Agent instrumentation", () => { let backgroundLogger: ReturnType< typeof _exportsForTestingOnly.useTestBackgroundLogger >; @@ -37,7 +37,7 @@ describe("OpenRouter Agent Plugin", () => { beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ - projectName: "openrouter-agent-plugin.test.ts", + projectName: "openrouter-agent-instrumentation.test.ts", projectId: "test-project-id", }); }); diff --git a/js/src/instrumentation/plugins/openrouter-agent-plugin.ts b/js/src/instrumentation/providers/openrouter-agent-instrumentation.ts similarity index 91% rename from js/src/instrumentation/plugins/openrouter-agent-plugin.ts rename to js/src/instrumentation/providers/openrouter-agent-instrumentation.ts index facda63c0..7f77cccfc 100644 --- a/js/src/instrumentation/plugins/openrouter-agent-plugin.ts +++ b/js/src/instrumentation/providers/openrouter-agent-instrumentation.ts @@ -1,9 +1,8 @@ -import { BasePlugin, toLoggedError } from "../core"; +import { toLoggedError } from "../core"; import { traceAsyncChannel, traceStreamingChannel, traceSyncStreamChannel, - unsubscribeAll, } from "../core/channel-tracing"; import type { ChannelMessage } from "../core/channel-definitions"; import { @@ -26,107 +25,93 @@ import type { OpenRouterAgentToolTurnContext, } from "../../vendor-sdk-types/openrouter-agent"; -export class OpenRouterAgentPlugin extends BasePlugin { - protected onEnable(): void { +class OpenRouterAgentInstrumentationConsumer { + public register(): void { this.subscribeToOpenRouterAgentChannels(); } - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } - private subscribeToOpenRouterAgentChannels(): void { - this.unsubscribers.push( - traceSyncStreamChannel(openRouterAgentChannels.callModel, { - name: "openrouter.callModel", - type: SpanTypeAttribute.TASK, - extractInput: (args) => { - const request = getOpenRouterCallModelRequestArg(args); - return { - input: request - ? extractOpenRouterCallModelInput(request) - : undefined, - metadata: request - ? extractOpenRouterCallModelMetadata(request) - : { provider: "openrouter" }, - }; - }, - patchResult: ({ endEvent, result, span }) => { - return patchOpenRouterCallModelResult({ - request: getOpenRouterCallModelRequestArg(endEvent.arguments), - result, - span, - }); - }, - }), - ); - - this.unsubscribers.push( - traceAsyncChannel(openRouterAgentChannels.callModelTurn, { - name: "openrouter.beta.responses.send", - type: SpanTypeAttribute.LLM, - extractInput: (args, event) => { - const request = getOpenRouterCallModelRequestArg(args); - const metadata = request + traceSyncStreamChannel(openRouterAgentChannels.callModel, { + name: "openrouter.callModel", + type: SpanTypeAttribute.TASK, + extractInput: (args) => { + const request = getOpenRouterCallModelRequestArg(args); + return { + input: request ? extractOpenRouterCallModelInput(request) : undefined, + metadata: request ? extractOpenRouterCallModelMetadata(request) - : { provider: "openrouter" }; + : { provider: "openrouter" }, + }; + }, + patchResult: ({ endEvent, result, span }) => { + return patchOpenRouterCallModelResult({ + request: getOpenRouterCallModelRequestArg(endEvent.arguments), + result, + span, + }); + }, + }); - if (isObject(metadata) && "tools" in metadata) { - delete (metadata as Record).tools; - } + traceAsyncChannel(openRouterAgentChannels.callModelTurn, { + name: "openrouter.beta.responses.send", + type: SpanTypeAttribute.LLM, + extractInput: (args, event) => { + const request = getOpenRouterCallModelRequestArg(args); + const metadata = request + ? extractOpenRouterCallModelMetadata(request) + : { provider: "openrouter" }; + + if (isObject(metadata) && "tools" in metadata) { + delete (metadata as Record).tools; + } + return { + input: request ? extractOpenRouterCallModelInput(request) : undefined, + metadata: { + ...metadata, + step: event.step, + step_type: event.stepType, + }, + }; + }, + extractOutput: (result) => + extractOpenRouterResponseOutput(result as Record), + extractMetadata: (result, event) => { + if (!isObject(result)) { return { - input: request - ? extractOpenRouterCallModelInput(request) - : undefined, - metadata: { - ...metadata, - step: event.step, - step_type: event.stepType, - }, + step: event?.step, + step_type: event?.stepType, }; - }, - extractOutput: (result) => - extractOpenRouterResponseOutput(result as Record), - extractMetadata: (result, event) => { - if (!isObject(result)) { - return { - step: event?.step, - step_type: event?.stepType, - }; - } + } - return { - ...(extractOpenRouterResponseMetadata(result) || {}), - ...(event?.step !== undefined ? { step: event.step } : {}), - ...(event?.stepType ? { step_type: event.stepType } : {}), - }; + return { + ...(extractOpenRouterResponseMetadata(result) || {}), + ...(event?.step !== undefined ? { step: event.step } : {}), + ...(event?.stepType ? { step_type: event.stepType } : {}), + }; + }, + extractMetrics: (result) => + isObject(result) ? parseOpenRouterMetricsFromUsage(result.usage) : {}, + }); + + traceStreamingChannel(openRouterAgentChannels.toolExecute, { + name: "openrouter.tool", + type: SpanTypeAttribute.TOOL, + extractInput: (args, event) => ({ + input: args[0], + metadata: { + provider: "openrouter", + tool_name: event.toolName, + ...(event.toolCallId ? { tool_call_id: event.toolCallId } : {}), }, - extractMetrics: (result) => - isObject(result) ? parseOpenRouterMetricsFromUsage(result.usage) : {}, }), - ); - - this.unsubscribers.push( - traceStreamingChannel(openRouterAgentChannels.toolExecute, { - name: "openrouter.tool", - type: SpanTypeAttribute.TOOL, - extractInput: (args, event) => ({ - input: args[0], - metadata: { - provider: "openrouter", - tool_name: event.toolName, - ...(event.toolCallId ? { tool_call_id: event.toolCallId } : {}), - }, - }), - extractOutput: (result) => result, - extractMetrics: () => ({}), - aggregateChunks: (chunks) => ({ - output: chunks.length > 0 ? chunks[chunks.length - 1] : undefined, - metrics: {}, - }), + extractOutput: (result) => result, + extractMetrics: () => ({}), + aggregateChunks: (chunks) => ({ + output: chunks.length > 0 ? chunks[chunks.length - 1] : undefined, + metrics: {}, }), - ); + }); const callModelChannel = openRouterAgentChannels.callModel.tracingChannel(); const callModelHandlers = { @@ -141,9 +126,6 @@ export class OpenRouterAgentPlugin extends BasePlugin { }; callModelChannel.subscribe(callModelHandlers); - this.unsubscribers.push(() => { - callModelChannel.unsubscribe(callModelHandlers); - }); } } @@ -1326,3 +1308,13 @@ function normalizeError(error: unknown): Error { } export { parseOpenRouterMetricsFromUsage }; + +let openRouterAgentInstrumentationConsumer: + | OpenRouterAgentInstrumentationConsumer + | undefined; + +export function registerOpenRouterAgentInstrumentation(): void { + openRouterAgentInstrumentationConsumer ??= + new OpenRouterAgentInstrumentationConsumer(); + openRouterAgentInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/openrouter-channels.ts b/js/src/instrumentation/providers/openrouter-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/openrouter-channels.ts rename to js/src/instrumentation/providers/openrouter-channels.ts diff --git a/js/src/instrumentation/plugins/openrouter-plugin.test.ts b/js/src/instrumentation/providers/openrouter-instrumentation.test.ts similarity index 99% rename from js/src/instrumentation/plugins/openrouter-plugin.test.ts rename to js/src/instrumentation/providers/openrouter-instrumentation.test.ts index 93627512e..c656d5e6d 100644 --- a/js/src/instrumentation/plugins/openrouter-plugin.test.ts +++ b/js/src/instrumentation/providers/openrouter-instrumentation.test.ts @@ -14,7 +14,7 @@ import { aggregateOpenRouterChatChunks, aggregateOpenRouterResponseStreamEvents, parseOpenRouterMetricsFromUsage, -} from "./openrouter-plugin"; +} from "./openrouter-instrumentation"; const TEST_PROVIDER = "openai"; const TEST_MODEL = "gpt-4.1-mini"; @@ -27,7 +27,7 @@ try { // Best-effort initialization for test environments. } -describe("OpenRouter Plugin", () => { +describe("OpenRouter instrumentation", () => { let backgroundLogger: ReturnType< typeof _exportsForTestingOnly.useTestBackgroundLogger >; @@ -39,7 +39,7 @@ describe("OpenRouter Plugin", () => { beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ - projectName: "openrouter-plugin.test.ts", + projectName: "openrouter-instrumentation.test.ts", projectId: "test-project-id", }); }); diff --git a/js/src/instrumentation/plugins/openrouter-plugin.ts b/js/src/instrumentation/providers/openrouter-instrumentation.ts similarity index 80% rename from js/src/instrumentation/plugins/openrouter-plugin.ts rename to js/src/instrumentation/providers/openrouter-instrumentation.ts index c89b8f9a7..fa8e54aaa 100644 --- a/js/src/instrumentation/plugins/openrouter-plugin.ts +++ b/js/src/instrumentation/providers/openrouter-instrumentation.ts @@ -1,9 +1,8 @@ -import { BasePlugin, toLoggedError } from "../core"; +import { toLoggedError } from "../core"; import { traceAsyncChannel, traceStreamingChannel, traceSyncStreamChannel, - unsubscribeAll, } from "../core/channel-tracing"; import type { ChannelMessage } from "../core/channel-definitions"; import { @@ -28,247 +27,225 @@ import type { OpenRouterToolTurnContext, } from "../../vendor-sdk-types/openrouter"; -export class OpenRouterPlugin extends BasePlugin { - protected onEnable(): void { +class OpenRouterInstrumentationConsumer { + public register(): void { this.subscribeToOpenRouterChannels(); } - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } - private subscribeToOpenRouterChannels(): void { - this.unsubscribers.push( - traceStreamingChannel(openRouterChannels.chatSend, { - name: "openrouter.chat.send", - type: SpanTypeAttribute.LLM, - extractInput: (args) => { - const request = getOpenRouterRequestArg(args); - const chatGenerationParams = isObject(request?.chatGenerationParams) - ? request.chatGenerationParams - : {}; - const httpReferer = request?.httpReferer; - const xTitle = request?.xTitle; - const { messages, ...metadata } = chatGenerationParams; - return { - input: messages, - metadata: buildOpenRouterMetadata(metadata, httpReferer, xTitle), - }; - }, - extractOutput: (result) => { - return isObject(result) ? result.choices : undefined; - }, - extractMetrics: (result, startTime) => { - const metrics = parseOpenRouterMetricsFromUsage(result?.usage); - if (startTime) { - metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; - } - return metrics; - }, - aggregateChunks: aggregateOpenRouterChatChunks, - }), - ); - - this.unsubscribers.push( - traceAsyncChannel(openRouterChannels.embeddingsGenerate, { - name: "openrouter.embeddings.generate", - type: SpanTypeAttribute.LLM, - extractInput: (args) => { - const request = getOpenRouterRequestArg(args); - const requestBody = isObject(request?.requestBody) - ? request.requestBody - : {}; - const httpReferer = request?.httpReferer; - const xTitle = request?.xTitle; - const { input, ...metadata } = requestBody; - return { - input, - metadata: buildOpenRouterEmbeddingMetadata( - metadata, - httpReferer, - xTitle, - ), - }; - }, - extractOutput: (result) => { - if (!isObject(result)) { - return undefined; - } + traceStreamingChannel(openRouterChannels.chatSend, { + name: "openrouter.chat.send", + type: SpanTypeAttribute.LLM, + extractInput: (args) => { + const request = getOpenRouterRequestArg(args); + const chatGenerationParams = isObject(request?.chatGenerationParams) + ? request.chatGenerationParams + : {}; + const httpReferer = request?.httpReferer; + const xTitle = request?.xTitle; + const { messages, ...metadata } = chatGenerationParams; + return { + input: messages, + metadata: buildOpenRouterMetadata(metadata, httpReferer, xTitle), + }; + }, + extractOutput: (result) => { + return isObject(result) ? result.choices : undefined; + }, + extractMetrics: (result, startTime) => { + const metrics = parseOpenRouterMetricsFromUsage(result?.usage); + if (startTime) { + metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; + } + return metrics; + }, + aggregateChunks: aggregateOpenRouterChatChunks, + }); - const embedding = result.data?.[0]?.embedding; - return Array.isArray(embedding) - ? { embedding_length: embedding.length } - : undefined; - }, - extractMetadata: (result) => { - if (!isObject(result)) { - return undefined; - } + traceAsyncChannel(openRouterChannels.embeddingsGenerate, { + name: "openrouter.embeddings.generate", + type: SpanTypeAttribute.LLM, + extractInput: (args) => { + const request = getOpenRouterRequestArg(args); + const requestBody = isObject(request?.requestBody) + ? request.requestBody + : {}; + const httpReferer = request?.httpReferer; + const xTitle = request?.xTitle; + const { input, ...metadata } = requestBody; + return { + input, + metadata: buildOpenRouterEmbeddingMetadata( + metadata, + httpReferer, + xTitle, + ), + }; + }, + extractOutput: (result) => { + if (!isObject(result)) { + return undefined; + } - return extractOpenRouterResponseMetadata(result); - }, - extractMetrics: (result) => { - return isObject(result) - ? parseOpenRouterMetricsFromUsage(result.usage) - : {}; - }, - }), - ); + const embedding = result.data?.[0]?.embedding; + return Array.isArray(embedding) + ? { embedding_length: embedding.length } + : undefined; + }, + extractMetadata: (result) => { + if (!isObject(result)) { + return undefined; + } - this.unsubscribers.push( - traceAsyncChannel(openRouterChannels.rerankRerank, { - name: "openrouter.rerank.rerank", - type: SpanTypeAttribute.LLM, - extractInput: (args) => { - const request = getOpenRouterRequestArg(args); - const requestBody = isObject(request?.requestBody) - ? request.requestBody - : {}; - const httpReferer = request?.httpReferer; - const xTitle = request?.xTitle ?? request?.appTitle; - const { documents, query, ...metadata } = requestBody; - return { - input: { - documents, - query, - }, - metadata: buildOpenRouterRerankMetadata( - metadata, - documents, - httpReferer, - xTitle, - ), - }; - }, - extractOutput: (result) => extractOpenRouterRerankOutput(result), - extractMetadata: (result) => extractOpenRouterResponseMetadata(result), - extractMetrics: (result) => - isObject(result) ? parseOpenRouterMetricsFromUsage(result.usage) : {}, - }), - ); + return extractOpenRouterResponseMetadata(result); + }, + extractMetrics: (result) => { + return isObject(result) + ? parseOpenRouterMetricsFromUsage(result.usage) + : {}; + }, + }); - this.unsubscribers.push( - traceStreamingChannel(openRouterChannels.betaResponsesSend, { - name: "openrouter.beta.responses.send", - type: SpanTypeAttribute.LLM, - extractInput: (args) => { - const request = getOpenRouterRequestArg(args); - const openResponsesRequest = isObject(request?.openResponsesRequest) - ? request.openResponsesRequest - : {}; - const httpReferer = request?.httpReferer; - const xTitle = request?.xTitle; - const { input, ...metadata } = openResponsesRequest; - return { - input, - metadata: buildOpenRouterMetadata(metadata, httpReferer, xTitle), - }; - }, - extractOutput: (result) => - extractOpenRouterResponseOutput(result as Record), - extractMetadata: (result) => extractOpenRouterResponseMetadata(result), - extractMetrics: (result, startTime) => { - const metrics = parseOpenRouterMetricsFromUsage(result?.usage); - if (startTime) { - metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; - } - return metrics; - }, - aggregateChunks: aggregateOpenRouterResponseStreamEvents, - }), - ); + traceAsyncChannel(openRouterChannels.rerankRerank, { + name: "openrouter.rerank.rerank", + type: SpanTypeAttribute.LLM, + extractInput: (args) => { + const request = getOpenRouterRequestArg(args); + const requestBody = isObject(request?.requestBody) + ? request.requestBody + : {}; + const httpReferer = request?.httpReferer; + const xTitle = request?.xTitle ?? request?.appTitle; + const { documents, query, ...metadata } = requestBody; + return { + input: { + documents, + query, + }, + metadata: buildOpenRouterRerankMetadata( + metadata, + documents, + httpReferer, + xTitle, + ), + }; + }, + extractOutput: (result) => extractOpenRouterRerankOutput(result), + extractMetadata: (result) => extractOpenRouterResponseMetadata(result), + extractMetrics: (result) => + isObject(result) ? parseOpenRouterMetricsFromUsage(result.usage) : {}, + }); - this.unsubscribers.push( - traceSyncStreamChannel(openRouterChannels.callModel, { - name: "openrouter.callModel", - type: SpanTypeAttribute.TASK, - extractInput: (args) => { - const request = getOpenRouterCallModelRequestArg(args); - return { - input: request - ? extractOpenRouterCallModelInput(request) - : undefined, - metadata: request - ? extractOpenRouterCallModelMetadata(request) - : { provider: "openrouter" }, - }; - }, - patchResult: ({ endEvent, result, span }) => { - return patchOpenRouterCallModelResult({ - request: getOpenRouterCallModelRequestArg(endEvent.arguments), - result, - span, - }); - }, - }), - ); + traceStreamingChannel(openRouterChannels.betaResponsesSend, { + name: "openrouter.beta.responses.send", + type: SpanTypeAttribute.LLM, + extractInput: (args) => { + const request = getOpenRouterRequestArg(args); + const openResponsesRequest = isObject(request?.openResponsesRequest) + ? request.openResponsesRequest + : {}; + const httpReferer = request?.httpReferer; + const xTitle = request?.xTitle; + const { input, ...metadata } = openResponsesRequest; + return { + input, + metadata: buildOpenRouterMetadata(metadata, httpReferer, xTitle), + }; + }, + extractOutput: (result) => + extractOpenRouterResponseOutput(result as Record), + extractMetadata: (result) => extractOpenRouterResponseMetadata(result), + extractMetrics: (result, startTime) => { + const metrics = parseOpenRouterMetricsFromUsage(result?.usage); + if (startTime) { + metrics.time_to_first_token = getCurrentUnixTimestamp() - startTime; + } + return metrics; + }, + aggregateChunks: aggregateOpenRouterResponseStreamEvents, + }); - this.unsubscribers.push( - traceAsyncChannel(openRouterChannels.callModelTurn, { - name: "openrouter.beta.responses.send", - type: SpanTypeAttribute.LLM, - extractInput: (args, event) => { - const request = getOpenRouterCallModelRequestArg(args); - const metadata = request + traceSyncStreamChannel(openRouterChannels.callModel, { + name: "openrouter.callModel", + type: SpanTypeAttribute.TASK, + extractInput: (args) => { + const request = getOpenRouterCallModelRequestArg(args); + return { + input: request ? extractOpenRouterCallModelInput(request) : undefined, + metadata: request ? extractOpenRouterCallModelMetadata(request) - : { provider: "openrouter" }; + : { provider: "openrouter" }, + }; + }, + patchResult: ({ endEvent, result, span }) => { + return patchOpenRouterCallModelResult({ + request: getOpenRouterCallModelRequestArg(endEvent.arguments), + result, + span, + }); + }, + }); - if (isObject(metadata) && "tools" in metadata) { - delete (metadata as Record).tools; - } + traceAsyncChannel(openRouterChannels.callModelTurn, { + name: "openrouter.beta.responses.send", + type: SpanTypeAttribute.LLM, + extractInput: (args, event) => { + const request = getOpenRouterCallModelRequestArg(args); + const metadata = request + ? extractOpenRouterCallModelMetadata(request) + : { provider: "openrouter" }; + + if (isObject(metadata) && "tools" in metadata) { + delete (metadata as Record).tools; + } + return { + input: request ? extractOpenRouterCallModelInput(request) : undefined, + metadata: { + ...metadata, + step: event.step, + step_type: event.stepType, + }, + }; + }, + extractOutput: (result) => + extractOpenRouterResponseOutput(result as Record), + extractMetadata: (result, event) => { + if (!isObject(result)) { return { - input: request - ? extractOpenRouterCallModelInput(request) - : undefined, - metadata: { - ...metadata, - step: event.step, - step_type: event.stepType, - }, + step: event?.step, + step_type: event?.stepType, }; - }, - extractOutput: (result) => - extractOpenRouterResponseOutput(result as Record), - extractMetadata: (result, event) => { - if (!isObject(result)) { - return { - step: event?.step, - step_type: event?.stepType, - }; - } + } - return { - ...(extractOpenRouterResponseMetadata(result) || {}), - ...(event?.step !== undefined ? { step: event.step } : {}), - ...(event?.stepType ? { step_type: event.stepType } : {}), - }; + return { + ...(extractOpenRouterResponseMetadata(result) || {}), + ...(event?.step !== undefined ? { step: event.step } : {}), + ...(event?.stepType ? { step_type: event.stepType } : {}), + }; + }, + extractMetrics: (result) => + isObject(result) ? parseOpenRouterMetricsFromUsage(result.usage) : {}, + }); + + traceStreamingChannel(openRouterChannels.toolExecute, { + name: "openrouter.tool", + type: SpanTypeAttribute.TOOL, + extractInput: (args, event) => ({ + input: args[0], + metadata: { + provider: "openrouter", + tool_name: event.toolName, + ...(event.toolCallId ? { tool_call_id: event.toolCallId } : {}), }, - extractMetrics: (result) => - isObject(result) ? parseOpenRouterMetricsFromUsage(result.usage) : {}, }), - ); - - this.unsubscribers.push( - traceStreamingChannel(openRouterChannels.toolExecute, { - name: "openrouter.tool", - type: SpanTypeAttribute.TOOL, - extractInput: (args, event) => ({ - input: args[0], - metadata: { - provider: "openrouter", - tool_name: event.toolName, - ...(event.toolCallId ? { tool_call_id: event.toolCallId } : {}), - }, - }), - extractOutput: (result) => result, - extractMetrics: () => ({}), - aggregateChunks: (chunks) => ({ - output: chunks.length > 0 ? chunks[chunks.length - 1] : undefined, - metrics: {}, - }), + extractOutput: (result) => result, + extractMetrics: () => ({}), + aggregateChunks: (chunks) => ({ + output: chunks.length > 0 ? chunks[chunks.length - 1] : undefined, + metrics: {}, }), - ); + }); const callModelChannel = openRouterChannels.callModel.tracingChannel(); const callModelHandlers = { @@ -283,9 +260,6 @@ export class OpenRouterPlugin extends BasePlugin { }; callModelChannel.subscribe(callModelHandlers); - this.unsubscribers.push(() => { - callModelChannel.unsubscribe(callModelHandlers); - }); } } @@ -1508,3 +1482,12 @@ function normalizeError(error: unknown): Error { } export { parseOpenRouterMetricsFromUsage }; + +let openRouterInstrumentationConsumer: + | OpenRouterInstrumentationConsumer + | undefined; + +export function registerOpenRouterInstrumentation(): void { + openRouterInstrumentationConsumer ??= new OpenRouterInstrumentationConsumer(); + openRouterInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/pi-coding-agent-channels.ts b/js/src/instrumentation/providers/pi-coding-agent-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/pi-coding-agent-channels.ts rename to js/src/instrumentation/providers/pi-coding-agent-channels.ts diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts b/js/src/instrumentation/providers/pi-coding-agent-instrumentation.test.ts similarity index 91% rename from js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts rename to js/src/instrumentation/providers/pi-coding-agent-instrumentation.test.ts index 433361974..7e7d406ed 100644 --- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.test.ts +++ b/js/src/instrumentation/providers/pi-coding-agent-instrumentation.test.ts @@ -30,7 +30,7 @@ vi.mock("../../logger", () => ({ })); import { isAutoInstrumentationSuppressed } from "../auto-instrumentation-suppression"; -import { PiCodingAgentPlugin } from "./pi-coding-agent-plugin"; +import { registerPiCodingAgentInstrumentation } from "./pi-coding-agent-instrumentation"; const PROMPT_CHANNEL = "orchestrion:@earendil-works/pi-coding-agent:AgentSession.prompt"; @@ -43,12 +43,10 @@ type TestSpan = { name?: string; }; -describe("PiCodingAgentPlugin", () => { - let plugins: PiCodingAgentPlugin[]; +describe("registerPiCodingAgentInstrumentation", () => { let spans: TestSpan[]; beforeEach(() => { - plugins = []; spans = []; interceptorsByName.clear(); mockNewTracingChannel.mockImplementation((name: string) => ({ @@ -74,25 +72,19 @@ describe("PiCodingAgentPlugin", () => { }); afterEach(() => { - for (const plugin of plugins) { - plugin.disable(); - } vi.clearAllMocks(); }); - it("registers and removes the prompt interceptor", () => { - const plugin = enablePlugin(plugins); + it("registers the prompt interceptor for the process lifetime", () => { + registerPiCodingAgentInstrumentation(); expect(interceptorsByName.has(PROMPT_CHANNEL)).toBe(true); expect(mockRemoveInterceptor).not.toHaveBeenCalled(); - - plugin.disable(); - - expect(mockRemoveInterceptor).toHaveBeenCalledTimes(1); }); it("uses interceptor ALS for prompt, LLM, and tool spans", async () => { - const interceptor = promptInterceptor(enablePlugin(plugins)); + registerPiCodingAgentInstrumentation(); + const interceptor = promptInterceptor(); const finalMessage = makeAssistantMessage("done"); const originalStreamFn = vi.fn(async () => { expect(isAutoInstrumentationSuppressed()).toBe(true); @@ -200,7 +192,8 @@ describe("PiCodingAgentPlugin", () => { }); it("isolates overlapping prompts on the same agent without prompt matching", async () => { - const interceptor = promptInterceptor(enablePlugin(plugins)); + registerPiCodingAgentInstrumentation(); + const interceptor = promptInterceptor(); const originalStreamFn = vi.fn(async () => makeStream(makeAssistantMessage("done")), ); @@ -258,7 +251,8 @@ describe("PiCodingAgentPlugin", () => { }); it("keeps deferred follow-up prompts open for their ALS-owned turn", async () => { - const interceptor = promptInterceptor(enablePlugin(plugins)); + registerPiCodingAgentInstrumentation(); + const interceptor = promptInterceptor(); const agent = makeAgent( vi.fn(async () => makeStream(makeAssistantMessage("done"))), ); @@ -296,7 +290,8 @@ describe("PiCodingAgentPlugin", () => { }); it("preserves full-iterator stream behavior", async () => { - const interceptor = promptInterceptor(enablePlugin(plugins)); + registerPiCodingAgentInstrumentation(); + const interceptor = promptInterceptor(); const message = makeAssistantMessage("done"); const { result, stream } = makeIteratorBackedStream([ { partial: message, type: "start" }, @@ -331,7 +326,8 @@ describe("PiCodingAgentPlugin", () => { }); it("forwards iterator failures and closes the underlying iterator", async () => { - const interceptor = promptInterceptor(enablePlugin(plugins)); + registerPiCodingAgentInstrumentation(); + const interceptor = promptInterceptor(); const { iterator, stream } = makeIteratorBackedStream([]); iterator.next.mockRejectedValueOnce(new Error("stream next failed")); const agent = makeAgent(vi.fn(async () => stream)); @@ -364,7 +360,8 @@ describe("PiCodingAgentPlugin", () => { }); it("closes active prompt and LLM spans when the target rejects", async () => { - const interceptor = promptInterceptor(enablePlugin(plugins)); + registerPiCodingAgentInstrumentation(); + const interceptor = promptInterceptor(); const agent = makeAgent( vi.fn(async () => makeStream(makeAssistantMessage("unused"))), ); @@ -396,9 +393,9 @@ describe("PiCodingAgentPlugin", () => { expect(taskSpan?.end).toHaveBeenCalledTimes(1); }); - it("ends in-flight prompts on disable while leaving the safe patch installed", async () => { - const plugin = enablePlugin(plugins); - const interceptor = promptInterceptor(plugin); + it("keeps in-flight prompts active until the call completes", async () => { + registerPiCodingAgentInstrumentation(); + const interceptor = promptInterceptor(); const originalStreamFn = vi.fn(async () => makeStream(makeAssistantMessage("done")), ); @@ -416,12 +413,10 @@ describe("PiCodingAgentPlugin", () => { ); const patchedStreamFn = agent.streamFn; - plugin.disable(); - const taskSpan = findSpan(spans, "AgentSession.prompt"); expect(agent.streamFn).toBe(patchedStreamFn); expect(agent.streamFn).not.toBe(originalStreamFn); - expect(taskSpan?.end).toHaveBeenCalledTimes(1); + expect(taskSpan?.end).not.toHaveBeenCalled(); targetGate.resolve(); await result; @@ -429,14 +424,7 @@ describe("PiCodingAgentPlugin", () => { }); }); -function enablePlugin(plugins: PiCodingAgentPlugin[]): PiCodingAgentPlugin { - const plugin = new PiCodingAgentPlugin(); - plugins.push(plugin); - plugin.enable(); - return plugin; -} - -function promptInterceptor(_plugin: PiCodingAgentPlugin): any { +function promptInterceptor(): any { const interceptor = interceptorsByName.get(PROMPT_CHANNEL); expect(interceptor).toEqual(expect.any(Function)); return interceptor; diff --git a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts b/js/src/instrumentation/providers/pi-coding-agent-instrumentation.ts similarity index 93% rename from js/src/instrumentation/plugins/pi-coding-agent-plugin.ts rename to js/src/instrumentation/providers/pi-coding-agent-instrumentation.ts index da6c50b30..11829e004 100644 --- a/js/src/instrumentation/plugins/pi-coding-agent-plugin.ts +++ b/js/src/instrumentation/providers/pi-coding-agent-instrumentation.ts @@ -1,4 +1,4 @@ -import { BasePlugin, toLoggedError } from "../core"; +import { toLoggedError } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import iso, { type IsoAsyncLocalStorage } from "../../isomorph"; import { debugLogger } from "../../debug-logger"; @@ -44,7 +44,6 @@ type PiPromptState = { deferCompletionUntilTurnEnd: boolean; finalized: boolean; metrics: Record; - onFinalize?: (state: PiPromptState) => void; metadata: Record; output?: unknown; promptCallEnded: boolean; @@ -77,74 +76,52 @@ let piPromptContextStore: | IsoAsyncLocalStorage | undefined; -export class PiCodingAgentPlugin extends BasePlugin { - private readonly activePromptStates = new Set(); - - protected onEnable(): void { +class PiCodingAgentInstrumentationConsumer { + public register(): void { this.interceptPrompt(); } - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - - for (const state of [...this.activePromptStates]) { - finishPiPromptRun(state); - } - } - private interceptPrompt(): void { - this.unsubscribers.push( - piCodingAgentChannels.prompt.intercept( - (target, thisArg, args, additional) => { - const invokeTarget = () => Reflect.apply(target, thisArg, args); - let state: PiPromptState | undefined; - try { - state = startPiPromptRun( - { - ...additional, - arguments: args, - self: thisArg, - }, - (finalizedState) => { - this.activePromptStates.delete(finalizedState); - }, - ); - } catch (error) { - logInstrumentationError("Pi Coding Agent prompt start", error); - } + void piCodingAgentChannels.prompt.intercept( + (target, thisArg, args, additional) => { + const invokeTarget = () => Reflect.apply(target, thisArg, args); + let state: PiPromptState | undefined; + try { + state = startPiPromptRun({ + ...additional, + arguments: args, + self: thisArg, + }); + } catch (error) { + logInstrumentationError("Pi Coding Agent prompt start", error); + } - if (!state) { - return runWithAutoInstrumentationSuppressed(invokeTarget); - } - this.activePromptStates.add(state); + if (!state) { + return runWithAutoInstrumentationSuppressed(invokeTarget); + } + return promptContextStore().run(state, () => + runWithAutoInstrumentationSuppressed(() => { + let result: PromiseLike; + try { + result = invokeTarget(); + } catch (error) { + finishPiPromptRun(state, error); + throw error; + } - return promptContextStore().run(state, () => - runWithAutoInstrumentationSuppressed(() => { - let result: PromiseLike; - try { - result = invokeTarget(); - } catch (error) { + return Promise.resolve(result).then( + (value) => { + finishPiPromptCall(state); + return value; + }, + (error) => { finishPiPromptRun(state, error); throw error; - } - - return Promise.resolve(result).then( - (value) => { - finishPiPromptCall(state); - return value; - }, - (error) => { - finishPiPromptRun(state, error); - throw error; - }, - ); - }), - ); - }, - ), + }, + ); + }), + ); + }, ); } } @@ -159,7 +136,6 @@ function finishPiPromptCall(state: PiPromptState): void { function startPiPromptRun( event: ChannelMessage, - onFinalize?: (state: PiPromptState) => void, ): PiPromptState | undefined { const session = extractSession(event); const agent = session?.agent; @@ -206,7 +182,6 @@ function startPiPromptRun( finalized: false, metadata, metrics: {}, - onFinalize, promptCallEnded: false, span, startTime: getCurrentUnixTimestamp(), @@ -595,7 +570,6 @@ function finishPiPromptRun(state: PiPromptState, error?: unknown): void { return; } state.finalized = true; - state.onFinalize?.(state); finishOpenLlmSpans(state, error); finishOpenToolSpans(state, error); @@ -1085,3 +1059,13 @@ function safeLog(span: Span, event: Parameters[0]): void { function logInstrumentationError(context: string, error: unknown): void { debugLogger.debug(`${context}:`, error); } + +let piCodingAgentInstrumentationConsumer: + | PiCodingAgentInstrumentationConsumer + | undefined; + +export function registerPiCodingAgentInstrumentation(): void { + piCodingAgentInstrumentationConsumer ??= + new PiCodingAgentInstrumentationConsumer(); + piCodingAgentInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/strands-agent-sdk-channels.ts b/js/src/instrumentation/providers/strands-agent-sdk-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/strands-agent-sdk-channels.ts rename to js/src/instrumentation/providers/strands-agent-sdk-channels.ts diff --git a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts b/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.test.ts similarity index 95% rename from js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts rename to js/src/instrumentation/providers/strands-agent-sdk-instrumentation.test.ts index 54340d1a3..5c0986ca8 100644 --- a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.test.ts +++ b/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.test.ts @@ -5,7 +5,6 @@ const { mockWithCurrent, mockNewAsyncLocalStorage, mockStartSpan, - mockUnbindStore, } = vi.hoisted(() => ({ mockBindStore: vi.fn(), mockWithCurrent: vi.fn(), @@ -28,7 +27,6 @@ const { }; }), mockStartSpan: vi.fn(), - mockUnbindStore: vi.fn(), })); vi.mock("../../isomorph", () => ({ @@ -51,11 +49,11 @@ vi.mock("../../logger", async (importOriginal) => { import iso from "../../isomorph"; import { Attachment } from "../../logger"; import { isAutoInstrumentationSuppressed } from "../auto-instrumentation-suppression"; -import { StrandsAgentSDKPlugin } from "./strands-agent-sdk-plugin"; +import { registerStrandsAgentSDKInstrumentation } from "./strands-agent-sdk-instrumentation"; const mockNewTracingChannel = iso.newTracingChannel as ReturnType; -describe("StrandsAgentSDKPlugin", () => { +describe("registerStrandsAgentSDKInstrumentation", () => { let handlersByName: Map; let spans: Array<{ args: any; @@ -77,11 +75,9 @@ describe("StrandsAgentSDKPlugin", () => { mockNewTracingChannel.mockImplementation((name: string) => ({ start: { bindStore: mockBindStore, - unbindStore: mockUnbindStore, }, subscribe: vi.fn((handlers) => handlersByName.set(name, handlers)), traceSync: vi.fn((fn) => fn()), - unsubscribe: vi.fn(), })); currentSpan = undefined; mockWithCurrent.mockImplementation((span: any, callback: () => unknown) => { @@ -123,8 +119,7 @@ describe("StrandsAgentSDKPlugin", () => { }); it("subscribes to Strands stream channels and binds suppression", () => { - const plugin = new StrandsAgentSDKPlugin(); - plugin.enable(); + registerStrandsAgentSDKInstrumentation(); expect( handlersByName.has("orchestrion:@strands-agents/sdk:Agent.stream"), @@ -136,15 +131,10 @@ describe("StrandsAgentSDKPlugin", () => { handlersByName.has("orchestrion:@strands-agents/sdk:Swarm.stream"), ).toBe(true); expect(mockBindStore).toHaveBeenCalledTimes(3); - - plugin.disable(); - - expect(mockUnbindStore).toHaveBeenCalledTimes(3); }); it("records agent model and tool spans from stream events", async () => { - const plugin = new StrandsAgentSDKPlugin(); - plugin.enable(); + registerStrandsAgentSDKInstrumentation(); const handlers = handlersByName.get( "orchestrion:@strands-agents/sdk:Agent.stream", @@ -292,8 +282,7 @@ describe("StrandsAgentSDKPlugin", () => { ])( "converts media from %s to one attachment shared by agent and model spans", async (_description, bytes) => { - const plugin = new StrandsAgentSDKPlugin(); - plugin.enable(); + registerStrandsAgentSDKInstrumentation(); const handlers = handlersByName.get( "orchestrion:@strands-agents/sdk:Agent.stream", @@ -353,8 +342,7 @@ describe("StrandsAgentSDKPlugin", () => { ); it("parents nested agent spans under active graph nodes", async () => { - const plugin = new StrandsAgentSDKPlugin(); - plugin.enable(); + registerStrandsAgentSDKInstrumentation(); const graphHandlers = handlersByName.get( "orchestrion:@strands-agents/sdk:Graph.stream", @@ -441,8 +429,7 @@ describe("StrandsAgentSDKPlugin", () => { }); it("parents nested graph spans under active graph nodes", async () => { - const plugin = new StrandsAgentSDKPlugin(); - plugin.enable(); + registerStrandsAgentSDKInstrumentation(); const graphHandlers = handlersByName.get( "orchestrion:@strands-agents/sdk:Graph.stream", @@ -530,8 +517,7 @@ describe("StrandsAgentSDKPlugin", () => { }); it("ends open child spans when a stream errors", async () => { - const plugin = new StrandsAgentSDKPlugin(); - plugin.enable(); + registerStrandsAgentSDKInstrumentation(); const handlers = handlersByName.get( "orchestrion:@strands-agents/sdk:Agent.stream", diff --git a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts b/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.ts similarity index 97% rename from js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts rename to js/src/instrumentation/providers/strands-agent-sdk-instrumentation.ts index 255578a13..5142f2d5c 100644 --- a/js/src/instrumentation/plugins/strands-agent-sdk-plugin.ts +++ b/js/src/instrumentation/providers/strands-agent-sdk-instrumentation.ts @@ -1,4 +1,4 @@ -import { BasePlugin, toLoggedError } from "../core"; +import { toLoggedError } from "../core"; import type { ChannelMessage } from "../core/channel-definitions"; import { isAsyncIterable, patchStreamIfNeeded } from "../core/stream-patcher"; import type { IsoChannelHandlers } from "../../isomorph"; @@ -101,10 +101,10 @@ type StrandsAttachmentCache = { strings: LRUCache>; }; -export class StrandsAgentSDKPlugin extends BasePlugin { +class StrandsAgentSDKInstrumentationConsumer { private readonly activeChildParents: ActiveChildParents = new WeakMap(); - protected onEnable(): void { + public register(): void { this.subscribeToAgentStream(); this.subscribeToMultiAgentStream( strandsAgentSDKChannels.graphStream, @@ -116,18 +116,10 @@ export class StrandsAgentSDKPlugin extends BasePlugin { ); } - protected onDisable(): void { - for (const unsubscribe of this.unsubscribers) { - unsubscribe(); - } - this.unsubscribers = []; - } - private subscribeToAgentStream(): void { const channel = strandsAgentSDKChannels.agentStream.tracingChannel(); const states = new WeakMap(); - const unbindAutoInstrumentationSuppression = - bindAutoInstrumentationSuppressionToStart(channel); + bindAutoInstrumentationSuppressionToStart(channel); const handlers: IsoChannelHandlers< ChannelMessage @@ -176,10 +168,6 @@ export class StrandsAgentSDKPlugin extends BasePlugin { }; channel.subscribe(handlers); - this.unsubscribers.push(() => { - unbindAutoInstrumentationSuppression?.(); - channel.unsubscribe(handlers); - }); } private subscribeToMultiAgentStream( @@ -188,8 +176,7 @@ export class StrandsAgentSDKPlugin extends BasePlugin { ): void { const tracingChannel = channel.tracingChannel(); const states = new WeakMap(); - const unbindAutoInstrumentationSuppression = - bindAutoInstrumentationSuppressionToStart(tracingChannel); + bindAutoInstrumentationSuppressionToStart(tracingChannel); const handlers: IsoChannelHandlers> = { start: (event) => { @@ -250,10 +237,6 @@ export class StrandsAgentSDKPlugin extends BasePlugin { }; tracingChannel.subscribe(handlers); - this.unsubscribers.push(() => { - unbindAutoInstrumentationSuppression?.(); - tracingChannel.unsubscribe(handlers); - }); } } @@ -1393,3 +1376,13 @@ function safeLog(span: Span, event: Parameters[0]): void { function logInstrumentationError(context: string, error: unknown): void { debugLogger.debug(`${context}:`, error); } + +let strandsAgentSDKInstrumentationConsumer: + | StrandsAgentSDKInstrumentationConsumer + | undefined; + +export function registerStrandsAgentSDKInstrumentation(): void { + strandsAgentSDKInstrumentationConsumer ??= + new StrandsAgentSDKInstrumentationConsumer(); + strandsAgentSDKInstrumentationConsumer.register(); +} diff --git a/js/src/instrumentation/plugins/voyageai-channels.ts b/js/src/instrumentation/providers/voyageai-channels.ts similarity index 100% rename from js/src/instrumentation/plugins/voyageai-channels.ts rename to js/src/instrumentation/providers/voyageai-channels.ts diff --git a/js/src/instrumentation/plugins/voyageai-plugin.test.ts b/js/src/instrumentation/providers/voyageai-instrumentation.test.ts similarity index 98% rename from js/src/instrumentation/plugins/voyageai-plugin.test.ts rename to js/src/instrumentation/providers/voyageai-instrumentation.test.ts index 65cfca93c..2cdb491b4 100644 --- a/js/src/instrumentation/plugins/voyageai-plugin.test.ts +++ b/js/src/instrumentation/providers/voyageai-instrumentation.test.ts @@ -9,7 +9,7 @@ try { // Best-effort initialization for test environments. } -describe("VoyageAIPlugin", () => { +describe("registerVoyageAIInstrumentation", () => { let backgroundLogger: ReturnType< typeof _exportsForTestingOnly.useTestBackgroundLogger >; @@ -21,7 +21,7 @@ describe("VoyageAIPlugin", () => { beforeEach(() => { backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); initLogger({ - projectName: "voyageai-plugin.test.ts", + projectName: "voyageai-instrumentation.test.ts", projectId: "test-project-id", }); }); diff --git a/js/src/instrumentation/plugins/voyageai-plugin.ts b/js/src/instrumentation/providers/voyageai-instrumentation.ts similarity index 89% rename from js/src/instrumentation/plugins/voyageai-plugin.ts rename to js/src/instrumentation/providers/voyageai-instrumentation.ts index 37f864b26..ac2304e81 100644 --- a/js/src/instrumentation/plugins/voyageai-plugin.ts +++ b/js/src/instrumentation/providers/voyageai-instrumentation.ts @@ -17,8 +17,6 @@ import { isAutoInstrumentationSuppressed, runWithAutoInstrumentationSuppressed, } from "../auto-instrumentation-suppression"; -import { BasePlugin } from "../core"; -import { unsubscribeAll } from "../core/channel-tracing"; import { voyageAIChannels } from "./voyageai-channels"; const RERANK_METADATA_ALLOWLIST = new Set([ @@ -28,42 +26,34 @@ const RERANK_METADATA_ALLOWLIST = new Set([ "truncation", ]); -export class VoyageAIPlugin extends BasePlugin { - protected onEnable(): void { - this.unsubscribers.push( - interceptVoyageAICall( - voyageAIChannels.embed, - "voyageai.embed", - extractTextEmbeddingInput, - summarizeEmbeddingOutput, - extractEmbeddingUsageMetrics, - ), - interceptVoyageAICall( - voyageAIChannels.multimodalEmbed, - "voyageai.multimodalEmbed", - extractMultimodalEmbeddingInput, - summarizeEmbeddingOutput, - extractEmbeddingUsageMetrics, - ), - interceptVoyageAICall( - voyageAIChannels.rerank, - "voyageai.rerank", - extractRerankInput, - summarizeRerankOutput, - ), - interceptVoyageAICall( - voyageAIChannels.contextualizedEmbed, - "voyageai.contextualizedEmbed", - extractContextualizedEmbeddingInput, - summarizeContextualizedEmbeddingOutput, - extractEmbeddingUsageMetrics, - ), - ); - } - - protected onDisable(): void { - this.unsubscribers = unsubscribeAll(this.unsubscribers); - } +export function registerVoyageAIInstrumentation(): void { + interceptVoyageAICall( + voyageAIChannels.embed, + "voyageai.embed", + extractTextEmbeddingInput, + summarizeEmbeddingOutput, + extractEmbeddingUsageMetrics, + ); + interceptVoyageAICall( + voyageAIChannels.multimodalEmbed, + "voyageai.multimodalEmbed", + extractMultimodalEmbeddingInput, + summarizeEmbeddingOutput, + extractEmbeddingUsageMetrics, + ); + interceptVoyageAICall( + voyageAIChannels.rerank, + "voyageai.rerank", + extractRerankInput, + summarizeRerankOutput, + ); + interceptVoyageAICall( + voyageAIChannels.contextualizedEmbed, + "voyageai.contextualizedEmbed", + extractContextualizedEmbeddingInput, + summarizeContextualizedEmbeddingOutput, + extractEmbeddingUsageMetrics, + ); } type VoyageAIResult = @@ -455,13 +445,10 @@ function extractUsageMetrics(result: VoyageAIResult): Record { const rawResponse = isObject(result.rawResponse) ? result.rawResponse : undefined; - const usage = ( - isObject(result.usage) - ? result.usage - : isObject(rawResponse?.usage) - ? rawResponse.usage - : undefined - ) as VoyageAIUsage | undefined; + const usageValue = isObject(result.usage) ? result.usage : rawResponse?.usage; + const usage = isObject(usageValue) + ? (usageValue as VoyageAIUsage) + : undefined; const tokens = typeof result.totalTokens === "number" ? result.totalTokens diff --git a/js/src/instrumentation/registry.test.ts b/js/src/instrumentation/registry.test.ts index fdd49a5c2..97d325cea 100644 --- a/js/src/instrumentation/registry.test.ts +++ b/js/src/instrumentation/registry.test.ts @@ -19,32 +19,25 @@ import iso from "../isomorph"; const mockNewTracingChannel = iso.newTracingChannel as ReturnType; -describe("Plugin Registry", () => { +describe("Instrumentation Registry", () => { beforeEach(() => { // Setup mock channel const mockChannel = { subscribe: vi.fn(), - unsubscribe: vi.fn(), hasSubscribers: false, }; mockNewTracingChannel.mockReturnValue(mockChannel); }); - // Clean up after each test - afterEach(() => { - registry.disable(); - }); - it("should not be enabled by default", () => { const testRegistry = new (registry.constructor as any)(); expect(testRegistry.isEnabled()).toBe(false); }); - it("should enable plugins when enable() is called", () => { + it("should enable instrumentation consumers when enable() is called", () => { const testRegistry = new (registry.constructor as any)(); testRegistry.enable(); expect(testRegistry.isEnabled()).toBe(true); - testRegistry.disable(); }); it("should be idempotent (calling enable() multiple times)", () => { @@ -52,12 +45,11 @@ describe("Plugin Registry", () => { testRegistry.enable(); testRegistry.enable(); // Should not throw expect(testRegistry.isEnabled()).toBe(true); - testRegistry.disable(); }); it("should block a second instance from subscribing when another is already enabled", () => { // Regression test for BT-5139: when the SDK is loaded from two different - // module paths in the same process, each gets its own PluginRegistry + // module paths in the same process, each gets its own instrumentation registry // instance. Without cross-instance deduplication, both would subscribe to // the same global hook, causing every OpenAI call to produce two // LLM spans. @@ -80,8 +72,6 @@ describe("Plugin Registry", () => { instanceB.enable(); expect(instanceB.isEnabled()).toBe(false); } finally { - instanceA.disable(); - instanceB.disable(); delete (globalThis as any)[stateKey]; } }); @@ -97,7 +87,6 @@ describe("Plugin Registry", () => { testRegistry.enable(); expect(testRegistry.isEnabled()).toBe(true); } finally { - testRegistry.disable(); delete (globalThis as any)[stateKey]; } }); @@ -116,7 +105,6 @@ describe("Plugin Registry", () => { expect(warnSpy[0]).toContain("Cannot configure instrumentation"); } finally { console.warn = originalWarn; - testRegistry.disable(); } }); @@ -133,33 +121,11 @@ describe("Plugin Registry", () => { expect(warnSpy.length).toBe(0); } finally { console.warn = originalWarn; - testRegistry.disable(); } }); - - it("should disable plugins when disable() is called", () => { - const testRegistry = new (registry.constructor as any)(); - testRegistry.enable(); - expect(testRegistry.isEnabled()).toBe(true); - - testRegistry.disable(); - expect(testRegistry.isEnabled()).toBe(false); - }); - - it("should be idempotent (calling disable() multiple times)", () => { - const testRegistry = new (registry.constructor as any)(); - testRegistry.enable(); - testRegistry.disable(); - testRegistry.disable(); // Should not throw - expect(testRegistry.isEnabled()).toBe(false); - }); }); describe("configureInstrumentation API", () => { - afterEach(() => { - registry.disable(); - }); - it("should export configureInstrumentation function", () => { expect(typeof configureInstrumentation).toBe("function"); }); @@ -195,7 +161,6 @@ describe("Environment Variable Configuration", () => { afterEach(async () => { const iso = (await import("../isomorph")).default; iso.getEnv = originalGetEnv; - registry.disable(); }); it("should parse BRAINTRUST_DISABLE_INSTRUMENTATION with single SDK", async () => { @@ -212,7 +177,6 @@ describe("Environment Variable Configuration", () => { // OpenAI should be disabled, others enabled by default expect(testRegistry.isEnabled()).toBe(true); - testRegistry.disable(); }); it("should parse BRAINTRUST_DISABLE_INSTRUMENTATION with multiple SDKs", async () => { @@ -229,7 +193,6 @@ describe("Environment Variable Configuration", () => { // Both should be disabled expect(testRegistry.isEnabled()).toBe(true); - testRegistry.disable(); }); it("should handle whitespace in BRAINTRUST_DISABLE_INSTRUMENTATION", async () => { @@ -245,7 +208,6 @@ describe("Environment Variable Configuration", () => { testRegistry.enable(); expect(testRegistry.isEnabled()).toBe(true); - testRegistry.disable(); }); it("should handle empty BRAINTRUST_DISABLE_INSTRUMENTATION", async () => { @@ -262,7 +224,6 @@ describe("Environment Variable Configuration", () => { // All should be enabled (nothing disabled) expect(testRegistry.isEnabled()).toBe(true); - testRegistry.disable(); }); it("should be case-insensitive for SDK names", async () => { @@ -278,6 +239,5 @@ describe("Environment Variable Configuration", () => { testRegistry.enable(); expect(testRegistry.isEnabled()).toBe(true); - testRegistry.disable(); }); }); diff --git a/js/src/instrumentation/registry.ts b/js/src/instrumentation/registry.ts index 0ed195ae8..010aee634 100644 --- a/js/src/instrumentation/registry.ts +++ b/js/src/instrumentation/registry.ts @@ -1,11 +1,11 @@ /** - * Plugin registry and configuration for auto-instrumentation. + * Registry and configuration for auto-instrumentation. * - * Plugins are automatically enabled when the Braintrust library is loaded. + * Instrumentation consumers are automatically enabled when Braintrust loads. * Users can disable specific integrations programmatically or via environment variables. */ -import { BraintrustPlugin } from "./braintrust-plugin"; +import { registerInstrumentationConsumers } from "./instrumentation-consumers"; import iso from "../isomorph"; import { getDefaultInstrumentationIntegrations, @@ -16,7 +16,7 @@ import { GLOBAL_INSTRUMENTATION_HOOKS_PROTOCOL_VERSION } from "../global-instrum export type { InstrumentationConfig } from "./config"; -// Key used to stamp the active PluginRegistry instance onto the shared +// Key used to stamp the active InstrumentationRegistry instance onto the shared // braintrust state object (globalThis[Symbol.for("braintrust-state")]). // // The braintrust state is already shared across all SDK instances loaded in @@ -43,8 +43,7 @@ function getSharedState(): Record | undefined { : undefined; } -class PluginRegistry { - private braintrustPlugin: BraintrustPlugin | null = null; +class InstrumentationRegistry { private config: InstrumentationConfig = {}; private enabled = false; @@ -65,7 +64,7 @@ class PluginRegistry { } /** - * Enable all configured plugins. + * Enable all configured instrumentation consumers. * Called automatically when the library is loaded. */ enable(): void { @@ -73,7 +72,7 @@ class PluginRegistry { return; } - // If another SDK instance in the same process already registered plugins, + // If another SDK instance already registered instrumentation consumers, // skip to avoid duplicate global hook subscriptions. const sharedState = getSharedState(); if (sharedState) { @@ -95,31 +94,8 @@ class PluginRegistry { }, }; - // Enable BraintrustPlugin with configuration - this.braintrustPlugin = new BraintrustPlugin(finalConfig); - this.braintrustPlugin.enable(); - } - - /** - * Disable all plugins. - * Primarily used for testing. - */ - disable(): void { - if (!this.enabled) { - return; - } - - this.enabled = false; - - const sharedState = getSharedState(); - if (sharedState && sharedState[REGISTRY_STATE_KEY] === this) { - delete sharedState[REGISTRY_STATE_KEY]; - } - - if (this.braintrustPlugin) { - this.braintrustPlugin.disable(); - this.braintrustPlugin = null; - } + // Enable the configured instrumentation consumers. + registerInstrumentationConsumers(finalConfig); } /** @@ -148,9 +124,9 @@ class PluginRegistry { } /** - * Global plugin registry instance. + * Global instrumentation registry instance. */ -export const registry = new PluginRegistry(); +export const registry = new InstrumentationRegistry(); /** * Configure auto-instrumentation. diff --git a/js/src/wrappers/ai-sdk/ai-sdk.ts b/js/src/wrappers/ai-sdk/ai-sdk.ts index 63de58ea6..dc6c7b09d 100644 --- a/js/src/wrappers/ai-sdk/ai-sdk.ts +++ b/js/src/wrappers/ai-sdk/ai-sdk.ts @@ -4,7 +4,7 @@ import { SpanTypeAttribute } from "../../../util"; import { aiSDKChannels, harnessAgentChannels, -} from "../../instrumentation/plugins/ai-sdk-channels"; +} from "../../instrumentation/providers/ai-sdk-channels"; import type { AISDK, AISDKAgentClass, diff --git a/js/src/wrappers/ai-sdk/harness-agent-context.ts b/js/src/wrappers/ai-sdk/harness-agent-context.ts index fc0b6171a..f17f2a442 100644 --- a/js/src/wrappers/ai-sdk/harness-agent-context.ts +++ b/js/src/wrappers/ai-sdk/harness-agent-context.ts @@ -335,10 +335,10 @@ export function currentHarnessTurnParent(): HarnessTurnParent | undefined { export function bindHarnessTurnParentToStart( tracingChannel: IsoTracingChannel, parentFromEvent: (event: T) => HarnessTurnParent | undefined, -): () => void { +): void { const startChannel = tracingChannel.start; if (!startChannel) { - return () => {}; + return; } harnessTurnParentStore ??= iso.newAsyncLocalStorage< @@ -349,9 +349,6 @@ export function bindHarnessTurnParentToStart( store, (event) => parentFromEvent(event) ?? store.getStore(), ); - return () => { - startChannel.unbindStore(store); - }; } export function startHarnessTurnChildSpan( diff --git a/js/src/wrappers/ai-sdk/telemetry.ts b/js/src/wrappers/ai-sdk/telemetry.ts index 1fb2c9812..110f8e221 100644 --- a/js/src/wrappers/ai-sdk/telemetry.ts +++ b/js/src/wrappers/ai-sdk/telemetry.ts @@ -20,7 +20,7 @@ import { processAISDKWorkflowAgentCallInput, processAISDKWorkflowAgentModelCallInput, serializeModelWithProvider, -} from "../../instrumentation/plugins/ai-sdk-plugin"; +} from "../../instrumentation/providers/ai-sdk-instrumentation"; import type { AISDKCallParams, AISDKEmbeddingResult, diff --git a/js/src/wrappers/anthropic.ts b/js/src/wrappers/anthropic.ts index 12aae0218..e4cfc708f 100644 --- a/js/src/wrappers/anthropic.ts +++ b/js/src/wrappers/anthropic.ts @@ -1,4 +1,4 @@ -import { anthropicChannels } from "../instrumentation/plugins/anthropic-channels"; +import { anthropicChannels } from "../instrumentation/providers/anthropic-channels"; import { TypedApplyProxy } from "../typed-instrumentation-helpers"; import type { AnthropicBeta, diff --git a/js/src/wrappers/bedrock-runtime.ts b/js/src/wrappers/bedrock-runtime.ts index 13fcd1844..4e60fd912 100644 --- a/js/src/wrappers/bedrock-runtime.ts +++ b/js/src/wrappers/bedrock-runtime.ts @@ -1,9 +1,9 @@ import { runWithAutoInstrumentationSuppressed } from "../instrumentation/auto-instrumentation-suppression"; -import { bedrockRuntimeChannels } from "../instrumentation/plugins/bedrock-runtime-channels"; +import { bedrockRuntimeChannels } from "../instrumentation/providers/bedrock-runtime-channels"; import { buildBedrockRuntimeSpanInfo, getBedrockRuntimeOperation, -} from "../instrumentation/plugins/bedrock-runtime-common"; +} from "../instrumentation/providers/bedrock-runtime-common"; import type { BedrockRuntimeClient, BedrockRuntimeCommandLike, diff --git a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts index db24f0bcf..fae63e64d 100644 --- a/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts +++ b/js/src/wrappers/claude-agent-sdk/claude-agent-sdk.ts @@ -1,6 +1,6 @@ -import { claudeAgentSDKChannels } from "../../instrumentation/plugins/claude-agent-sdk-channels"; -import { CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION } from "../../instrumentation/plugins/claude-agent-sdk-instrumentation-constants"; -import { wrapLocalClaudeToolHandler } from "../../instrumentation/plugins/claude-agent-sdk-local-tool-spans"; +import { claudeAgentSDKChannels } from "../../instrumentation/providers/claude-agent-sdk-channels"; +import { CLAUDE_AGENT_SDK_SKIP_LOCAL_TOOL_HOOKS_OPTION } from "../../instrumentation/providers/claude-agent-sdk-instrumentation-constants"; +import { wrapLocalClaudeToolHandler } from "../../instrumentation/providers/claude-agent-sdk-local-tool-spans"; import type { ClaudeAgentSDKModule, ClaudeAgentSDKQueryParams, diff --git a/js/src/wrappers/cloudflare-agent.ts b/js/src/wrappers/cloudflare-agent.ts index c34bb6134..baee1d63d 100644 --- a/js/src/wrappers/cloudflare-agent.ts +++ b/js/src/wrappers/cloudflare-agent.ts @@ -1,5 +1,5 @@ import { debugLogger } from "../debug-logger"; -import { cloudflareAgentsChannels } from "../instrumentation/plugins/cloudflare-agents-channels"; +import { cloudflareAgentsChannels } from "../instrumentation/providers/cloudflare-agents-channels"; import type { CloudflareAgent } from "../vendor-sdk-types/cloudflare-agents"; const WRAPPED_RUN_AGENT_TOOL = Symbol.for( diff --git a/js/src/wrappers/cloudflare-ai-chat.ts b/js/src/wrappers/cloudflare-ai-chat.ts index 87df27775..11f834feb 100644 --- a/js/src/wrappers/cloudflare-ai-chat.ts +++ b/js/src/wrappers/cloudflare-ai-chat.ts @@ -1,5 +1,5 @@ import { debugLogger } from "../debug-logger"; -import { instrumentCloudflareAIChatAgent } from "../instrumentation/plugins/cloudflare-ai-chat-instrumentation"; +import { instrumentCloudflareAIChatAgent } from "../instrumentation/providers/cloudflare-ai-chat-instrumentation"; import type { CloudflareAIChatAgent, CloudflareAIChatAgentConstructor, diff --git a/js/src/wrappers/cloudflare-think.ts b/js/src/wrappers/cloudflare-think.ts index 9c04400ce..2b1e371bd 100644 --- a/js/src/wrappers/cloudflare-think.ts +++ b/js/src/wrappers/cloudflare-think.ts @@ -1,4 +1,4 @@ -import { cloudflareThinkChannels } from "../instrumentation/plugins/cloudflare-think-channels"; +import { cloudflareThinkChannels } from "../instrumentation/providers/cloudflare-think-channels"; import type { CloudflareThinkConstructor, CloudflareThinkInstance, diff --git a/js/src/wrappers/cohere.ts b/js/src/wrappers/cohere.ts index aa6038030..6c55f65dd 100644 --- a/js/src/wrappers/cohere.ts +++ b/js/src/wrappers/cohere.ts @@ -1,4 +1,4 @@ -import { cohereChannels } from "../instrumentation/plugins/cohere-channels"; +import { cohereChannels } from "../instrumentation/providers/cohere-channels"; import type { CohereChatRequest, CohereChatResponse, diff --git a/js/src/wrappers/cursor-sdk.ts b/js/src/wrappers/cursor-sdk.ts index e20f41fba..c20897b82 100644 --- a/js/src/wrappers/cursor-sdk.ts +++ b/js/src/wrappers/cursor-sdk.ts @@ -1,4 +1,4 @@ -import { cursorSDKChannels } from "../instrumentation/plugins/cursor-sdk-channels"; +import { cursorSDKChannels } from "../instrumentation/providers/cursor-sdk-channels"; import type { CursorSDKAgent, CursorSDKAgentClass, diff --git a/js/src/wrappers/genkit.test.ts b/js/src/wrappers/genkit.test.ts index 5363a64eb..3677dc99d 100644 --- a/js/src/wrappers/genkit.test.ts +++ b/js/src/wrappers/genkit.test.ts @@ -1,7 +1,7 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import type { IsoChannelHandlers } from "../isomorph"; import type { ChannelMessage } from "../instrumentation/core/channel-definitions"; -import { genkitChannels } from "../instrumentation/plugins/genkit-channels"; +import { genkitChannels } from "../instrumentation/providers/genkit-channels"; import { configureNode } from "../node/config"; import type { GenkitAction, diff --git a/js/src/wrappers/genkit.ts b/js/src/wrappers/genkit.ts index fc35a8426..a4900ae9f 100644 --- a/js/src/wrappers/genkit.ts +++ b/js/src/wrappers/genkit.ts @@ -1,4 +1,4 @@ -import { genkitChannels } from "../instrumentation/plugins/genkit-channels"; +import { genkitChannels } from "../instrumentation/providers/genkit-channels"; import type { GenkitAction, GenkitEmbedManyParams, diff --git a/js/src/wrappers/github-copilot.ts b/js/src/wrappers/github-copilot.ts index 35e68fa87..20e76dbda 100644 --- a/js/src/wrappers/github-copilot.ts +++ b/js/src/wrappers/github-copilot.ts @@ -1,4 +1,4 @@ -import { gitHubCopilotChannels } from "../instrumentation/plugins/github-copilot-channels"; +import { gitHubCopilotChannels } from "../instrumentation/providers/github-copilot-channels"; import type { GitHubCopilotClient, GitHubCopilotResumeSessionConfig, diff --git a/js/src/wrappers/google-adk.ts b/js/src/wrappers/google-adk.ts index 28af7395e..6e908222e 100644 --- a/js/src/wrappers/google-adk.ts +++ b/js/src/wrappers/google-adk.ts @@ -1,4 +1,4 @@ -import { googleADKChannels } from "../instrumentation/plugins/google-adk-channels"; +import { googleADKChannels } from "../instrumentation/providers/google-adk-channels"; import type { GoogleADKRunner, GoogleADKRunnerConstructor, diff --git a/js/src/wrappers/google-genai.ts b/js/src/wrappers/google-genai.ts index 192f79c1a..7d21e8eea 100644 --- a/js/src/wrappers/google-genai.ts +++ b/js/src/wrappers/google-genai.ts @@ -1,4 +1,4 @@ -import { googleGenAIChannels } from "../instrumentation/plugins/google-genai-channels"; +import { googleGenAIChannels } from "../instrumentation/providers/google-genai-channels"; import { isObject } from "../util"; import type { GoogleGenAIClient, diff --git a/js/src/wrappers/groq.ts b/js/src/wrappers/groq.ts index 68312fa81..c736e00c8 100644 --- a/js/src/wrappers/groq.ts +++ b/js/src/wrappers/groq.ts @@ -1,4 +1,4 @@ -import { groqChannels } from "../instrumentation/plugins/groq-channels"; +import { groqChannels } from "../instrumentation/providers/groq-channels"; import type { GroqChat, GroqChatCompletion, diff --git a/js/src/wrappers/huggingface-transformers.ts b/js/src/wrappers/huggingface-transformers.ts index 3965f1234..52d8e2d56 100644 --- a/js/src/wrappers/huggingface-transformers.ts +++ b/js/src/wrappers/huggingface-transformers.ts @@ -1,7 +1,7 @@ import { huggingFaceTransformersChannels, isSupportedHuggingFaceTransformersTask, -} from "../instrumentation/plugins/huggingface-transformers-channels"; +} from "../instrumentation/providers/huggingface-transformers-channels"; import type { HuggingFaceTransformersModule, HuggingFaceTransformersPipeline, diff --git a/js/src/wrappers/huggingface.ts b/js/src/wrappers/huggingface.ts index 6ff98a259..13a58f515 100644 --- a/js/src/wrappers/huggingface.ts +++ b/js/src/wrappers/huggingface.ts @@ -1,4 +1,4 @@ -import { huggingFaceChannels } from "../instrumentation/plugins/huggingface-channels"; +import { huggingFaceChannels } from "../instrumentation/providers/huggingface-channels"; import { isObject } from "../../util"; import type { HuggingFaceChatCompletion, diff --git a/js/src/wrappers/langsmith.ts b/js/src/wrappers/langsmith.ts index a8486c573..983f42dd8 100644 --- a/js/src/wrappers/langsmith.ts +++ b/js/src/wrappers/langsmith.ts @@ -1,5 +1,5 @@ import { debugLogger } from "../debug-logger"; -import { langSmithChannels } from "../instrumentation/plugins/langsmith-channels"; +import { langSmithChannels } from "../instrumentation/providers/langsmith-channels"; import type { LangSmithClient, LangSmithClientConstructor, diff --git a/js/src/wrappers/mistral.ts b/js/src/wrappers/mistral.ts index c449997f7..4771dda38 100644 --- a/js/src/wrappers/mistral.ts +++ b/js/src/wrappers/mistral.ts @@ -1,4 +1,4 @@ -import { mistralChannels } from "../instrumentation/plugins/mistral-channels"; +import { mistralChannels } from "../instrumentation/providers/mistral-channels"; import type { MistralAgents, MistralAgentsCompletionResponse, diff --git a/js/src/wrappers/oai.ts b/js/src/wrappers/oai.ts index b2fa0850e..743460b7e 100644 --- a/js/src/wrappers/oai.ts +++ b/js/src/wrappers/oai.ts @@ -10,7 +10,7 @@ import type { ArgsOf, ResultOf, } from "../instrumentation/core/channel-definitions"; -import { openAIChannels } from "../instrumentation/plugins/openai-channels"; +import { openAIChannels } from "../instrumentation/providers/openai-channels"; import type { OpenAIChatCompletion, OpenAIChatCreateParams, diff --git a/js/src/wrappers/oai_responses.ts b/js/src/wrappers/oai_responses.ts index 8e936ebbc..432b1281a 100644 --- a/js/src/wrappers/oai_responses.ts +++ b/js/src/wrappers/oai_responses.ts @@ -3,7 +3,7 @@ import type { ResultOf, } from "../instrumentation/core/channel-definitions"; import type { ChannelSpanInfo } from "../instrumentation/core/types"; -import { openAIChannels } from "../instrumentation/plugins/openai-channels"; +import { openAIChannels } from "../instrumentation/providers/openai-channels"; import { parseMetricsFromUsage } from "../openai-utils"; import { APIPromise, diff --git a/js/src/wrappers/ollama.test.ts b/js/src/wrappers/ollama.test.ts index 9a8ebff46..b07dac5f4 100644 --- a/js/src/wrappers/ollama.test.ts +++ b/js/src/wrappers/ollama.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it, vi } from "vitest"; import { debugLogger } from "../debug-logger"; -import { ollamaChannels } from "../instrumentation/plugins/ollama-channels"; +import { ollamaChannels } from "../instrumentation/providers/ollama-channels"; import type { OllamaClient } from "../vendor-sdk-types/ollama"; import { wrapOllama } from "./ollama"; diff --git a/js/src/wrappers/ollama.ts b/js/src/wrappers/ollama.ts index 3f90067a9..bcc48a937 100644 --- a/js/src/wrappers/ollama.ts +++ b/js/src/wrappers/ollama.ts @@ -1,5 +1,5 @@ import { debugLogger } from "../debug-logger"; -import { ollamaChannels } from "../instrumentation/plugins/ollama-channels"; +import { ollamaChannels } from "../instrumentation/providers/ollama-channels"; import { isObject } from "../../util"; import type { OllamaChatRequest, diff --git a/js/src/wrappers/openai-codex.ts b/js/src/wrappers/openai-codex.ts index 8b5a435b3..c4225e9c0 100644 --- a/js/src/wrappers/openai-codex.ts +++ b/js/src/wrappers/openai-codex.ts @@ -1,4 +1,4 @@ -import { openAICodexChannels } from "../instrumentation/plugins/openai-codex-channels"; +import { openAICodexChannels } from "../instrumentation/providers/openai-codex-channels"; import type { OpenAICodexClass, OpenAICodexClient, diff --git a/js/src/wrappers/openai-promise-utils.ts b/js/src/wrappers/openai-promise-utils.ts index afa23586c..b41e57f4b 100644 --- a/js/src/wrappers/openai-promise-utils.ts +++ b/js/src/wrappers/openai-promise-utils.ts @@ -6,7 +6,7 @@ import type { OpenAIAsyncChannel, OpenAIChannel, OpenAIStartContext, -} from "../instrumentation/plugins/openai-channels"; +} from "../instrumentation/providers/openai-channels"; export type EnhancedResponse = { response: Response; diff --git a/js/src/wrappers/openrouter-agent.test.ts b/js/src/wrappers/openrouter-agent.test.ts index 02f01e040..c1c94e12f 100644 --- a/js/src/wrappers/openrouter-agent.test.ts +++ b/js/src/wrappers/openrouter-agent.test.ts @@ -1,6 +1,6 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { wrapOpenRouterAgent } from "./openrouter-agent"; -import { openRouterAgentChannels } from "../instrumentation/plugins/openrouter-agent-channels"; +import { openRouterAgentChannels } from "../instrumentation/providers/openrouter-agent-channels"; describe("wrapOpenRouterAgent", () => { afterEach(() => { diff --git a/js/src/wrappers/openrouter-agent.ts b/js/src/wrappers/openrouter-agent.ts index 07328249b..336b2d84f 100644 --- a/js/src/wrappers/openrouter-agent.ts +++ b/js/src/wrappers/openrouter-agent.ts @@ -1,4 +1,4 @@ -import { openRouterAgentChannels } from "../instrumentation/plugins/openrouter-agent-channels"; +import { openRouterAgentChannels } from "../instrumentation/providers/openrouter-agent-channels"; import type { OpenRouterAgentClient, OpenRouterAgentCallModelRequest, diff --git a/js/src/wrappers/openrouter.ts b/js/src/wrappers/openrouter.ts index bfd39d2c1..ab7231b0c 100644 --- a/js/src/wrappers/openrouter.ts +++ b/js/src/wrappers/openrouter.ts @@ -1,4 +1,4 @@ -import { openRouterChannels } from "../instrumentation/plugins/openrouter-channels"; +import { openRouterChannels } from "../instrumentation/providers/openrouter-channels"; import type { OpenRouterBeta, OpenRouterCallModelRequest, diff --git a/js/src/wrappers/pi-coding-agent.ts b/js/src/wrappers/pi-coding-agent.ts index 60565c596..f01f6a75a 100644 --- a/js/src/wrappers/pi-coding-agent.ts +++ b/js/src/wrappers/pi-coding-agent.ts @@ -1,4 +1,4 @@ -import { piCodingAgentChannels } from "../instrumentation/plugins/pi-coding-agent-channels"; +import { piCodingAgentChannels } from "../instrumentation/providers/pi-coding-agent-channels"; import type { PiAgentSession, PiAgentSessionClass, diff --git a/js/src/wrappers/strands-agent-sdk.ts b/js/src/wrappers/strands-agent-sdk.ts index 6c6173034..3b975c75c 100644 --- a/js/src/wrappers/strands-agent-sdk.ts +++ b/js/src/wrappers/strands-agent-sdk.ts @@ -1,4 +1,4 @@ -import { strandsAgentSDKChannels } from "../instrumentation/plugins/strands-agent-sdk-channels"; +import { strandsAgentSDKChannels } from "../instrumentation/providers/strands-agent-sdk-channels"; import type { StrandsAgent, StrandsAgentConstructor, diff --git a/js/src/wrappers/voyageai.ts b/js/src/wrappers/voyageai.ts index 6abc4bb1a..8285282f5 100644 --- a/js/src/wrappers/voyageai.ts +++ b/js/src/wrappers/voyageai.ts @@ -1,4 +1,4 @@ -import { voyageAIChannels } from "../instrumentation/plugins/voyageai-channels"; +import { voyageAIChannels } from "../instrumentation/providers/voyageai-channels"; import type { VoyageAIClient, VoyageAIContextualizedEmbedRequest,