diff --git a/e2e/config/pr-comment-scenarios.json b/e2e/config/pr-comment-scenarios.json index 94a29964f..8a84f7f32 100644 --- a/e2e/config/pr-comment-scenarios.json +++ b/e2e/config/pr-comment-scenarios.json @@ -52,6 +52,14 @@ "metadataScenario": "openrouter-instrumentation", "variants": [{ "variantKey": "openrouter-current", "label": "Current" }] }, + { + "scenarioDirName": "openrouter-agent-instrumentation", + "label": "OpenRouter Agent Instrumentation", + "metadataScenario": "openrouter-agent-instrumentation", + "variants": [ + { "variantKey": "openrouter-agent-current", "label": "Current" } + ] + }, { "scenarioDirName": "ai-sdk-instrumentation", "label": "AI SDK Instrumentation", diff --git a/e2e/scenarios/openrouter-agent-instrumentation/assertions.ts b/e2e/scenarios/openrouter-agent-instrumentation/assertions.ts new file mode 100644 index 000000000..f6ab2ba9e --- /dev/null +++ b/e2e/scenarios/openrouter-agent-instrumentation/assertions.ts @@ -0,0 +1,136 @@ +import { beforeAll, describe, expect, test } from "vitest"; +import type { CapturedLogEvent } from "../../helpers/mock-braintrust-server"; +import { withScenarioHarness } from "../../helpers/scenario-harness"; +import { findChildSpans, findLatestSpan } from "../../helpers/trace-selectors"; +import { CHAT_MODEL, ROOT_NAME, SCENARIO_NAME } from "./constants.mjs"; + +const CHAT_MODEL_NAME = CHAT_MODEL.split("/").at(-1) ?? CHAT_MODEL; +const OPENROUTER_MODEL_PROVIDER = "openai"; + +type RunOpenRouterAgentScenario = (harness: { + runNodeScenarioDir: (options: { + entry: string; + nodeArgs: string[]; + runContext?: { variantKey: string }; + scenarioDir: string; + timeoutMs: number; + }) => Promise; + runScenarioDir: (options: { + entry: string; + runContext?: { variantKey: string }; + scenarioDir: string; + timeoutMs: number; + }) => Promise; +}) => Promise; + +function findOpenRouterSpans( + events: CapturedLogEvent[], + parentId: string | undefined, + names: string[], +) { + for (const name of names) { + const spans = findChildSpans(events, name, parentId); + if (spans.length > 0) { + return spans; + } + } + + return []; +} + +function findOpenRouterSpan( + events: CapturedLogEvent[], + parentId: string | undefined, + names: string[], +) { + const spans = findOpenRouterSpans(events, parentId, names); + return spans.find((candidate) => candidate.output !== undefined) ?? spans[0]; +} + +export function defineOpenRouterAgentTraceAssertions(options: { + name: string; + runScenario: RunOpenRouterAgentScenario; + timeoutMs: number; +}): void { + const testConfig = { + timeout: options.timeoutMs, + }; + + describe(options.name, () => { + let events: CapturedLogEvent[] = []; + + beforeAll(async () => { + await withScenarioHarness(async (harness) => { + await options.runScenario(harness); + events = harness.events(); + }); + }, options.timeoutMs); + + test("captures the root trace for the scenario", testConfig, () => { + const root = findLatestSpan(events, ROOT_NAME); + + expect(root).toBeDefined(); + expect(root?.row.metadata).toMatchObject({ + scenario: SCENARIO_NAME, + }); + }); + + test( + "captures trace for client.callModel(request) and nested tool/turn spans", + testConfig, + () => { + const root = findLatestSpan(events, ROOT_NAME); + const operation = findLatestSpan( + events, + "openrouter-agent-call-model-operation", + ); + const span = findOpenRouterSpan(events, operation?.span.id, [ + "openrouter.callModel", + ]); + const nestedLlmSpans = findOpenRouterSpans(events, span?.span.id, [ + "openrouter.beta.responses.send", + ]); + const nestedToolSpan = findOpenRouterSpan(events, span?.span.id, [ + "lookup_weather", + "openrouter.tool", + ]); + + expect(operation).toBeDefined(); + expect(span).toBeDefined(); + expect(operation?.span.parentIds).toEqual([root?.span.id ?? ""]); + expect(span?.span.type).toBe("llm"); + expect(span?.row.metadata).toMatchObject({ + provider: OPENROUTER_MODEL_PROVIDER, + }); + expect(String(span?.row.metadata?.model)).toContain(CHAT_MODEL_NAME); + expect(span?.output).toBeDefined(); + + expect(nestedLlmSpans.length).toBeGreaterThanOrEqual(2); + for (const [index, nestedLlmSpan] of nestedLlmSpans.entries()) { + expect(nestedLlmSpan?.span.type).toBe("llm"); + expect(nestedLlmSpan?.row.metadata).toMatchObject({ + provider: OPENROUTER_MODEL_PROVIDER, + step: index + 1, + }); + expect(String(nestedLlmSpan?.row.metadata?.model)).toContain( + CHAT_MODEL_NAME, + ); + expect(nestedLlmSpan?.output).toBeDefined(); + } + + expect(nestedToolSpan).toBeDefined(); + expect(nestedToolSpan?.span.type).toBe("tool"); + expect(nestedToolSpan?.input).toMatchObject({ + city: "Vienna", + }); + expect(nestedToolSpan?.row.metadata).toMatchObject({ + provider: "openrouter", + tool_name: "lookup_weather", + }); + expect(nestedToolSpan?.output).toMatchObject({ + forecast: "Sunny in Vienna", + }); + }, + ); + }); +} diff --git a/e2e/scenarios/openrouter-agent-instrumentation/constants.mjs b/e2e/scenarios/openrouter-agent-instrumentation/constants.mjs new file mode 100644 index 000000000..2c362eb5d --- /dev/null +++ b/e2e/scenarios/openrouter-agent-instrumentation/constants.mjs @@ -0,0 +1,5 @@ +const CHAT_MODEL = "openai/gpt-4o-mini"; +const ROOT_NAME = "openrouter-agent-root"; +const SCENARIO_NAME = "openrouter-agent-instrumentation"; + +export { CHAT_MODEL, ROOT_NAME, SCENARIO_NAME }; diff --git a/e2e/scenarios/openrouter-agent-instrumentation/package.json b/e2e/scenarios/openrouter-agent-instrumentation/package.json new file mode 100644 index 000000000..86b45fd6c --- /dev/null +++ b/e2e/scenarios/openrouter-agent-instrumentation/package.json @@ -0,0 +1,15 @@ +{ + "name": "@braintrust/e2e-openrouter-agent-instrumentation", + "private": true, + "braintrustScenario": { + "canary": { + "dependencies": { + "@openrouter/agent": "latest" + } + } + }, + "dependencies": { + "@openrouter/agent": "0.1.2", + "zod": "4.1.11" + } +} diff --git a/e2e/scenarios/openrouter-agent-instrumentation/pnpm-lock.yaml b/e2e/scenarios/openrouter-agent-instrumentation/pnpm-lock.yaml new file mode 100644 index 000000000..6449e9207 --- /dev/null +++ b/e2e/scenarios/openrouter-agent-instrumentation/pnpm-lock.yaml @@ -0,0 +1,40 @@ +lockfileVersion: '9.0' + +settings: + autoInstallPeers: true + excludeLinksFromLockfile: false + +importers: + + .: + dependencies: + '@openrouter/agent': + specifier: 0.1.2 + version: 0.1.2 + zod: + specifier: 4.1.11 + version: 4.1.11 + +packages: + + '@openrouter/agent@0.1.2': + resolution: {integrity: sha512-nlbatensa2OpQTVJHKAEnzkfcO/aBOIz+ao9YZ/gh6R3cxCXBHejp0sH7ZC3afXMNKtxUikF+gLfxHyrIky2Ng==} + + '@openrouter/sdk@0.11.2': + resolution: {integrity: sha512-uu8zu7vd4hA2l4vUD1UiZuefxqaH2/ixFcUG8GIO9+qcEJkWX4AYAil7SpGmZOTgy8STLFTEk4M4MmyUW0YMLg==} + + zod@4.1.11: + resolution: {integrity: sha512-WPsqwxITS2tzx1bzhIKsEs19ABD5vmCVa4xBo2tq/SrV4RNZtfws1EnCWQXM6yh8bD08a1idvkB5MZSBiZsjwg==} + +snapshots: + + '@openrouter/agent@0.1.2': + dependencies: + '@openrouter/sdk': 0.11.2 + zod: 4.1.11 + + '@openrouter/sdk@0.11.2': + dependencies: + zod: 4.1.11 + + zod@4.1.11: {} diff --git a/e2e/scenarios/openrouter-agent-instrumentation/scenario.impl.mjs b/e2e/scenarios/openrouter-agent-instrumentation/scenario.impl.mjs new file mode 100644 index 000000000..ba48296a7 --- /dev/null +++ b/e2e/scenarios/openrouter-agent-instrumentation/scenario.impl.mjs @@ -0,0 +1,73 @@ +import { tool } from "@openrouter/agent"; +import { wrapOpenRouterAgent } from "braintrust"; +import { z } from "zod"; +import { + runOperation, + runTracedScenario, +} from "../../helpers/provider-runtime.mjs"; +import { CHAT_MODEL, ROOT_NAME, SCENARIO_NAME } from "./constants.mjs"; + +function createWeatherTool(toolFactory) { + return toolFactory({ + name: "lookup_weather", + description: "Look up the weather forecast for a city.", + inputSchema: z.object({ + city: z.string(), + }), + outputSchema: z.object({ + forecast: z.string(), + }), + execute: async ({ city }) => ({ + forecast: `Sunny in ${city}`, + }), + }); +} + +async function runOpenRouterAgentInstrumentationScenario( + OpenRouter, + { decorateClient } = {}, +) { + const baseClient = new OpenRouter({ + apiKey: process.env.OPENROUTER_API_KEY, + }); + const client = decorateClient ? decorateClient(baseClient) : baseClient; + const weatherTool = createWeatherTool(tool); + + await runTracedScenario({ + callback: async () => { + await runOperation( + "openrouter-agent-call-model-operation", + "call-model", + async () => { + const result = client.callModel({ + input: + "Use the lookup_weather tool for Vienna exactly once, then answer with only the forecast.", + maxOutputTokens: 16, + maxToolCalls: 1, + model: CHAT_MODEL, + temperature: 0, + toolChoice: "required", + tools: [weatherTool], + }); + + await result.getText(); + }, + ); + }, + metadata: { + scenario: SCENARIO_NAME, + }, + projectNameBase: "e2e-openrouter-agent-instrumentation", + rootName: ROOT_NAME, + }); +} + +export async function runWrappedOpenRouterAgentInstrumentation(OpenRouter) { + await runOpenRouterAgentInstrumentationScenario(OpenRouter, { + decorateClient: wrapOpenRouterAgent, + }); +} + +export async function runAutoOpenRouterAgentInstrumentation(OpenRouter) { + await runOpenRouterAgentInstrumentationScenario(OpenRouter); +} diff --git a/e2e/scenarios/openrouter-agent-instrumentation/scenario.mjs b/e2e/scenarios/openrouter-agent-instrumentation/scenario.mjs new file mode 100644 index 000000000..58a3f18b6 --- /dev/null +++ b/e2e/scenarios/openrouter-agent-instrumentation/scenario.mjs @@ -0,0 +1,5 @@ +import { OpenRouter } from "@openrouter/agent"; +import { runMain } from "../../helpers/provider-runtime.mjs"; +import { runAutoOpenRouterAgentInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => runAutoOpenRouterAgentInstrumentation(OpenRouter)); diff --git a/e2e/scenarios/openrouter-agent-instrumentation/scenario.test.ts b/e2e/scenarios/openrouter-agent-instrumentation/scenario.test.ts new file mode 100644 index 000000000..f4499c6c0 --- /dev/null +++ b/e2e/scenarios/openrouter-agent-instrumentation/scenario.test.ts @@ -0,0 +1,46 @@ +import { describe } from "vitest"; +import { + prepareScenarioDir, + readInstalledPackageVersion, + resolveScenarioDir, +} from "../../helpers/scenario-harness"; +import { defineOpenRouterAgentTraceAssertions } from "./assertions"; + +const scenarioDir = await prepareScenarioDir({ + scenarioDir: resolveScenarioDir(import.meta.url), +}); +const openrouterAgentVersion = await readInstalledPackageVersion( + scenarioDir, + "@openrouter/agent", +); +const OPENROUTER_AGENT_VARIANT_KEY = "openrouter-agent-current"; +const TIMEOUT_MS = 90_000; + +describe(`openrouter agent ${openrouterAgentVersion}`, () => { + defineOpenRouterAgentTraceAssertions({ + name: "wrapped instrumentation", + runScenario: async ({ runScenarioDir }) => { + await runScenarioDir({ + entry: "scenario.ts", + runContext: { variantKey: OPENROUTER_AGENT_VARIANT_KEY }, + scenarioDir, + timeoutMs: TIMEOUT_MS, + }); + }, + timeoutMs: TIMEOUT_MS, + }); + + defineOpenRouterAgentTraceAssertions({ + name: "auto-hook instrumentation", + runScenario: async ({ runNodeScenarioDir }) => { + await runNodeScenarioDir({ + entry: "scenario.mjs", + nodeArgs: ["--import", "braintrust/hook.mjs"], + runContext: { variantKey: OPENROUTER_AGENT_VARIANT_KEY }, + scenarioDir, + timeoutMs: TIMEOUT_MS, + }); + }, + timeoutMs: TIMEOUT_MS, + }); +}); diff --git a/e2e/scenarios/openrouter-agent-instrumentation/scenario.ts b/e2e/scenarios/openrouter-agent-instrumentation/scenario.ts new file mode 100644 index 000000000..e25b29a83 --- /dev/null +++ b/e2e/scenarios/openrouter-agent-instrumentation/scenario.ts @@ -0,0 +1,5 @@ +import { OpenRouter } from "@openrouter/agent"; +import { runMain } from "../../helpers/scenario-runtime"; +import { runWrappedOpenRouterAgentInstrumentation } from "./scenario.impl.mjs"; + +runMain(async () => runWrappedOpenRouterAgentInstrumentation(OpenRouter)); diff --git a/js/src/auto-instrumentations/bundler/plugin.ts b/js/src/auto-instrumentations/bundler/plugin.ts index dd2f1efe2..c63ebd652 100644 --- a/js/src/auto-instrumentations/bundler/plugin.ts +++ b/js/src/auto-instrumentations/bundler/plugin.ts @@ -25,6 +25,7 @@ import { anthropicConfigs } from "../configs/anthropic"; import { aiSDKConfigs } from "../configs/ai-sdk"; import { claudeAgentSDKConfigs } from "../configs/claude-agent-sdk"; import { googleGenAIConfigs } from "../configs/google-genai"; +import { openRouterAgentConfigs } from "../configs/openrouter-agent"; import { openRouterConfigs } from "../configs/openrouter"; import { mistralConfigs } from "../configs/mistral"; @@ -74,6 +75,7 @@ export const unplugin = createUnplugin((options = {}) => { ...claudeAgentSDKConfigs, ...googleGenAIConfigs, ...openRouterConfigs, + ...openRouterAgentConfigs, ...mistralConfigs, ...(options.instrumentations || []), ]; diff --git a/js/src/auto-instrumentations/bundler/webpack-loader.ts b/js/src/auto-instrumentations/bundler/webpack-loader.ts index 29f3fc7ac..1431a02cd 100644 --- a/js/src/auto-instrumentations/bundler/webpack-loader.ts +++ b/js/src/auto-instrumentations/bundler/webpack-loader.ts @@ -34,6 +34,7 @@ import { anthropicConfigs } from "../configs/anthropic"; import { aiSDKConfigs } from "../configs/ai-sdk"; import { claudeAgentSDKConfigs } from "../configs/claude-agent-sdk"; import { googleGenAIConfigs } from "../configs/google-genai"; +import { openRouterAgentConfigs } from "../configs/openrouter-agent"; import { openRouterConfigs } from "../configs/openrouter"; import { mistralConfigs } from "../configs/mistral"; import { type BundlerPluginOptions } from "./plugin"; @@ -69,6 +70,7 @@ function getMatcher(options: BundlerPluginOptions): InstrumentationMatcher { ...claudeAgentSDKConfigs, ...googleGenAIConfigs, ...openRouterConfigs, + ...openRouterAgentConfigs, ...mistralConfigs, ...(options.instrumentations ?? []), ]; diff --git a/js/src/auto-instrumentations/configs/openrouter-agent.test.ts b/js/src/auto-instrumentations/configs/openrouter-agent.test.ts new file mode 100644 index 000000000..7dca3f433 --- /dev/null +++ b/js/src/auto-instrumentations/configs/openrouter-agent.test.ts @@ -0,0 +1,20 @@ +import { describe, expect, it } from "vitest"; +import { openRouterAgentConfigs } from "./openrouter-agent"; +import { openRouterAgentChannels } from "../../instrumentation/plugins/openrouter-agent-channels"; + +describe("openRouterAgentConfigs", () => { + it("registers auto-instrumentation for @openrouter/agent callModel()", () => { + expect(openRouterAgentConfigs).toContainEqual({ + channelName: openRouterAgentChannels.callModel.channelName, + module: { + name: "@openrouter/agent", + versionRange: ">=0.1.2", + filePath: "esm/inner-loop/call-model.js", + }, + functionQuery: { + functionName: "callModel", + kind: "Sync", + }, + }); + }); +}); diff --git a/js/src/auto-instrumentations/configs/openrouter-agent.ts b/js/src/auto-instrumentations/configs/openrouter-agent.ts new file mode 100644 index 000000000..f27411cb0 --- /dev/null +++ b/js/src/auto-instrumentations/configs/openrouter-agent.ts @@ -0,0 +1,17 @@ +import type { InstrumentationConfig } from "@apm-js-collab/code-transformer"; +import { openRouterAgentChannels } from "../../instrumentation/plugins/openrouter-agent-channels"; + +export const openRouterAgentConfigs: InstrumentationConfig[] = [ + { + channelName: openRouterAgentChannels.callModel.channelName, + module: { + name: "@openrouter/agent", + versionRange: ">=0.1.2", + filePath: "esm/inner-loop/call-model.js", + }, + functionQuery: { + functionName: "callModel", + kind: "Sync", + }, + }, +]; diff --git a/js/src/auto-instrumentations/hook.mts b/js/src/auto-instrumentations/hook.mts index db202a0a6..ec2083687 100644 --- a/js/src/auto-instrumentations/hook.mts +++ b/js/src/auto-instrumentations/hook.mts @@ -19,6 +19,7 @@ import { anthropicConfigs } from "./configs/anthropic.js"; import { aiSDKConfigs } from "./configs/ai-sdk.js"; import { claudeAgentSDKConfigs } from "./configs/claude-agent-sdk.js"; import { googleGenAIConfigs } from "./configs/google-genai.js"; +import { openRouterAgentConfigs } from "./configs/openrouter-agent.js"; import { openRouterConfigs } from "./configs/openrouter.js"; import { mistralConfigs } from "./configs/mistral.js"; import { ModulePatch } from "./loader/cjs-patch.js"; @@ -39,6 +40,7 @@ const allConfigs = [ ...claudeAgentSDKConfigs, ...googleGenAIConfigs, ...openRouterConfigs, + ...openRouterAgentConfigs, ...mistralConfigs, ]; diff --git a/js/src/auto-instrumentations/index.ts b/js/src/auto-instrumentations/index.ts index 87040081e..4177ae061 100644 --- a/js/src/auto-instrumentations/index.ts +++ b/js/src/auto-instrumentations/index.ts @@ -33,6 +33,7 @@ export { anthropicConfigs } from "./configs/anthropic"; export { aiSDKConfigs } from "./configs/ai-sdk"; export { claudeAgentSDKConfigs } from "./configs/claude-agent-sdk"; export { googleGenAIConfigs } from "./configs/google-genai"; +export { openRouterAgentConfigs } from "./configs/openrouter-agent"; export { openRouterConfigs } from "./configs/openrouter"; export { mistralConfigs } from "./configs/mistral"; diff --git a/js/src/exports.ts b/js/src/exports.ts index 9f072fc7a..d4a1a61c8 100644 --- a/js/src/exports.ts +++ b/js/src/exports.ts @@ -175,6 +175,7 @@ export { wrapAnthropic } from "./wrappers/anthropic"; export { wrapMastraAgent } from "./wrappers/mastra"; export { wrapClaudeAgentSDK } from "./wrappers/claude-agent-sdk/claude-agent-sdk"; export { wrapGoogleGenAI } from "./wrappers/google-genai"; +export { wrapOpenRouterAgent } from "./wrappers/openrouter-agent"; export { wrapOpenRouter } from "./wrappers/openrouter"; export { wrapMistral } from "./wrappers/mistral"; export { wrapVitest } from "./wrappers/vitest"; diff --git a/js/src/instrumentation/braintrust-plugin.test.ts b/js/src/instrumentation/braintrust-plugin.test.ts index d8f6a5904..07d91337d 100644 --- a/js/src/instrumentation/braintrust-plugin.test.ts +++ b/js/src/instrumentation/braintrust-plugin.test.ts @@ -5,6 +5,7 @@ import { AnthropicPlugin } from "./plugins/anthropic-plugin"; import { AISDKPlugin } from "./plugins/ai-sdk-plugin"; import { ClaudeAgentSDKPlugin } from "./plugins/claude-agent-sdk-plugin"; import { GoogleGenAIPlugin } from "./plugins/google-genai-plugin"; +import { OpenRouterAgentPlugin } from "./plugins/openrouter-agent-plugin"; import { OpenRouterPlugin } from "./plugins/openrouter-plugin"; import { MistralPlugin } from "./plugins/mistral-plugin"; @@ -49,6 +50,10 @@ vi.mock("./plugins/openrouter-plugin", () => ({ OpenRouterPlugin: createPluginClassMock(), })); +vi.mock("./plugins/openrouter-agent-plugin", () => ({ + OpenRouterAgentPlugin: createPluginClassMock(), +})); + vi.mock("./plugins/mistral-plugin", () => ({ MistralPlugin: createPluginClassMock(), })); @@ -110,10 +115,21 @@ describe("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(); @@ -133,6 +149,7 @@ describe("BraintrustPlugin", () => { expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); }); @@ -146,6 +163,7 @@ describe("BraintrustPlugin", () => { expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); }); @@ -159,6 +177,7 @@ describe("BraintrustPlugin", () => { expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); }); }); @@ -177,6 +196,7 @@ describe("BraintrustPlugin", () => { expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); }); it("should not create Anthropic plugin when anthropic: false", () => { @@ -192,6 +212,7 @@ describe("BraintrustPlugin", () => { expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); }); it("should not create AI SDK plugin when aisdk: false", () => { @@ -207,6 +228,7 @@ describe("BraintrustPlugin", () => { expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); }); @@ -223,6 +245,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); }); @@ -239,6 +262,7 @@ describe("BraintrustPlugin", () => { expect(AISDKPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); }); @@ -249,6 +273,7 @@ describe("BraintrustPlugin", () => { plugin.enable(); expect(OpenRouterPlugin).not.toHaveBeenCalled(); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(OpenAIPlugin).toHaveBeenCalledTimes(1); expect(AnthropicPlugin).toHaveBeenCalledTimes(1); expect(AISDKPlugin).toHaveBeenCalledTimes(1); @@ -272,6 +297,16 @@ describe("BraintrustPlugin", () => { expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); }); + 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: { @@ -281,6 +316,7 @@ describe("BraintrustPlugin", () => { claudeAgentSDK: false, googleGenAI: false, openrouter: false, + openrouterAgent: false, mistral: false, }, }); @@ -292,6 +328,7 @@ describe("BraintrustPlugin", () => { expect(ClaudeAgentSDKPlugin).not.toHaveBeenCalled(); expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); expect(OpenRouterPlugin).not.toHaveBeenCalled(); + expect(OpenRouterAgentPlugin).not.toHaveBeenCalled(); expect(MistralPlugin).not.toHaveBeenCalled(); }); @@ -312,6 +349,7 @@ describe("BraintrustPlugin", () => { expect(OpenAIPlugin).toHaveBeenCalledTimes(1); expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(AnthropicPlugin).not.toHaveBeenCalled(); expect(AISDKPlugin).not.toHaveBeenCalled(); expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); @@ -420,6 +458,8 @@ describe("BraintrustPlugin", () => { const googleGenAIMock = vi.mocked(GoogleGenAIPlugin).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; expect(openaiMock.enable).toHaveBeenCalledTimes(1); @@ -428,6 +468,7 @@ describe("BraintrustPlugin", () => { expect(claudeAgentSDKMock.enable).toHaveBeenCalledTimes(1); expect(googleGenAIMock.enable).toHaveBeenCalledTimes(1); expect(openRouterMock.enable).toHaveBeenCalledTimes(1); + expect(openRouterAgentMock.enable).toHaveBeenCalledTimes(1); expect(mistralMock.enable).toHaveBeenCalledTimes(1); }); @@ -443,6 +484,8 @@ describe("BraintrustPlugin", () => { const googleGenAIMock = vi.mocked(GoogleGenAIPlugin).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; plugin.disable(); @@ -453,6 +496,7 @@ describe("BraintrustPlugin", () => { expect(claudeAgentSDKMock.disable).toHaveBeenCalledTimes(1); expect(googleGenAIMock.disable).toHaveBeenCalledTimes(1); expect(openRouterMock.disable).toHaveBeenCalledTimes(1); + expect(openRouterAgentMock.disable).toHaveBeenCalledTimes(1); expect(mistralMock.disable).toHaveBeenCalledTimes(1); }); @@ -493,6 +537,7 @@ describe("BraintrustPlugin", () => { expect(ClaudeAgentSDKPlugin).not.toHaveBeenCalled(); expect(GoogleGenAIPlugin).not.toHaveBeenCalled(); expect(OpenRouterPlugin).not.toHaveBeenCalled(); + expect(OpenRouterAgentPlugin).not.toHaveBeenCalled(); expect(MistralPlugin).not.toHaveBeenCalled(); }); @@ -511,6 +556,7 @@ describe("BraintrustPlugin", () => { expect(ClaudeAgentSDKPlugin).toHaveBeenCalledTimes(1); expect(GoogleGenAIPlugin).toHaveBeenCalledTimes(1); expect(OpenRouterPlugin).toHaveBeenCalledTimes(1); + expect(OpenRouterAgentPlugin).toHaveBeenCalledTimes(1); expect(MistralPlugin).toHaveBeenCalledTimes(1); }); @@ -523,6 +569,7 @@ describe("BraintrustPlugin", () => { claudeAgentSDK: false, googleGenAI: true, openrouter: true, + openrouterAgent: true, mistral: false, }, }); @@ -533,6 +580,8 @@ describe("BraintrustPlugin", () => { const googleGenAIMock = vi.mocked(GoogleGenAIPlugin).mock.results[0].value; const openRouterMock = vi.mocked(OpenRouterPlugin).mock.results[0].value; + const openRouterAgentMock = vi.mocked(OpenRouterAgentPlugin).mock + .results[0].value; plugin.disable(); @@ -540,6 +589,9 @@ describe("BraintrustPlugin", () => { expect(aiSDKMock.disable).toHaveBeenCalledTimes(1); expect(googleGenAIMock.disable).toHaveBeenCalledTimes(1); expect(openRouterMock.disable).toHaveBeenCalledTimes(1); + expect(openRouterAgentMock.disable).toHaveBeenCalledTimes(1); + expect(MistralPlugin).not.toHaveBeenCalled(); + expect(openRouterAgentMock.disable).toHaveBeenCalledTimes(1); expect(MistralPlugin).not.toHaveBeenCalled(); }); }); diff --git a/js/src/instrumentation/braintrust-plugin.ts b/js/src/instrumentation/braintrust-plugin.ts index 726be1da2..cc6523de3 100644 --- a/js/src/instrumentation/braintrust-plugin.ts +++ b/js/src/instrumentation/braintrust-plugin.ts @@ -4,6 +4,7 @@ import { AnthropicPlugin } from "./plugins/anthropic-plugin"; import { AISDKPlugin } from "./plugins/ai-sdk-plugin"; import { ClaudeAgentSDKPlugin } from "./plugins/claude-agent-sdk-plugin"; import { GoogleGenAIPlugin } from "./plugins/google-genai-plugin"; +import { OpenRouterAgentPlugin } from "./plugins/openrouter-agent-plugin"; import { OpenRouterPlugin } from "./plugins/openrouter-plugin"; import { MistralPlugin } from "./plugins/mistral-plugin"; @@ -17,6 +18,7 @@ export interface BraintrustPluginConfig { googleGenAI?: boolean; claudeAgentSDK?: boolean; openrouter?: boolean; + openrouterAgent?: boolean; mistral?: boolean; }; } @@ -43,6 +45,7 @@ export class BraintrustPlugin extends BasePlugin { private claudeAgentSDKPlugin: ClaudeAgentSDKPlugin | null = null; private googleGenAIPlugin: GoogleGenAIPlugin | null = null; private openRouterPlugin: OpenRouterPlugin | null = null; + private openRouterAgentPlugin: OpenRouterAgentPlugin | null = null; private mistralPlugin: MistralPlugin | null = null; constructor(config: BraintrustPluginConfig = {}) { @@ -90,6 +93,11 @@ export class BraintrustPlugin extends BasePlugin { 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(); @@ -127,6 +135,11 @@ export class BraintrustPlugin extends BasePlugin { this.openRouterPlugin = null; } + if (this.openRouterAgentPlugin) { + this.openRouterAgentPlugin.disable(); + this.openRouterAgentPlugin = null; + } + if (this.mistralPlugin) { this.mistralPlugin.disable(); this.mistralPlugin = null; diff --git a/js/src/instrumentation/plugins/openrouter-agent-channels.ts b/js/src/instrumentation/plugins/openrouter-agent-channels.ts new file mode 100644 index 000000000..b95f40307 --- /dev/null +++ b/js/src/instrumentation/plugins/openrouter-agent-channels.ts @@ -0,0 +1,40 @@ +import { channel, defineChannels } from "../core/channel-definitions"; +import type { + OpenRouterAgentCallModelArgs, + OpenRouterAgentCallModelRequest, +} from "../../vendor-sdk-types/openrouter-agent"; + +export const openRouterAgentChannels = defineChannels("@openrouter/agent", { + callModel: channel({ + channelName: "callModel", + kind: "sync-stream", + }), + + callModelTurn: channel< + [OpenRouterAgentCallModelRequest | undefined], + unknown, + { + step: number; + stepType: "initial" | "continue"; + } + >({ + channelName: "callModel.turn", + kind: "async", + }), + + toolExecute: channel< + [unknown], + unknown | AsyncIterable, + { + span_info?: { + name?: string; + }; + toolCallId?: string; + toolName: string; + }, + unknown + >({ + channelName: "tool.execute", + kind: "async", + }), +}); diff --git a/js/src/instrumentation/plugins/openrouter-agent-plugin.test.ts b/js/src/instrumentation/plugins/openrouter-agent-plugin.test.ts new file mode 100644 index 000000000..cb0261670 --- /dev/null +++ b/js/src/instrumentation/plugins/openrouter-agent-plugin.test.ts @@ -0,0 +1,580 @@ +import { + afterEach, + beforeAll, + beforeEach, + describe, + expect, + it, + vi, +} from "vitest"; +import { configureNode } from "../../node/config"; +import { _exportsForTestingOnly, initLogger } from "../../logger"; +import { openRouterAgentChannels } from "./openrouter-agent-channels"; +import { + aggregateOpenRouterChatChunks, + aggregateOpenRouterResponseStreamEvents, + parseOpenRouterMetricsFromUsage, +} from "./openrouter-agent-plugin"; + +const TEST_PROVIDER = "openai"; +const TEST_MODEL = "gpt-4.1-mini"; + +try { + configureNode(); +} catch { + // Best-effort initialization for test environments. +} + +describe("OpenRouter Agent Plugin", () => { + let backgroundLogger: ReturnType< + typeof _exportsForTestingOnly.useTestBackgroundLogger + >; + + beforeAll(async () => { + await _exportsForTestingOnly.simulateLoginForTests(); + }); + + beforeEach(() => { + backgroundLogger = _exportsForTestingOnly.useTestBackgroundLogger(); + initLogger({ + projectName: "openrouter-agent-plugin.test.ts", + projectId: "test-project-id", + }); + }); + + afterEach(() => { + _exportsForTestingOnly.clearTestBackgroundLogger(); + vi.restoreAllMocks(); + }); + + describe("parseOpenRouterMetricsFromUsage", () => { + it("should parse chat token usage", () => { + expect( + parseOpenRouterMetricsFromUsage({ + promptTokens: 10, + completionTokens: 20, + totalTokens: 30, + promptTokensDetails: { + cachedTokens: 4, + cacheWriteTokens: 2, + audioTokens: 1, + }, + completionTokensDetails: { + reasoningTokens: 6, + acceptedPredictionTokens: 3, + }, + }), + ).toEqual({ + prompt_tokens: 10, + completion_tokens: 20, + tokens: 30, + prompt_cached_tokens: 4, + prompt_cache_write_tokens: 2, + prompt_audio_tokens: 1, + completion_reasoning_tokens: 6, + completion_accepted_prediction_tokens: 3, + }); + }); + + it("should parse responses usage with cost details", () => { + expect( + parseOpenRouterMetricsFromUsage({ + inputTokens: 11, + outputTokens: 7, + totalTokens: 18, + cost: 0.0021, + inputTokensDetails: { + cachedTokens: 5, + }, + outputTokensDetails: { + reasoningTokens: 2, + }, + costDetails: { + upstreamInferenceCost: 0.001, + upstreamInferenceInputCost: 0.0004, + upstreamInferenceOutputCost: 0.0006, + }, + isByok: true, + }), + ).toEqual({ + prompt_tokens: 11, + completion_tokens: 7, + tokens: 18, + cost: 0.0021, + prompt_cached_tokens: 5, + completion_reasoning_tokens: 2, + cost_upstream_inference_cost: 0.001, + cost_upstream_inference_input_cost: 0.0004, + cost_upstream_inference_output_cost: 0.0006, + }); + }); + + it("should ignore non-object usage", () => { + expect(parseOpenRouterMetricsFromUsage(undefined)).toEqual({}); + expect(parseOpenRouterMetricsFromUsage(null)).toEqual({}); + expect(parseOpenRouterMetricsFromUsage("nope")).toEqual({}); + }); + }); + + describe("aggregateOpenRouterResponseStreamEvents", () => { + it("should use the terminal response payload for output, metrics, and metadata", () => { + expect( + aggregateOpenRouterResponseStreamEvents([ + { type: "response.created" }, + { + type: "response.output_text.delta", + delta: "Hello", + }, + { + type: "response.completed", + response: { + id: "resp_123", + model: "openai/gpt-4.1-mini", + output: [{ type: "message", role: "assistant" }], + usage: { + inputTokens: 9, + outputTokens: 4, + totalTokens: 13, + cost: 0.0012, + isByok: true, + }, + }, + }, + ]), + ).toEqual({ + output: [{ type: "message", role: "assistant" }], + metrics: { + prompt_tokens: 9, + completion_tokens: 4, + tokens: 13, + cost: 0.0012, + }, + metadata: { + id: "resp_123", + model: TEST_MODEL, + provider: TEST_PROVIDER, + is_byok: true, + }, + }); + }); + + it("should ignore non-terminal events", () => { + expect( + aggregateOpenRouterResponseStreamEvents([ + { type: "response.created" }, + { type: "response.output_text.delta", delta: "Hello" }, + ]), + ).toEqual({ + output: undefined, + metrics: {}, + }); + }); + }); + + describe("aggregateOpenRouterChatChunks", () => { + it("should aggregate assistant content, tool calls, and usage", () => { + expect( + aggregateOpenRouterChatChunks([ + { + choices: [ + { + delta: { + role: "assistant", + content: "Hello", + tool_calls: [ + { + id: "call_1", + type: "function", + function: { + name: "lookup", + arguments: '{"cit', + }, + }, + ], + }, + }, + ], + }, + { + choices: [ + { + delta: { + content: " world", + tool_calls: [ + { + function: { + arguments: 'y":"Vienna"}', + }, + }, + ], + }, + finish_reason: "tool_calls", + }, + ], + usage: { + promptTokens: 12, + completionTokens: 8, + totalTokens: 20, + }, + }, + ]), + ).toEqual({ + output: [ + { + index: 0, + message: { + role: "assistant", + content: "Hello world", + tool_calls: [ + { + id: "call_1", + index: 0, + type: "function", + function: { + name: "lookup", + arguments: '{"city":"Vienna"}', + }, + }, + ], + }, + logprobs: null, + finish_reason: "tool_calls", + }, + ], + metrics: { + prompt_tokens: 12, + completion_tokens: 8, + tokens: 20, + }, + }); + }); + + it("should return an empty assistant message for empty chunks", () => { + expect(aggregateOpenRouterChatChunks([])).toEqual({ + output: [ + { + index: 0, + message: { + role: undefined, + content: undefined, + }, + logprobs: null, + finish_reason: undefined, + }, + ], + metrics: {}, + }); + }); + + it("should aggregate actual SDK camelCase chunk shapes", () => { + expect( + aggregateOpenRouterChatChunks([ + { + choices: [ + { + delta: { + role: "assistant", + content: "Let me check ", + toolCalls: [ + { + index: 0, + id: "call_1", + type: "function", + function: { + name: "lookup_weather", + arguments: '{"city":"Vie', + }, + }, + ], + }, + }, + ], + }, + { + choices: [ + { + delta: { + content: "that.", + toolCalls: [ + { + index: 0, + function: { + arguments: 'nna"}', + }, + }, + ], + }, + finishReason: "tool_calls", + }, + ], + }, + ]), + ).toEqual({ + output: [ + { + index: 0, + message: { + role: "assistant", + content: "Let me check that.", + tool_calls: [ + { + id: "call_1", + index: 0, + type: "function", + function: { + name: "lookup_weather", + arguments: '{"city":"Vienna"}', + }, + }, + ], + }, + logprobs: null, + finish_reason: "tool_calls", + }, + ], + metrics: {}, + }); + }); + }); + + describe("callModel tool patching", () => { + it("patches tools on callModel and records child llm turns plus the final response", async () => { + const tool = { + type: "function", + function: { + name: "lookup_weather", + execute: async (params: { city: string }) => ({ + forecast: `Sunny in ${params.city}`, + }), + }, + }; + const initialResponse = { + id: "resp_initial", + model: "openai/gpt-4.1-mini", + status: "completed", + output: [ + { + type: "function_call", + id: "fc_1", + callId: "call_1", + name: "lookup_weather", + arguments: '{"city":"Vienna"}', + }, + ], + usage: { + inputTokens: 10, + outputTokens: 4, + totalTokens: 14, + }, + }; + const finalResponse = { + id: "resp_final", + model: "openai/gpt-4.1-mini", + status: "completed", + output: [ + { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "Sunny in Vienna" }], + }, + ], + usage: { + inputTokens: 12, + outputTokens: 3, + totalTokens: 15, + }, + }; + const request = { + input: + "Use the lookup_weather tool for Vienna exactly once, then answer with only the forecast.", + maxOutputTokens: 32, + model: "openai/gpt-4.1-mini", + tools: [tool], + }; + + const result = openRouterAgentChannels.callModel.traceSync( + () => { + const modelResult = { + allToolExecutionRounds: [], + finalResponse, + resolvedRequest: { + input: + "Use the lookup_weather tool for Vienna exactly once, then answer with only the forecast.", + maxOutputTokens: 32, + model: "openai/gpt-4.1-mini", + tools: [{ type: "function", name: "lookup_weather" }], + }, + async getInitialResponse() { + return initialResponse; + }, + async makeFollowupRequest( + currentResponse: unknown, + toolResults: unknown[], + ) { + modelResult.allToolExecutionRounds.push({ + round: 0, + response: currentResponse, + toolResults, + }); + return finalResponse; + }, + async getResponse() { + return finalResponse; + }, + async getText() { + const currentResponse = await modelResult.getInitialResponse(); + const toolResult = await request.tools[0].function.execute( + { city: "Vienna" }, + { + toolCall: { + id: "call_1", + name: "lookup_weather", + }, + }, + ); + await modelResult.makeFollowupRequest(currentResponse, [ + { + type: "function_call_output", + id: "output_call_1", + callId: "call_1", + output: JSON.stringify(toolResult), + }, + ]); + return "Sunny in Vienna"; + }, + }; + + return modelResult; + }, + { + arguments: [request], + }, + ); + expect(request.tools[0]).not.toBe(tool); + + await expect(result.getText()).resolves.toBe("Sunny in Vienna"); + expect(request.tools[0]).not.toBe(tool); + + const spans = await backgroundLogger.drain(); + expect(spans).toHaveLength(4); + const callModelSpan = spans.find( + (span) => span.span_attributes?.name === "openrouter.callModel", + ) as Record | undefined; + const turnSpans = spans.filter( + (span) => + span.span_attributes?.name === "openrouter.beta.responses.send", + ) as Array>; + const toolSpan = spans.find( + (span) => span.span_attributes?.name === "lookup_weather", + ) as Record | undefined; + + expect(callModelSpan?.span_attributes).toMatchObject({ + name: "openrouter.callModel", + type: "llm", + }); + expect(callModelSpan?.metadata).toMatchObject({ + provider: TEST_PROVIDER, + model: TEST_MODEL, + maxOutputTokens: 32, + turn_count: 2, + }); + expect(callModelSpan?.output).toMatchObject(finalResponse.output); + expect(callModelSpan?.metrics).toMatchObject({ + prompt_tokens: 22, + completion_tokens: 7, + tokens: 29, + }); + + expect(turnSpans).toHaveLength(2); + expect(turnSpans[0]?.metadata).toMatchObject({ + provider: TEST_PROVIDER, + model: TEST_MODEL, + id: "resp_initial", + status: "completed", + step: 1, + step_type: "initial", + }); + expect(turnSpans[1]?.metadata).toMatchObject({ + provider: TEST_PROVIDER, + model: TEST_MODEL, + id: "resp_final", + status: "completed", + step: 2, + step_type: "continue", + }); + + expect(toolSpan?.span_attributes).toMatchObject({ + name: "lookup_weather", + type: "tool", + }); + expect(toolSpan?.input).toMatchObject({ + city: "Vienna", + }); + expect(toolSpan?.metadata).toMatchObject({ + provider: "openrouter", + tool_name: "lookup_weather", + tool_call_id: "call_1", + }); + expect(toolSpan?.output).toMatchObject({ + forecast: "Sunny in Vienna", + }); + }); + + it("supports @openrouter/agent OpenRouter.callModel(request)", async () => { + const finalResponse = { + id: "resp_agent", + model: "openai/gpt-4.1-mini", + status: "completed", + output: [ + { + type: "message", + role: "assistant", + content: [{ type: "output_text", text: "ok" }], + }, + ], + usage: { + inputTokens: 3, + outputTokens: 2, + totalTokens: 5, + }, + }; + const request = { + input: "reply with ok", + model: "openai/gpt-4.1-mini", + }; + + const result = openRouterAgentChannels.callModel.traceSync( + () => ({ + async getResponse() { + return finalResponse; + }, + async getText() { + return "ok"; + }, + }), + { + arguments: [request], + }, + ); + + await expect(result.getText()).resolves.toBe("ok"); + + const spans = await backgroundLogger.drain(); + const callModelSpan = spans.find( + (span) => span.span_attributes?.name === "openrouter.callModel", + ) as Record | undefined; + + expect(callModelSpan).toBeDefined(); + expect(callModelSpan?.metadata).toMatchObject({ + provider: TEST_PROVIDER, + model: TEST_MODEL, + }); + expect(callModelSpan?.output).toMatchObject(finalResponse.output); + expect(callModelSpan?.metrics).toMatchObject({ + prompt_tokens: 3, + completion_tokens: 2, + tokens: 5, + }); + }); + }); +}); diff --git a/js/src/instrumentation/plugins/openrouter-agent-plugin.ts b/js/src/instrumentation/plugins/openrouter-agent-plugin.ts new file mode 100644 index 000000000..7d3eb878b --- /dev/null +++ b/js/src/instrumentation/plugins/openrouter-agent-plugin.ts @@ -0,0 +1,1301 @@ +import { BasePlugin } from "../core"; +import { + traceAsyncChannel, + traceStreamingChannel, + traceSyncStreamChannel, + unsubscribeAll, +} from "../core/channel-tracing"; +import type { ChannelMessage } from "../core/channel-definitions"; +import { SpanTypeAttribute, isObject } from "../../../util/index"; +import { withCurrent } from "../../logger"; +import type { Span } from "../../logger"; +import { zodToJsonSchema } from "../../zod/utils"; +import { openRouterAgentChannels } from "./openrouter-agent-channels"; +import type { + OpenRouterAgentChatChoice, + OpenRouterAgentChatCompletionChunk, + OpenRouterAgentCallModelRequest, + OpenRouterAgentEmbeddingResponse, + OpenRouterAgentResponse, + OpenRouterAgentResponseStreamEvent, + OpenRouterAgentTool, + OpenRouterAgentToolTurnContext, +} from "../../vendor-sdk-types/openrouter-agent"; + +export class OpenRouterAgentPlugin extends BasePlugin { + protected onEnable(): 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.LLM, + 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 + ? 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 { + step: event?.step, + 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) : {}, + }), + ); + + 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: {}, + }), + }), + ); + + const callModelChannel = openRouterAgentChannels.callModel.tracingChannel(); + const callModelHandlers = { + start: (event: { arguments: unknown[] }) => { + const request = getOpenRouterCallModelRequestArg(event.arguments); + if (!request) { + return; + } + + patchOpenRouterCallModelRequestTools(request); + }, + }; + + callModelChannel.subscribe(callModelHandlers); + this.unsubscribers.push(() => { + callModelChannel.unsubscribe(callModelHandlers); + }); + } +} + +function normalizeArgs(args: unknown[] | unknown): unknown[] { + if (Array.isArray(args)) { + return args; + } + + if (isArrayLike(args)) { + return Array.from(args); + } + + return [args]; +} + +function isArrayLike(value: unknown): value is ArrayLike { + return ( + isObject(value) && + "length" in value && + typeof value.length === "number" && + Number.isInteger(value.length) && + value.length >= 0 + ); +} + +function getOpenRouterCallModelRequestArg( + args: unknown[] | unknown, +): OpenRouterAgentCallModelRequest | undefined { + const normalizedArgs = normalizeArgs(args); + + const keyedRequestArg = normalizedArgs.find( + (arg) => + isObject(arg) && ("input" in arg || "model" in arg || "tools" in arg), + ); + if (isObject(keyedRequestArg)) { + return keyedRequestArg as OpenRouterAgentCallModelRequest; + } + + const firstObjectArg = normalizedArgs.find((arg) => isObject(arg)); + return isObject(firstObjectArg) + ? (firstObjectArg as OpenRouterAgentCallModelRequest) + : undefined; +} + +const TOKEN_NAME_MAP: Record = { + promptTokens: "prompt_tokens", + inputTokens: "prompt_tokens", + completionTokens: "completion_tokens", + outputTokens: "completion_tokens", + totalTokens: "tokens", + prompt_tokens: "prompt_tokens", + input_tokens: "prompt_tokens", + completion_tokens: "completion_tokens", + output_tokens: "completion_tokens", + total_tokens: "tokens", +}; + +const TOKEN_DETAIL_PREFIX_MAP: Record = { + promptTokensDetails: "prompt", + inputTokensDetails: "prompt", + completionTokensDetails: "completion", + outputTokensDetails: "completion", + costDetails: "cost", + prompt_tokens_details: "prompt", + input_tokens_details: "prompt", + completion_tokens_details: "completion", + output_tokens_details: "completion", + cost_details: "cost", +}; + +function camelToSnake(value: string): string { + return value.replace(/[A-Z]/g, (match) => `_${match.toLowerCase()}`); +} + +function parseOpenRouterMetricsFromUsage( + usage: unknown, +): Record { + if (!isObject(usage)) { + return {}; + } + + const metrics: Record = {}; + + for (const [name, value] of Object.entries(usage)) { + if (typeof value === "number") { + metrics[TOKEN_NAME_MAP[name] || camelToSnake(name)] = value; + continue; + } + + if (!isObject(value)) { + continue; + } + + const prefix = TOKEN_DETAIL_PREFIX_MAP[name]; + if (!prefix) { + continue; + } + + for (const [nestedName, nestedValue] of Object.entries(value)) { + if (typeof nestedValue !== "number") { + continue; + } + + metrics[`${prefix}_${camelToSnake(nestedName)}`] = nestedValue; + } + } + + return metrics; +} + +function extractOpenRouterUsageMetadata( + usage: unknown, +): Record | undefined { + if (!isObject(usage)) { + return undefined; + } + + const metadata: Record = {}; + + if (typeof usage.isByok === "boolean") { + metadata.is_byok = usage.isByok; + } else if (typeof usage.is_byok === "boolean") { + metadata.is_byok = usage.is_byok; + } + + return Object.keys(metadata).length > 0 ? metadata : undefined; +} + +const OMITTED_OPENROUTER_KEYS = new Set([ + "execute", + "render", + "nextTurnParams", + "requireApproval", +]); + +function parseOpenRouterModelString(model: unknown): { + model: unknown; + provider?: string; +} { + if (typeof model !== "string") { + return { model }; + } + + const slashIndex = model.indexOf("/"); + if (slashIndex > 0 && slashIndex < model.length - 1) { + return { + provider: model.substring(0, slashIndex), + model: model.substring(slashIndex + 1), + }; + } + + return { model }; +} + +function isZodSchema(value: unknown): boolean { + return ( + value != null && + typeof value === "object" && + "_def" in value && + typeof (value as { _def?: unknown })._def === "object" + ); +} + +function serializeZodSchema(schema: unknown): Record { + try { + return zodToJsonSchema(schema as any) as Record; + } catch { + return { + type: "object", + description: "Zod schema (conversion failed)", + }; + } +} + +function serializeOpenRouterTool( + tool: OpenRouterAgentTool, +): OpenRouterAgentTool { + if (!isObject(tool)) { + return tool; + } + + const serialized: Record = {}; + for (const [key, value] of Object.entries(tool)) { + if (OMITTED_OPENROUTER_KEYS.has(key)) { + continue; + } + + if (key === "function" && isObject(value)) { + serialized.function = sanitizeOpenRouterLoggedValue(value); + continue; + } + + serialized[key] = sanitizeOpenRouterLoggedValue(value); + } + + return serialized as OpenRouterAgentTool; +} + +function serializeOpenRouterToolsForLogging( + tools: readonly OpenRouterAgentTool[] | undefined, +): OpenRouterAgentTool[] | undefined { + if (!Array.isArray(tools)) { + return undefined; + } + + return tools.map((tool) => serializeOpenRouterTool(tool)); +} + +function sanitizeOpenRouterLoggedValue(value: unknown): unknown { + if (isZodSchema(value)) { + return serializeZodSchema(value); + } + + if (typeof value === "function") { + return "[Function]"; + } + + if (Array.isArray(value)) { + return value.map((entry) => sanitizeOpenRouterLoggedValue(entry)); + } + + if (!isObject(value)) { + return value; + } + + const sanitized: Record = {}; + for (const [key, entry] of Object.entries(value)) { + if (OMITTED_OPENROUTER_KEYS.has(key)) { + continue; + } + + if (key === "tools" && Array.isArray(entry)) { + sanitized.tools = serializeOpenRouterToolsForLogging(entry); + continue; + } + + sanitized[key] = sanitizeOpenRouterLoggedValue(entry); + } + + return sanitized; +} + +function buildOpenRouterMetadata( + metadata: Record, + httpReferer: unknown, + appTitle: unknown, + appCategories: unknown, + xTitle: unknown, +): Record { + const sanitized = sanitizeOpenRouterLoggedValue(metadata); + const metadataRecord = isObject(sanitized) ? sanitized : {}; + const { model, provider: providerRouting, ...rest } = metadataRecord; + const normalizedModel = parseOpenRouterModelString(model); + + return { + ...rest, + ...(normalizedModel.model !== undefined + ? { model: normalizedModel.model } + : {}), + ...(providerRouting !== undefined ? { providerRouting } : {}), + ...(httpReferer !== undefined ? { httpReferer } : {}), + ...(appTitle !== undefined ? { appTitle } : {}), + ...(appCategories !== undefined ? { appCategories } : {}), + ...(xTitle !== undefined ? { xTitle } : {}), + provider: normalizedModel.provider || "openrouter", + }; +} + +function extractOpenRouterCallModelInput( + request: OpenRouterAgentCallModelRequest, +): unknown { + return isObject(request) && "input" in request + ? sanitizeOpenRouterLoggedValue(request.input) + : undefined; +} + +function extractOpenRouterCallModelMetadata( + request: OpenRouterAgentCallModelRequest, +): Record { + if (!isObject(request)) { + return { provider: "openrouter" }; + } + + const { input: _input, ...metadata } = request; + return buildOpenRouterMetadata( + metadata, + undefined, + undefined, + undefined, + undefined, + ); +} + +function extractOpenRouterResponseMetadata( + result: + | OpenRouterAgentResponse + | OpenRouterAgentEmbeddingResponse + | undefined, +): Record | undefined { + if (!isObject(result)) { + return undefined; + } + + const { output: _output, data: _data, usage, ...metadata } = result; + const sanitized = sanitizeOpenRouterLoggedValue(metadata); + const metadataRecord = isObject(sanitized) ? sanitized : {}; + const { model, provider, ...rest } = metadataRecord; + const normalizedModel = parseOpenRouterModelString(model); + const normalizedProvider = + (typeof provider === "string" ? provider : undefined) || + normalizedModel.provider; + const usageMetadata = extractOpenRouterUsageMetadata(usage); + const combined = { + ...rest, + ...(normalizedModel.model !== undefined + ? { model: normalizedModel.model } + : {}), + ...(usageMetadata || {}), + ...(normalizedProvider !== undefined + ? { provider: normalizedProvider } + : {}), + }; + + return Object.keys(combined).length > 0 ? combined : undefined; +} + +function extractOpenRouterResponseOutput( + response: Record | undefined, + fallbackOutput?: unknown, +): unknown { + if ( + isObject(response) && + "output" in response && + response.output !== undefined + ) { + return sanitizeOpenRouterLoggedValue(response.output); + } + + if (fallbackOutput !== undefined) { + return sanitizeOpenRouterLoggedValue(fallbackOutput); + } + + return undefined; +} + +const OPENROUTER_WRAPPED_TOOL = Symbol("braintrust.openrouter.wrappedTool"); + +type OpenRouterToolTraceContext = ChannelMessage< + typeof openRouterAgentChannels.toolExecute +>; + +type WrappedOpenRouterTool = OpenRouterAgentTool & { + [OPENROUTER_WRAPPED_TOOL]?: true; +}; + +function patchOpenRouterCallModelRequestTools( + request: OpenRouterAgentCallModelRequest, +): (() => void) | undefined { + if (!Array.isArray(request.tools) || request.tools.length === 0) { + return undefined; + } + + const originalTools = request.tools; + const wrappedTools = originalTools.map((tool) => wrapOpenRouterTool(tool)); + const didPatch = wrappedTools.some( + (tool, index) => tool !== originalTools[index], + ); + if (!didPatch) { + return undefined; + } + + (request as { tools?: readonly OpenRouterAgentTool[] }).tools = wrappedTools; + return () => { + (request as { tools?: readonly OpenRouterAgentTool[] }).tools = + originalTools; + }; +} + +function wrapOpenRouterTool(tool: OpenRouterAgentTool): OpenRouterAgentTool { + if ( + isWrappedTool(tool) || + !tool.function || + typeof tool.function !== "object" || + typeof tool.function.execute !== "function" + ) { + return tool; + } + + const toolName = tool.function.name || "tool"; + const originalExecute = tool.function.execute; + const wrappedTool: WrappedOpenRouterTool = { + ...tool, + function: { + ...tool.function, + execute(this: unknown, ...args: unknown[]) { + return traceToolExecution({ + args, + execute: () => Reflect.apply(originalExecute, this, args), + toolCallId: getToolCallId(args[1]), + toolName, + }); + }, + }, + }; + + Object.defineProperty(wrappedTool, OPENROUTER_WRAPPED_TOOL, { + value: true, + enumerable: false, + configurable: false, + }); + + return wrappedTool; +} + +function isWrappedTool(tool: OpenRouterAgentTool): boolean { + return Boolean((tool as WrappedOpenRouterTool)[OPENROUTER_WRAPPED_TOOL]); +} + +function traceToolExecution(args: { + args: unknown[]; + execute: () => unknown; + toolCallId?: string; + toolName: string; +}): unknown { + const tracingChannel = openRouterAgentChannels.toolExecute.tracingChannel(); + const input = args.args.length > 0 ? args.args[0] : undefined; + const event: OpenRouterToolTraceContext = { + arguments: [input], + span_info: { + name: args.toolName, + }, + toolCallId: args.toolCallId, + toolName: args.toolName, + }; + + tracingChannel.start!.publish(event); + + try { + const result = args.execute(); + return publishToolResult(tracingChannel, event, result); + } catch (error) { + event.error = normalizeError(error); + tracingChannel.error!.publish(event); + throw error; + } +} + +function publishToolResult( + tracingChannel: ReturnType< + typeof openRouterAgentChannels.toolExecute.tracingChannel + >, + event: OpenRouterToolTraceContext, + result: unknown, +): unknown { + if (isPromiseLike(result)) { + return result.then( + (resolved) => { + event.result = resolved; + tracingChannel.asyncEnd!.publish(event); + return resolved; + }, + (error) => { + event.error = normalizeError(error); + tracingChannel.error!.publish(event); + throw error; + }, + ); + } + + event.result = result; + tracingChannel.asyncEnd!.publish(event); + return result; +} + +function getToolCallId(context: unknown): string | undefined { + const toolContext = context as OpenRouterAgentToolTurnContext | undefined; + return typeof toolContext?.toolCall?.id === "string" + ? toolContext.toolCall.id + : undefined; +} + +function isPromiseLike(value: unknown): value is PromiseLike { + return ( + !!value && + (typeof value === "object" || typeof value === "function") && + "then" in value && + typeof value.then === "function" + ); +} + +export function aggregateOpenRouterChatChunks( + chunks: OpenRouterAgentChatCompletionChunk[], +): { + output: OpenRouterAgentChatChoice[]; + metrics: Record; +} { + let role: string | undefined; + let content = ""; + let toolCalls: + | Array<{ + index?: number; + id?: string; + type?: string; + function: { + name?: string; + arguments: string; + }; + }> + | undefined; + let finishReason: string | null | undefined; + let metrics: Record = {}; + + for (const chunk of chunks) { + metrics = { + ...metrics, + ...parseOpenRouterMetricsFromUsage(chunk?.usage), + }; + + const choice = chunk?.choices?.[0]; + const delta = choice?.delta; + if (!delta) { + if (choice?.finish_reason !== undefined) { + finishReason = choice.finish_reason; + } + continue; + } + + if (!role && delta.role) { + role = delta.role; + } + + if (typeof delta.content === "string") { + content += delta.content; + } + + const choiceFinishReason = + choice?.finishReason ?? choice?.finish_reason ?? undefined; + const deltaFinishReason = + delta.finishReason ?? delta.finish_reason ?? undefined; + + if (choiceFinishReason !== undefined) { + finishReason = choiceFinishReason; + } else if (deltaFinishReason !== undefined) { + finishReason = deltaFinishReason; + } + + const toolCallDeltas = Array.isArray(delta.toolCalls) + ? delta.toolCalls + : Array.isArray(delta.tool_calls) + ? delta.tool_calls + : undefined; + + if (!toolCallDeltas) { + continue; + } + + for (const toolDelta of toolCallDeltas) { + if (!toolDelta?.function) { + continue; + } + + const toolIndex = toolDelta.index ?? 0; + const existingToolCall = toolCalls?.[toolIndex]; + + if ( + !existingToolCall || + (toolDelta.id && + existingToolCall.id !== undefined && + existingToolCall.id !== toolDelta.id) + ) { + const nextToolCalls = [...(toolCalls || [])]; + nextToolCalls[toolIndex] = { + index: toolIndex, + id: toolDelta.id, + type: toolDelta.type, + function: { + name: toolDelta.function.name, + arguments: toolDelta.function.arguments || "", + }, + }; + toolCalls = nextToolCalls; + continue; + } + + const current = existingToolCall; + if (toolDelta.id && !current.id) { + current.id = toolDelta.id; + } + if (toolDelta.type && !current.type) { + current.type = toolDelta.type; + } + if (toolDelta.function.name && !current.function.name) { + current.function.name = toolDelta.function.name; + } + current.function.arguments += toolDelta.function.arguments || ""; + } + } + + return { + output: [ + { + index: 0, + message: { + role, + content: content || undefined, + ...(toolCalls ? { tool_calls: toolCalls } : {}), + }, + logprobs: null, + finish_reason: finishReason, + }, + ], + metrics, + }; +} + +export function aggregateOpenRouterResponseStreamEvents( + chunks: OpenRouterAgentResponseStreamEvent[], +): { + output: unknown; + metrics: Record; + metadata?: Record; +} { + let finalResponse: OpenRouterAgentResponse | undefined; + + for (const chunk of chunks) { + const response = chunk?.response; + if (!response) { + continue; + } + + if ( + chunk.type === "response.completed" || + chunk.type === "response.incomplete" || + chunk.type === "response.failed" + ) { + finalResponse = response; + } + } + + if (!finalResponse) { + return { + output: undefined, + metrics: {}, + }; + } + + return { + output: extractOpenRouterResponseOutput(finalResponse), + metrics: parseOpenRouterMetricsFromUsage(finalResponse.usage), + ...(extractOpenRouterResponseMetadata(finalResponse) + ? { metadata: extractOpenRouterResponseMetadata(finalResponse) } + : {}), + }; +} + +const OPENROUTER_WRAPPED_CALL_MODEL_RESULT = Symbol( + "braintrust.openrouter.wrappedCallModelResult", +); + +const OPENROUTER_CALL_MODEL_STREAM_METHODS = [ + "getFullResponsesStream", + "getItemsStream", + "getNewMessagesStream", + "getReasoningStream", + "getTextStream", + "getToolCallsStream", + "getToolStream", +] as const; + +const OPENROUTER_CALL_MODEL_CONTEXT_METHODS = [ + "cancel", + "getPendingToolCalls", + "getState", + "getToolCalls", + "requiresApproval", +] as const; + +type OpenRouterCallModelTurnTraceContext = ChannelMessage< + typeof openRouterAgentChannels.callModelTurn +>; + +type OpenRouterCallModelResultLike = { + [OPENROUTER_WRAPPED_CALL_MODEL_RESULT]?: true; + allToolExecutionRounds?: unknown; + finalResponse?: unknown; + getInitialResponse?: (...args: unknown[]) => Promise; + getResponse?: (...args: unknown[]) => Promise; + makeFollowupRequest?: (...args: unknown[]) => Promise; + resolvedRequest?: unknown; + getText?: (...args: unknown[]) => Promise; + [key: string]: unknown; +}; + +function patchOpenRouterCallModelResult(args: { + request?: OpenRouterAgentCallModelRequest; + result: unknown; + span: Span; +}): boolean { + const { request, result, span } = args; + if (!isObject(result) || isWrappedCallModelResult(result)) { + return false; + } + + const resultLike = result as OpenRouterCallModelResultLike; + const hasInstrumentableMethod = + typeof resultLike.getResponse === "function" || + typeof resultLike.getText === "function" || + OPENROUTER_CALL_MODEL_STREAM_METHODS.some( + (methodName) => typeof resultLike[methodName] === "function", + ); + + if (!hasInstrumentableMethod) { + return false; + } + + Object.defineProperty(resultLike, OPENROUTER_WRAPPED_CALL_MODEL_RESULT, { + value: true, + enumerable: false, + configurable: false, + }); + + const originalGetResponse = + typeof resultLike.getResponse === "function" + ? resultLike.getResponse.bind(resultLike) + : undefined; + const originalGetInitialResponse = + typeof resultLike.getInitialResponse === "function" + ? resultLike.getInitialResponse.bind(resultLike) + : undefined; + const originalMakeFollowupRequest = + typeof resultLike.makeFollowupRequest === "function" + ? resultLike.makeFollowupRequest.bind(resultLike) + : undefined; + + let ended = false; + let tracedTurnCount = 0; + + const endSpanWithResult = async ( + response?: unknown, + fallbackOutput?: unknown, + ) => { + if (ended) { + return; + } + ended = true; + + const finalResponse = getFinalOpenRouterCallModelResponse( + resultLike, + response, + ); + if (finalResponse) { + const rounds = getOpenRouterCallModelRounds(resultLike); + + const metadata = extractOpenRouterCallModelResultMetadata( + finalResponse, + rounds.length + 1, + ); + span.log({ + output: extractOpenRouterResponseOutput(finalResponse, fallbackOutput), + ...(metadata ? { metadata } : {}), + metrics: aggregateOpenRouterCallModelMetrics(rounds, finalResponse), + }); + span.end(); + return; + } + + if (fallbackOutput !== undefined) { + span.log({ + output: fallbackOutput, + }); + } + + span.end(); + }; + + const endSpanWithError = (error: unknown) => { + if (ended) { + return; + } + ended = true; + span.log({ + error: normalizeError(error).message, + }); + span.end(); + }; + + const finalizeFromResponse = async (fallbackOutput?: unknown) => { + if (!originalGetResponse) { + await endSpanWithResult(undefined, fallbackOutput); + return; + } + + try { + await endSpanWithResult(await originalGetResponse(), fallbackOutput); + } catch { + await endSpanWithResult(undefined, fallbackOutput); + } + }; + + if (originalGetResponse) { + resultLike.getResponse = async (...args: unknown[]) => { + return await withCurrent(span, async () => { + try { + const response = await originalGetResponse(...args); + await endSpanWithResult(response); + return response; + } catch (error) { + endSpanWithError(error); + throw error; + } + }); + }; + } + + if (typeof resultLike.getText === "function") { + const originalGetText = resultLike.getText.bind(resultLike); + resultLike.getText = async (...args: unknown[]) => { + return await withCurrent(span, async () => { + try { + const text = await originalGetText(...args); + await finalizeFromResponse(text); + return text; + } catch (error) { + endSpanWithError(error); + throw error; + } + }); + }; + } + + for (const methodName of OPENROUTER_CALL_MODEL_CONTEXT_METHODS) { + if (typeof resultLike[methodName] !== "function") { + continue; + } + + const originalMethod = resultLike[methodName] as ( + ...args: unknown[] + ) => Promise; + resultLike[methodName] = async (...args: unknown[]) => { + return await withCurrent(span, async () => { + return await originalMethod.apply(resultLike, args); + }); + }; + } + + for (const methodName of OPENROUTER_CALL_MODEL_STREAM_METHODS) { + if (typeof resultLike[methodName] !== "function") { + continue; + } + + const originalMethod = resultLike[methodName] as ( + ...args: unknown[] + ) => AsyncIterable; + resultLike[methodName] = (...args: unknown[]) => { + const stream = withCurrent(span, () => + originalMethod.apply(resultLike, args), + ); + if (!isAsyncIterable(stream)) { + return stream; + } + + return wrapAsyncIterableWithSpan({ + finalize: finalizeFromResponse, + iteratorFactory: () => stream[Symbol.asyncIterator](), + onError: endSpanWithError, + span, + }); + }; + } + + if (originalGetInitialResponse) { + let initialTurnTraced = false; + resultLike.getInitialResponse = async (...args: unknown[]) => { + if (initialTurnTraced) { + return await withCurrent(span, async () => { + return await originalGetInitialResponse(...args); + }); + } + + initialTurnTraced = true; + const step = tracedTurnCount + 1; + const stepType = tracedTurnCount === 0 ? "initial" : "continue"; + + const response = await traceOpenRouterCallModelTurn({ + fn: async () => { + const nextResponse = await originalGetInitialResponse(...args); + tracedTurnCount++; + return nextResponse; + }, + parentSpan: span, + request: getOpenRouterResolvedRequest(resultLike, request), + step, + stepType, + }); + + return response; + }; + } + + if (originalMakeFollowupRequest) { + resultLike.makeFollowupRequest = async (...args: unknown[]) => { + const currentResponse = args[0]; + const toolResults = Array.isArray(args[1]) ? args[1] : []; + const step = tracedTurnCount + 1; + + const response = await traceOpenRouterCallModelTurn({ + fn: async () => { + const nextResponse = await originalMakeFollowupRequest(...args); + tracedTurnCount++; + return nextResponse; + }, + parentSpan: span, + request: buildOpenRouterFollowupRequest( + getOpenRouterResolvedRequest(resultLike, request), + currentResponse, + toolResults, + ), + step, + stepType: "continue", + }); + + return response; + }; + } + + return true; +} + +async function traceOpenRouterCallModelTurn(args: { + fn: () => Promise; + parentSpan: Span; + request: OpenRouterAgentCallModelRequest | undefined; + step: number; + stepType: "initial" | "continue"; +}): Promise { + const context: OpenRouterCallModelTurnTraceContext = { + arguments: [args.request], + step: args.step, + stepType: args.stepType, + }; + + return await withCurrent(args.parentSpan, () => + openRouterAgentChannels.callModelTurn.tracePromise(args.fn, context), + ); +} + +function isWrappedCallModelResult(value: unknown): boolean { + return Boolean( + isObject(value) && + (value as OpenRouterCallModelResultLike)[ + OPENROUTER_WRAPPED_CALL_MODEL_RESULT + ], + ); +} + +function extractOpenRouterCallModelResultMetadata( + response: Record, + turnCount?: number, +): Record | undefined { + const combined = { + ...(extractOpenRouterResponseMetadata(response) || {}), + ...(turnCount !== undefined ? { turn_count: turnCount } : {}), + }; + + return Object.keys(combined).length > 0 ? combined : undefined; +} + +function getFinalOpenRouterCallModelResponse( + result: OpenRouterCallModelResultLike, + response: unknown, +): Record | undefined { + if (isObject(response)) { + return response; + } + + return isObject(result.finalResponse) + ? (result.finalResponse as Record) + : undefined; +} + +type OpenRouterCallModelRound = { + response?: Record; + round?: number; + toolResults?: unknown[]; +}; + +function getOpenRouterCallModelRounds( + result: OpenRouterCallModelResultLike, +): OpenRouterCallModelRound[] { + if (!Array.isArray(result.allToolExecutionRounds)) { + return []; + } + + return result.allToolExecutionRounds + .filter((round): round is Record => isObject(round)) + .map((round) => ({ + response: isObject(round.response) + ? (round.response as Record) + : undefined, + round: typeof round.round === "number" ? round.round : undefined, + toolResults: Array.isArray(round.toolResults) ? round.toolResults : [], + })) + .filter((round) => round.response !== undefined); +} + +function aggregateOpenRouterCallModelMetrics( + rounds: OpenRouterCallModelRound[], + finalResponse: Record, +): Record { + const metrics: Record = {}; + const responses = [ + ...rounds.map((round) => round.response).filter(isObject), + finalResponse, + ]; + + for (const response of responses) { + const responseMetrics = parseOpenRouterMetricsFromUsage(response.usage); + for (const [name, value] of Object.entries(responseMetrics)) { + metrics[name] = (metrics[name] || 0) + value; + } + } + + return metrics; +} + +function buildNextOpenRouterCallModelInput( + currentInput: unknown, + response: Record, + toolResults: unknown[], +): unknown { + const normalizedInput = Array.isArray(currentInput) + ? [...currentInput] + : currentInput === undefined + ? [] + : [currentInput]; + const responseOutput = Array.isArray(response.output) + ? response.output + : response.output === undefined + ? [] + : [response.output]; + + return [...normalizedInput, ...responseOutput, ...toolResults].map((entry) => + sanitizeOpenRouterLoggedValue(entry), + ); +} + +function getOpenRouterResolvedRequest( + result: OpenRouterCallModelResultLike, + request: OpenRouterAgentCallModelRequest | undefined, +): OpenRouterAgentCallModelRequest | undefined { + if (isObject(result.resolvedRequest)) { + return result.resolvedRequest as OpenRouterAgentCallModelRequest; + } + + return request; +} + +function buildOpenRouterFollowupRequest( + request: OpenRouterAgentCallModelRequest | undefined, + currentResponse: unknown, + toolResults: unknown[], +): OpenRouterAgentCallModelRequest | undefined { + if (!request) { + return undefined; + } + + return { + ...request, + input: buildNextOpenRouterCallModelInput( + extractOpenRouterCallModelInput(request), + isObject(currentResponse) + ? (currentResponse as Record) + : {}, + toolResults, + ), + stream: false, + }; +} + +function wrapAsyncIterableWithSpan(args: { + finalize: (fallbackOutput?: unknown) => Promise; + iteratorFactory: () => AsyncIterator; + onError: (error: unknown) => void; + span: Span; +}): AsyncIterable { + return { + [Symbol.asyncIterator]() { + const iterator = args.iteratorFactory(); + return { + next(value?: unknown) { + return withCurrent(args.span, () => + value === undefined + ? iterator.next() + : ( + iterator.next as ( + value: unknown, + ) => Promise> + )(value), + ).then( + async (result) => { + if (result.done) { + await args.finalize(); + } + return result; + }, + (error) => { + args.onError(error); + throw error; + }, + ); + }, + return(value?: unknown) { + if (typeof iterator.return !== "function") { + return args.finalize().then(() => ({ + done: true, + value, + })); + } + + return withCurrent(args.span, () => iterator.return!(value)).then( + async (result) => { + await args.finalize(); + return result; + }, + (error) => { + args.onError(error); + throw error; + }, + ); + }, + throw(error?: unknown) { + args.onError(error); + if (typeof iterator.throw !== "function") { + return Promise.reject(error); + } + return withCurrent(args.span, () => iterator.throw!(error)); + }, + [Symbol.asyncIterator]() { + return this; + }, + }; + }, + }; +} + +function isAsyncIterable(value: unknown): value is AsyncIterable { + return ( + !!value && + (typeof value === "object" || typeof value === "function") && + Symbol.asyncIterator in value && + typeof value[Symbol.asyncIterator] === "function" + ); +} + +function normalizeError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)); +} + +export { parseOpenRouterMetricsFromUsage }; diff --git a/js/src/instrumentation/registry.ts b/js/src/instrumentation/registry.ts index e3521bdc3..d506dcc0b 100644 --- a/js/src/instrumentation/registry.ts +++ b/js/src/instrumentation/registry.ts @@ -21,6 +21,7 @@ export interface InstrumentationConfig { google?: boolean; claudeAgentSDK?: boolean; openrouter?: boolean; + openrouterAgent?: boolean; mistral?: boolean; }; } @@ -108,6 +109,7 @@ class PluginRegistry { google: true, claudeAgentSDK: true, openrouter: true, + openrouterAgent: true, mistral: true, }; } diff --git a/js/src/vendor-sdk-types/openrouter-agent.ts b/js/src/vendor-sdk-types/openrouter-agent.ts new file mode 100644 index 000000000..289c853a5 --- /dev/null +++ b/js/src/vendor-sdk-types/openrouter-agent.ts @@ -0,0 +1,96 @@ +// @openrouter/agent types + +export type OpenRouterAgentChatToolCallDelta = { + index?: number; + id?: string; + type?: string; + function?: { + name?: string; + arguments?: string; + }; +}; + +export type OpenRouterAgentChatChoice = { + index?: number; + message?: { + role?: string; + content?: string | null; + tool_calls?: unknown; + }; + logprobs?: unknown; + finish_reason?: string | null; +}; + +export type OpenRouterAgentChatCompletionChunk = { + choices?: Array<{ + delta?: { + role?: string; + content?: string; + tool_calls?: OpenRouterAgentChatToolCallDelta[]; + toolCalls?: OpenRouterAgentChatToolCallDelta[]; + finish_reason?: string | null; + finishReason?: string | null; + }; + finish_reason?: string | null; + finishReason?: string | null; + }>; + usage?: unknown; + [key: string]: unknown; +}; + +export type OpenRouterAgentEmbeddingResponse = + | string + | { + data?: Array<{ + embedding?: number[] | string; + }>; + usage?: unknown; + [key: string]: unknown; + }; + +export type OpenRouterAgentResponse = { + output?: unknown; + usage?: unknown; + [key: string]: unknown; +}; + +export type OpenRouterAgentResponseStreamEvent = { + type?: string; + response?: OpenRouterAgentResponse; + [key: string]: unknown; +}; + +export type OpenRouterAgentToolTurnContext = { + toolCall?: { + id?: string; + name?: string; + [key: string]: unknown; + }; + [key: string]: unknown; +}; + +export type OpenRouterAgentTool = { + function?: { + name?: string; + execute?: (...args: unknown[]) => unknown; + [key: string]: unknown; + }; + [key: string]: unknown; +}; + +export type OpenRouterAgentCallModelRequest = { + input?: unknown; + model?: unknown; + tools?: readonly OpenRouterAgentTool[]; + [key: string]: unknown; +}; + +export type OpenRouterAgentCallModelArgs = [OpenRouterAgentCallModelRequest]; + +export type OpenRouterAgentClient = { + callModel?: ( + request: OpenRouterAgentCallModelRequest, + options?: unknown, + ) => unknown; + [key: string]: unknown; +}; diff --git a/js/src/vendor-sdk-types/openrouter.ts b/js/src/vendor-sdk-types/openrouter.ts index 5301e4701..aba3c9960 100644 --- a/js/src/vendor-sdk-types/openrouter.ts +++ b/js/src/vendor-sdk-types/openrouter.ts @@ -1,3 +1,5 @@ +// @openrouter/sdk types + export type OpenRouterChatCreateParams = { httpReferer?: string; xTitle?: string; diff --git a/js/src/wrappers/openrouter-agent.test.ts b/js/src/wrappers/openrouter-agent.test.ts new file mode 100644 index 000000000..02f01e040 --- /dev/null +++ b/js/src/wrappers/openrouter-agent.test.ts @@ -0,0 +1,49 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { wrapOpenRouterAgent } from "./openrouter-agent"; +import { openRouterAgentChannels } from "../instrumentation/plugins/openrouter-agent-channels"; + +describe("wrapOpenRouterAgent", () => { + afterEach(() => { + vi.restoreAllMocks(); + }); + + it("returns the original value and warns for unsupported inputs", () => { + const warnSpy = vi.spyOn(console, "warn").mockImplementation(() => {}); + const input = { notCallModel: true }; + + expect(wrapOpenRouterAgent(input as object)).toBe(input); + expect(warnSpy).toHaveBeenCalledWith( + "Unsupported OpenRouter Agent library. Not wrapping.", + ); + }); + + it("emits callModel tracing events and clones the request", () => { + const traceSpy = vi.spyOn(openRouterAgentChannels.callModel, "traceSync"); + const sdk = { + name: "agent-sdk", + callModel(request: Record, options?: unknown) { + return { + options, + request, + thisName: this.name, + }; + }, + }; + const wrapped = wrapOpenRouterAgent(sdk); + + const request = { input: "hello", model: "openai/gpt-4.1-mini" }; + const result = wrapped.callModel(request); + + expect(traceSpy).toHaveBeenCalledTimes(1); + const traceContext = traceSpy.mock.calls[0]?.[1] as { + arguments: unknown[]; + }; + expect(traceContext.arguments[0]).toMatchObject(request); + expect(traceContext.arguments[0]).not.toBe(request); + expect(result).toMatchObject({ + options: undefined, + request: traceContext.arguments[0], + thisName: "agent-sdk", + }); + }); +}); diff --git a/js/src/wrappers/openrouter-agent.ts b/js/src/wrappers/openrouter-agent.ts new file mode 100644 index 000000000..07328249b --- /dev/null +++ b/js/src/wrappers/openrouter-agent.ts @@ -0,0 +1,91 @@ +import { openRouterAgentChannels } from "../instrumentation/plugins/openrouter-agent-channels"; +import type { + OpenRouterAgentClient, + OpenRouterAgentCallModelRequest, +} from "../vendor-sdk-types/openrouter-agent"; + +/** + * Wrap an @openrouter/agent OpenRouter client so callModel() emits + * diagnostics-channel events consumed by the OpenRouter Agent plugin. + */ +export function wrapOpenRouterAgent(agent: T): T { + const candidate: unknown = agent; + if ( + candidate && + typeof candidate === "object" && + "callModel" in candidate && + typeof candidate.callModel === "function" + ) { + return openRouterAgentProxy(candidate as OpenRouterAgentClient) as T; + } + + // eslint-disable-next-line no-restricted-properties -- preserving intentional console usage. + console.warn("Unsupported OpenRouter Agent library. Not wrapping."); + return agent; +} + +function openRouterAgentProxy( + agent: OpenRouterAgentClient, +): OpenRouterAgentClient { + const cache = new Map(); + + return new Proxy(agent, { + get(target, prop, receiver) { + if (cache.has(prop)) { + return cache.get(prop); + } + + const value = Reflect.get(target, prop, receiver); + + if (prop === "callModel" && typeof value === "function") { + const wrapped = wrapCallModel( + value as NonNullable, + target, + ); + cache.set(prop, wrapped); + return wrapped; + } + + if (typeof value === "function") { + const bound = value.bind(target); + cache.set(prop, bound); + return bound; + } + + return value; + }, + }); +} + +function wrapCallModel( + callModelFn: NonNullable, + defaultThis?: unknown, +): NonNullable { + return new Proxy(callModelFn, { + apply(target, thisArg, argArray) { + const request = cloneCallModelRequest(argArray[0]); + const options = argArray[1] as Parameters< + NonNullable + >[1]; + const invocationTarget = + thisArg === undefined ? (defaultThis ?? thisArg) : thisArg; + + return openRouterAgentChannels.callModel.traceSync( + () => Reflect.apply(target, invocationTarget, [request, options]), + { + arguments: [request], + } as Parameters[1], + ); + }, + }); +} + +function cloneCallModelRequest( + request: unknown, +): OpenRouterAgentCallModelRequest { + if (!request || typeof request !== "object") { + return request as OpenRouterAgentCallModelRequest; + } + + return { ...(request as OpenRouterAgentCallModelRequest) }; +}