diff --git a/js/package.json b/js/package.json index 577942fd6..ae5eed913 100644 --- a/js/package.json +++ b/js/package.json @@ -176,7 +176,7 @@ }, "dependencies": { "@ai-sdk/provider": "^1.1.3", - "@apm-js-collab/code-transformer": "^0.8.2", + "@apm-js-collab/code-transformer": "^0.9.0", "@next/env": "^14.2.3", "@vercel/functions": "^1.0.2", "ajv": "^8.17.1", diff --git a/js/src/auto-instrumentations/configs/google-adk.test.ts b/js/src/auto-instrumentations/configs/google-adk.test.ts new file mode 100644 index 000000000..260781859 --- /dev/null +++ b/js/src/auto-instrumentations/configs/google-adk.test.ts @@ -0,0 +1,114 @@ +import { describe, it, expect } from "vitest"; +import { googleADKConfigs } from "./google-adk"; + +describe("Google ADK Instrumentation Configs", () => { + it("should have valid configs", () => { + expect(googleADKConfigs).toBeDefined(); + expect(Array.isArray(googleADKConfigs)).toBe(true); + expect(googleADKConfigs).toHaveLength(4); + }); + + it("should have runner.runAsync config", () => { + const config = googleADKConfigs.find( + (c) => c.channelName === "runner.runAsync", + ); + + expect(config).toBeDefined(); + expect(config?.module.name).toBe("@google/adk"); + expect(config?.module.versionRange).toBe(">=0.1.0"); + expect(config?.module.filePath).toBe("dist/esm/index.js"); + expect((config?.functionQuery as any).className).toBe("Runner"); + expect((config?.functionQuery as any).methodName).toBe("runAsync"); + expect((config?.functionQuery as any).kind).toBe("Async"); + expect((config?.functionQuery as any).isExportAlias).toBe(true); + }); + + it("should have agent.runAsync config", () => { + const config = googleADKConfigs.find( + (c) => c.channelName === "agent.runAsync", + ); + + expect(config).toBeDefined(); + expect(config?.module.name).toBe("@google/adk"); + expect(config?.module.versionRange).toBe(">=0.1.0"); + expect(config?.module.filePath).toBe("dist/esm/index.js"); + expect((config?.functionQuery as any).className).toBe("BaseAgent"); + expect((config?.functionQuery as any).methodName).toBe("runAsync"); + expect((config?.functionQuery as any).kind).toBe("Async"); + expect((config?.functionQuery as any).isExportAlias).toBe(true); + }); + + it("should have llm.callLlmAsync config", () => { + const config = googleADKConfigs.find( + (c) => c.channelName === "llm.callLlmAsync", + ); + + expect(config).toBeDefined(); + expect(config?.module.name).toBe("@google/adk"); + expect(config?.module.versionRange).toBe(">=0.1.0"); + expect(config?.module.filePath).toBe("dist/esm/index.js"); + expect((config?.functionQuery as any).className).toBe("LlmAgent"); + expect((config?.functionQuery as any).methodName).toBe("callLlmAsync"); + expect((config?.functionQuery as any).kind).toBe("Async"); + expect((config?.functionQuery as any).isExportAlias).toBe(true); + }); + + it("should have mcpTool.runAsync config", () => { + const config = googleADKConfigs.find( + (c) => c.channelName === "mcpTool.runAsync", + ); + + expect(config).toBeDefined(); + expect(config?.module.name).toBe("@google/adk"); + expect(config?.module.versionRange).toBe(">=0.1.0"); + expect(config?.module.filePath).toBe("dist/esm/index.js"); + expect((config?.functionQuery as any).className).toBe("MCPTool"); + expect((config?.functionQuery as any).methodName).toBe("runAsync"); + expect((config?.functionQuery as any).kind).toBe("Async"); + expect((config?.functionQuery as any).isExportAlias).toBe(true); + }); + + it("should NOT include braintrust: or orchestrion: prefix (code-transformer adds orchestrion:google-adk: prefix)", () => { + for (const config of googleADKConfigs) { + expect(config.channelName).not.toContain("braintrust:"); + expect(config.channelName).not.toContain("orchestrion:"); + } + }); + + it("should target @google/adk package for all configs", () => { + for (const config of googleADKConfigs) { + expect(config.module.name).toBe("@google/adk"); + } + }); + + it("should have valid version ranges", () => { + for (const config of googleADKConfigs) { + expect(config.module.versionRange).toMatch(/^>=\d+\.\d+\.\d+$/); + } + }); + + it("should have valid function kinds", () => { + const validKinds = ["Async", "Sync", "Callback"]; + for (const config of googleADKConfigs) { + expect(validKinds).toContain((config.functionQuery as any).kind); + } + }); + + it("should use Async kind for all methods", () => { + for (const config of googleADKConfigs) { + expect((config.functionQuery as any).kind).toBe("Async"); + } + }); + + it("should use isExportAlias for all configs", () => { + for (const config of googleADKConfigs) { + expect((config.functionQuery as any).isExportAlias).toBe(true); + } + }); + + it("should target dist/esm/index.js for all configs", () => { + for (const config of googleADKConfigs) { + expect(config.module.filePath).toBe("dist/esm/index.js"); + } + }); +}); diff --git a/js/src/auto-instrumentations/configs/google-adk.ts b/js/src/auto-instrumentations/configs/google-adk.ts new file mode 100644 index 000000000..08ae663dc --- /dev/null +++ b/js/src/auto-instrumentations/configs/google-adk.ts @@ -0,0 +1,78 @@ +import type { InstrumentationConfig } from "@apm-js-collab/code-transformer"; + +/** + * Instrumentation configurations for the Google ADK (Agent Development Kit). + * + * These configs define which functions to instrument and what channel + * to emit events on. They are used by orchestrion-js to perform AST + * transformation at build-time or load-time. + * + * NOTE: Channel names should NOT include the braintrust: prefix. The code-transformer + * will prepend "orchestrion:google-adk:" to these names, resulting in final channel names like: + * "orchestrion:google-adk:runner.runAsync" + */ +export const googleADKConfigs: InstrumentationConfig[] = [ + // Runner.runAsync - Top-level orchestration entry point + { + channelName: "runner.runAsync", + module: { + name: "@google/adk", + versionRange: ">=0.1.0", + filePath: "dist/esm/index.js", + }, + functionQuery: { + className: "Runner", + methodName: "runAsync", + kind: "Async", + isExportAlias: true, + }, + }, + + // BaseAgent.runAsync - Agent execution + { + channelName: "agent.runAsync", + module: { + name: "@google/adk", + versionRange: ">=0.1.0", + filePath: "dist/esm/index.js", + }, + functionQuery: { + className: "BaseAgent", + methodName: "runAsync", + kind: "Async", + isExportAlias: true, + }, + }, + + // LlmAgent.callLlmAsync - Actual LLM call + { + channelName: "llm.callLlmAsync", + module: { + name: "@google/adk", + versionRange: ">=0.1.0", + filePath: "dist/esm/index.js", + }, + functionQuery: { + className: "LlmAgent", + methodName: "callLlmAsync", + kind: "Async", + isExportAlias: true, + }, + }, + + // MCPTool.runAsync - MCP tool calls + { + channelName: "mcpTool.runAsync", + module: { + name: "@google/adk", + versionRange: ">=0.1.0", + filePath: "dist/esm/index.js", + }, + functionQuery: { + className: "MCPTool", + methodName: "runAsync", + kind: "Async", + isExportAlias: true, + }, + }, +]; diff --git a/js/src/auto-instrumentations/hook.mts b/js/src/auto-instrumentations/hook.mts index 5ffe5e1e4..266a40e1b 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 { googleADKConfigs } from "./configs/google-adk.js"; import { ModulePatch } from "./loader/cjs-patch.js"; // Combine all instrumentation configs @@ -28,6 +29,7 @@ const allConfigs = [ ...aiSDKConfigs, ...claudeAgentSDKConfigs, ...googleGenAIConfigs, + ...googleADKConfigs, ]; // 1. Register ESM loader for ESM modules diff --git a/js/src/auto-instrumentations/index.ts b/js/src/auto-instrumentations/index.ts index 20b78db77..a3e0aa813 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 { googleADKConfigs } from "./configs/google-adk"; // Re-export orchestrion configuration types // Note: ModuleMetadata and FunctionQuery are properties of InstrumentationConfig, diff --git a/js/src/instrumentation/braintrust-plugin.ts b/js/src/instrumentation/braintrust-plugin.ts index 5a8e5d32c..414f3ea7a 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 { GoogleADKPlugin } from "./plugins/google-adk-plugin"; export interface BraintrustPluginConfig { integrations?: { @@ -13,6 +14,7 @@ export interface BraintrustPluginConfig { aisdk?: boolean; google?: boolean; googleGenAI?: boolean; + googleADK?: boolean; claudeAgentSDK?: boolean; }; } @@ -26,6 +28,7 @@ export interface BraintrustPluginConfig { * - Claude Agent SDK (agent interactions) * - Vercel AI SDK (generateText, streamText, etc.) * - Google GenAI SDK + * - Google ADK (Agent Development Kit) * * The plugin is automatically enabled when the Braintrust library is loaded. * Individual integrations can be disabled via configuration. @@ -37,6 +40,7 @@ export class BraintrustPlugin extends BasePlugin { private aiSDKPlugin: AISDKPlugin | null = null; private claudeAgentSDKPlugin: ClaudeAgentSDKPlugin | null = null; private googleGenAIPlugin: GoogleGenAIPlugin | null = null; + private googleADKPlugin: GoogleADKPlugin | null = null; constructor(config: BraintrustPluginConfig = {}) { super(); @@ -77,6 +81,12 @@ export class BraintrustPlugin extends BasePlugin { this.googleGenAIPlugin = new GoogleGenAIPlugin(); this.googleGenAIPlugin.enable(); } + + // Enable Google ADK integration (default: true) + if (integrations.googleADK !== false) { + this.googleADKPlugin = new GoogleADKPlugin(); + this.googleADKPlugin.enable(); + } } protected onDisable(): void { @@ -104,6 +114,11 @@ export class BraintrustPlugin extends BasePlugin { this.googleGenAIPlugin.disable(); this.googleGenAIPlugin = null; } + + if (this.googleADKPlugin) { + this.googleADKPlugin.disable(); + this.googleADKPlugin = null; + } } } diff --git a/js/src/instrumentation/plugins/google-adk-plugin.test.ts b/js/src/instrumentation/plugins/google-adk-plugin.test.ts new file mode 100644 index 000000000..6191748d3 --- /dev/null +++ b/js/src/instrumentation/plugins/google-adk-plugin.test.ts @@ -0,0 +1,493 @@ +import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; +import { GoogleADKPlugin } from "./google-adk-plugin"; +import { tracingChannel } from "dc-browser"; + +// Mock dc-browser +vi.mock("dc-browser", () => ({ + tracingChannel: vi.fn(), +})); + +// Mock logger +vi.mock("../../logger", () => ({ + startSpan: vi.fn(() => ({ + log: vi.fn(), + end: vi.fn(), + export: vi.fn(() => Promise.resolve({})), + })), +})); + +// Mock utility modules +vi.mock("../../../util/index", () => ({ + SpanTypeAttribute: { + TASK: "task", + LLM: "llm", + TOOL: "tool", + }, + isObject: vi.fn((val: unknown) => val !== null && typeof val === "object"), +})); + +vi.mock("../../util", () => ({ + getCurrentUnixTimestamp: vi.fn(() => 1000), +})); + +vi.mock("../core", () => ({ + 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 stream; + }), +})); + +describe("GoogleADKPlugin", () => { + let plugin: GoogleADKPlugin; + let mockChannel: any; + let mockUnsubscribe: any; + + beforeEach(() => { + mockUnsubscribe = vi.fn(); + mockChannel = { + subscribe: vi.fn(), + unsubscribe: mockUnsubscribe, + }; + + (tracingChannel as any).mockReturnValue(mockChannel); + + plugin = new GoogleADKPlugin(); + }); + + afterEach(() => { + vi.clearAllMocks(); + }); + + describe("enable/disable lifecycle", () => { + it("should subscribe to all four channels when enabled", () => { + plugin.enable(); + + expect(tracingChannel).toHaveBeenCalledWith( + "orchestrion:google-adk:runner.runAsync", + ); + expect(tracingChannel).toHaveBeenCalledWith( + "orchestrion:google-adk:agent.runAsync", + ); + expect(tracingChannel).toHaveBeenCalledWith( + "orchestrion:google-adk:llm.callLlmAsync", + ); + expect(tracingChannel).toHaveBeenCalledWith( + "orchestrion:google-adk:mcpTool.runAsync", + ); + expect(mockChannel.subscribe).toHaveBeenCalledTimes(4); + }); + + it("should not subscribe multiple times if enabled twice", () => { + plugin.enable(); + const firstCallCount = mockChannel.subscribe.mock.calls.length; + + plugin.enable(); + const secondCallCount = mockChannel.subscribe.mock.calls.length; + + expect(firstCallCount).toBe(secondCallCount); + }); + + it("should unsubscribe from channels when disabled", () => { + plugin.enable(); + plugin.disable(); + + expect(mockUnsubscribe).toHaveBeenCalledTimes(4); + }); + + it("should clear unsubscribers array after disable", () => { + plugin.enable(); + plugin.disable(); + + expect((plugin as any).unsubscribers).toHaveLength(0); + }); + + it("should not crash when disabled without being enabled", () => { + expect(() => plugin.disable()).not.toThrow(); + }); + + it("should allow re-enabling after disable", () => { + plugin.enable(); + plugin.disable(); + plugin.enable(); + + expect(mockChannel.subscribe).toHaveBeenCalledTimes(8); // 4 + 4 + }); + }); + + describe("runner.runAsync channel subscription", () => { + let handlers: any; + + beforeEach(() => { + plugin.enable(); + // Find the handlers for the runner channel (first subscribe call) + handlers = mockChannel.subscribe.mock.calls[0][0]; + }); + + it("should have start, asyncEnd, and error handlers", () => { + expect(handlers).toHaveProperty("start"); + expect(handlers).toHaveProperty("asyncEnd"); + expect(handlers).toHaveProperty("error"); + }); + + it("should handle start event with runner params", () => { + const event = { + self: { appName: "my-app" }, + arguments: [ + { + userId: "user-1", + sessionId: "session-1", + newMessage: { parts: [{ text: "Hello" }] }, + }, + ], + }; + + expect(() => handlers.start(event)).not.toThrow(); + }); + + it("should handle start event with missing self", () => { + const event = { + self: null, + arguments: [{}], + }; + + expect(() => handlers.start(event)).not.toThrow(); + }); + + it("should handle asyncEnd without matching start", () => { + const event = { + arguments: [{}], + result: {}, + }; + + expect(() => handlers.asyncEnd(event)).not.toThrow(); + }); + + it("should handle error without matching start", () => { + const event = { + arguments: [{}], + error: new Error("Test error"), + }; + + expect(() => handlers.error(event)).not.toThrow(); + }); + + it("should handle error with matching start", () => { + const startEvent = { + self: { appName: "my-app" }, + arguments: [{ userId: "user-1" }], + }; + handlers.start(startEvent); + + const errorEvent = { + ...startEvent, + error: new Error("Runner error"), + }; + + expect(() => handlers.error(errorEvent)).not.toThrow(); + }); + }); + + describe("agent.runAsync channel subscription", () => { + let handlers: any; + + beforeEach(() => { + plugin.enable(); + // Second subscribe call is for agent channel + handlers = mockChannel.subscribe.mock.calls[1][0]; + }); + + it("should have start, asyncEnd, and error handlers", () => { + expect(handlers).toHaveProperty("start"); + expect(handlers).toHaveProperty("asyncEnd"); + expect(handlers).toHaveProperty("error"); + }); + + it("should handle start event with agent name", () => { + const event = { + self: { name: "weather-agent" }, + arguments: [{}], + }; + + expect(() => handlers.start(event)).not.toThrow(); + }); + + it("should handle start event with missing agent name", () => { + const event = { + self: {}, + arguments: [{}], + }; + + expect(() => handlers.start(event)).not.toThrow(); + }); + }); + + describe("llm.callLlmAsync channel subscription", () => { + let handlers: any; + + beforeEach(() => { + plugin.enable(); + // Third subscribe call is for LLM channel + handlers = mockChannel.subscribe.mock.calls[2][0]; + }); + + it("should have start, asyncEnd, and error handlers", () => { + expect(handlers).toHaveProperty("start"); + expect(handlers).toHaveProperty("asyncEnd"); + expect(handlers).toHaveProperty("error"); + }); + + it("should handle start event with llm request", () => { + const event = { + self: { + name: "llm-agent", + llm: { model: "gemini-2.0-flash" }, + }, + arguments: [ + {}, // invocationContext + { + // llmRequest + contents: [{ parts: [{ text: "Hello" }] }], + config: { temperature: 0.7 }, + }, + {}, // modelResponseEvent + ], + }; + + expect(() => handlers.start(event)).not.toThrow(); + }); + + it("should handle start event with model on self directly", () => { + const event = { + self: { + name: "llm-agent", + model: "gemini-2.0-flash", + }, + arguments: [{}, {}, {}], + }; + + expect(() => handlers.start(event)).not.toThrow(); + }); + + it("should handle non-streaming asyncEnd result", () => { + const startEvent = { + self: { name: "llm-agent", llm: { model: "gemini-2.0-flash" } }, + arguments: [{}, {}, {}], + }; + handlers.start(startEvent); + + const endEvent = { + ...startEvent, + result: { + content: { parts: [{ text: "Response" }] }, + usageMetadata: { + promptTokenCount: 10, + candidatesTokenCount: 20, + totalTokenCount: 30, + }, + }, + }; + + expect(() => handlers.asyncEnd(endEvent)).not.toThrow(); + }); + }); + + describe("mcpTool.runAsync channel subscription", () => { + let handlers: any; + + beforeEach(() => { + plugin.enable(); + // Fourth subscribe call is for MCP tool channel + handlers = mockChannel.subscribe.mock.calls[3][0]; + }); + + it("should have start, asyncEnd, and error handlers", () => { + expect(handlers).toHaveProperty("start"); + expect(handlers).toHaveProperty("asyncEnd"); + expect(handlers).toHaveProperty("error"); + }); + + it("should handle start event with tool name and args", () => { + const event = { + self: { name: "get_weather" }, + arguments: [{ args: { location: "NYC" } }], + }; + + expect(() => handlers.start(event)).not.toThrow(); + }); + + it("should handle asyncEnd with result", () => { + const startEvent = { + self: { name: "get_weather" }, + arguments: [{ args: { location: "NYC" } }], + }; + handlers.start(startEvent); + + const endEvent = { + ...startEvent, + result: { content: [{ type: "text", text: "Sunny, 72F" }] }, + }; + + expect(() => handlers.asyncEnd(endEvent)).not.toThrow(); + }); + + it("should handle error with tool call", () => { + const startEvent = { + self: { name: "get_weather" }, + arguments: [{ args: { location: "NYC" } }], + }; + handlers.start(startEvent); + + const errorEvent = { + ...startEvent, + error: new Error("Tool execution failed"), + }; + + expect(() => handlers.error(errorEvent)).not.toThrow(); + }); + }); +}); + +describe("Google ADK helper functions", () => { + describe("isFinalResponse logic", () => { + it("should identify final response events (no function calls/responses)", () => { + const event = { + content: { parts: [{ text: "Final answer" }] }, + actions: {}, + }; + + // No functionCall or functionResponse parts = final + const hasFunctionCalls = event.content.parts.some( + (p: any) => p.functionCall, + ); + const hasFunctionResponses = event.content.parts.some( + (p: any) => p.functionResponse, + ); + + expect(hasFunctionCalls).toBe(false); + expect(hasFunctionResponses).toBe(false); + }); + + it("should not identify events with function calls as final", () => { + const event = { + content: { + parts: [{ functionCall: { name: "get_weather", args: {} } }], + }, + actions: {}, + }; + + const hasFunctionCalls = event.content.parts.some( + (p: any) => p.functionCall, + ); + expect(hasFunctionCalls).toBe(true); + }); + + it("should identify events with skipSummarization as final", () => { + const event = { + content: { parts: [] }, + actions: { skipSummarization: true }, + }; + + expect(event.actions.skipSummarization).toBe(true); + }); + + it("should not identify partial events as final", () => { + const event = { + content: { parts: [{ text: "Partial..." }] }, + partial: true, + actions: {}, + }; + + expect(event.partial).toBe(true); + }); + }); + + describe("token metric extraction", () => { + it("should extract usage metadata correctly", () => { + const usageMetadata = { + promptTokenCount: 10, + candidatesTokenCount: 20, + totalTokenCount: 30, + }; + + expect(usageMetadata.promptTokenCount).toBe(10); + expect(usageMetadata.candidatesTokenCount).toBe(20); + expect(usageMetadata.totalTokenCount).toBe(30); + }); + + it("should handle cached content tokens", () => { + const usageMetadata = { + promptTokenCount: 100, + cachedContentTokenCount: 50, + }; + + expect(usageMetadata.cachedContentTokenCount).toBe(50); + }); + + it("should handle thoughts tokens", () => { + const usageMetadata = { + candidatesTokenCount: 80, + thoughtsTokenCount: 20, + }; + + expect(usageMetadata.thoughtsTokenCount).toBe(20); + }); + + it("should handle missing usage metadata", () => { + const response: any = {}; + expect(response.usageMetadata).toBeUndefined(); + }); + }); + + describe("edge cases", () => { + it("should handle events without content", () => { + const event = {}; + expect((event as any).content).toBeUndefined(); + }); + + it("should handle events with empty parts", () => { + const event = { + content: { parts: [] }, + }; + expect(event.content.parts).toHaveLength(0); + }); + + it("should handle MCPTool with missing args", () => { + const request = {}; + expect((request as any).args).toBeUndefined(); + }); + }); +}); diff --git a/js/src/instrumentation/plugins/google-adk-plugin.ts b/js/src/instrumentation/plugins/google-adk-plugin.ts new file mode 100644 index 000000000..dae5d4709 --- /dev/null +++ b/js/src/instrumentation/plugins/google-adk-plugin.ts @@ -0,0 +1,560 @@ +import { tracingChannel } from "dc-browser"; +import { BasePlugin, isAsyncIterable, patchStreamIfNeeded } from "../core"; +import type { StartEvent } from "../core"; +import { startSpan } from "../../logger"; +import type { Span } from "../../logger"; +import { SpanTypeAttribute } from "../../../util/index"; +import { getCurrentUnixTimestamp } from "../../util"; + +/** + * Auto-instrumentation plugin for the Google ADK (Agent Development Kit). + * + * This plugin subscribes to orchestrion channels for Google ADK methods + * and creates Braintrust spans to track: + * - Runner.runAsync (top-level invocation) + * - BaseAgent.runAsync (agent execution) + * - LlmAgent.callLlmAsync (LLM calls) + * - MCPTool.runAsync (MCP tool calls) + */ +export class GoogleADKPlugin extends BasePlugin { + protected unsubscribers: Array<() => void> = []; + + protected onEnable(): void { + this.subscribeToRunnerRunAsync(); + this.subscribeToAgentRunAsync(); + this.subscribeToLlmCallLlmAsync(); + this.subscribeToMCPToolRunAsync(); + } + + protected onDisable(): void { + for (const unsubscribe of this.unsubscribers) { + unsubscribe(); + } + this.unsubscribers = []; + } + + /** + * Runner.runAsync - Top-level orchestration. + * Creates a TASK span, patches the async generator stream, + * and outputs the last final response event. + */ + private subscribeToRunnerRunAsync(): void { + const channelName = "orchestrion:google-adk:runner.runAsync"; + const channel = tracingChannel(channelName); + + const spans = new WeakMap(); + + const handlers = { + start: (event: StartEvent) => { + const self = event.self as any; + const args = event.arguments as any[]; + const params = args[0] || {}; + + const appName = self?.appName || "unknown"; + + const span = startSpan({ + name: `invocation [${appName}]`, + spanAttributes: { + type: SpanTypeAttribute.TASK, + }, + }); + + const startTime = getCurrentUnixTimestamp(); + spans.set(event, { span, startTime }); + + try { + const input: any = {}; + if (params.newMessage) { + input.newMessage = params.newMessage; + } + + const metadata: any = { provider: "google-adk" }; + if (params.userId) { + metadata.userId = params.userId; + } + if (params.sessionId) { + metadata.sessionId = params.sessionId; + } + + span.log({ input, metadata }); + } catch (error) { + console.error(`Error extracting input for ${channelName}:`, error); + } + }, + + asyncEnd: (event: any) => { + const spanData = spans.get(event); + if (!spanData) { + return; + } + + const { span } = spanData; + + if (isAsyncIterable(event.result)) { + patchStreamIfNeeded(event.result, { + onComplete: (chunks: any[]) => { + try { + // Find the last final response event + let lastFinalResponse: any = null; + for (const chunk of chunks) { + if (isFinalResponse(chunk)) { + lastFinalResponse = chunk; + } + } + + span.log({ + output: + lastFinalResponse || + (chunks.length > 0 ? chunks[chunks.length - 1] : undefined), + }); + } catch (error) { + console.error( + `Error extracting output for ${channelName}:`, + error, + ); + } finally { + span.end(); + } + }, + onError: (error: Error) => { + span.log({ error: error.message }); + span.end(); + }, + }); + } else { + span.log({ output: event.result }); + 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); + this.unsubscribers.push(() => { + channel.unsubscribe(handlers); + }); + } + + /** + * BaseAgent.runAsync - Agent execution. + * Creates a TASK span, patches the async generator stream, + * and outputs the last event. + */ + private subscribeToAgentRunAsync(): void { + const channelName = "orchestrion:google-adk:agent.runAsync"; + const channel = tracingChannel(channelName); + + const spans = new WeakMap(); + + const handlers = { + start: (event: StartEvent) => { + const self = event.self as any; + const agentName = self?.name || "unknown"; + + const span = startSpan({ + name: `agent_run [${agentName}]`, + spanAttributes: { + type: SpanTypeAttribute.TASK, + }, + }); + + const startTime = getCurrentUnixTimestamp(); + spans.set(event, { span, startTime }); + + try { + const metadata: any = { provider: "google-adk" }; + span.log({ metadata }); + } catch (error) { + console.error(`Error extracting input for ${channelName}:`, error); + } + }, + + asyncEnd: (event: any) => { + const spanData = spans.get(event); + if (!spanData) { + return; + } + + const { span } = spanData; + + if (isAsyncIterable(event.result)) { + patchStreamIfNeeded(event.result, { + onComplete: (chunks: any[]) => { + try { + const lastEvent = + chunks.length > 0 ? chunks[chunks.length - 1] : undefined; + + span.log({ output: lastEvent }); + } catch (error) { + console.error( + `Error extracting output for ${channelName}:`, + error, + ); + } finally { + span.end(); + } + }, + onError: (error: Error) => { + span.log({ error: error.message }); + span.end(); + }, + }); + } else { + span.log({ output: event.result }); + 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); + this.unsubscribers.push(() => { + channel.unsubscribe(handlers); + }); + } + + /** + * LlmAgent.callLlmAsync - Actual LLM call. + * Creates an LLM span, patches the async generator stream, + * and extracts token metrics from the response's usageMetadata. + */ + private subscribeToLlmCallLlmAsync(): void { + const channelName = "orchestrion:google-adk:llm.callLlmAsync"; + const channel = tracingChannel(channelName); + + const spans = new WeakMap(); + + const handlers = { + start: (event: StartEvent) => { + const self = event.self as any; + const args = event.arguments as any[]; + const llmRequest = args[1]; + + const span = startSpan({ + name: "llm_call", + spanAttributes: { + type: SpanTypeAttribute.LLM, + }, + }); + + const startTime = getCurrentUnixTimestamp(); + spans.set(event, { span, startTime }); + + try { + const metadata: any = { provider: "google-adk" }; + + // Extract model name from the agent instance + const model = self?.llm?.model || self?.model; + if (model) { + metadata.model = model; + } + + span.log({ + input: llmRequest, + metadata, + }); + } catch (error) { + console.error(`Error extracting input for ${channelName}:`, error); + } + }, + + asyncEnd: (event: any) => { + const spanData = spans.get(event); + if (!spanData) { + return; + } + + const { span, startTime } = spanData; + + if (isAsyncIterable(event.result)) { + patchStreamIfNeeded(event.result, { + onComplete: (chunks: any[]) => { + try { + const lastEvent = + chunks.length > 0 ? chunks[chunks.length - 1] : undefined; + + const metrics = extractLlmMetrics(chunks, startTime); + + span.log({ + output: lastEvent, + metrics, + }); + } catch (error) { + console.error( + `Error extracting output for ${channelName}:`, + error, + ); + } finally { + span.end(); + } + }, + onError: (error: Error) => { + span.log({ error: error.message }); + span.end(); + }, + }); + } else { + try { + const metrics = extractGenerateContentMetrics( + event.result, + startTime, + ); + span.log({ + output: event.result, + metrics, + }); + } catch (error) { + 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); + this.unsubscribers.push(() => { + channel.unsubscribe(handlers); + }); + } + + /** + * MCPTool.runAsync - MCP tool calls. + * Creates a TOOL span for non-streaming tool invocations. + */ + private subscribeToMCPToolRunAsync(): void { + const channelName = "orchestrion:google-adk:mcpTool.runAsync"; + const channel = tracingChannel(channelName); + + const spans = new WeakMap(); + + const handlers = { + start: (event: StartEvent) => { + const self = event.self as any; + const args = event.arguments as any[]; + const request = args[0] || {}; + + const toolName = self?.name || "unknown"; + + const span = startSpan({ + name: `mcp_tool [${toolName}]`, + spanAttributes: { + type: SpanTypeAttribute.TOOL, + }, + }); + + const startTime = getCurrentUnixTimestamp(); + spans.set(event, { span, startTime }); + + try { + span.log({ + input: { + tool_name: toolName, + arguments: request.args, + }, + metadata: { provider: "google-adk" }, + }); + } catch (error) { + console.error(`Error extracting input for ${channelName}:`, error); + } + }, + + asyncEnd: (event: any) => { + const spanData = spans.get(event); + if (!spanData) { + return; + } + + const { span, startTime } = spanData; + + try { + const end = getCurrentUnixTimestamp(); + span.log({ + output: event.result, + metrics: { duration: end - startTime }, + }); + } catch (error) { + 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); + this.unsubscribers.push(() => { + channel.unsubscribe(handlers); + }); + } +} + +/** + * Check if an ADK event is a final response. + * Mirrors the logic from @google/adk's isFinalResponse. + */ +function isFinalResponse(event: any): boolean { + if (!event) { + return false; + } + + if (event.actions?.skipSummarization) { + return true; + } + + if (event.longRunningToolIds && event.longRunningToolIds.length > 0) { + return true; + } + + const functionCalls = getFunctionCalls(event); + const functionResponses = getFunctionResponses(event); + + return ( + functionCalls.length === 0 && + functionResponses.length === 0 && + !event.partial + ); +} + +function getFunctionCalls(event: any): any[] { + const funcCalls: any[] = []; + if (event.content?.parts) { + for (const part of event.content.parts) { + if (part.functionCall) { + funcCalls.push(part.functionCall); + } + } + } + return funcCalls; +} + +function getFunctionResponses(event: any): any[] { + const funcResponses: any[] = []; + if (event.content?.parts) { + for (const part of event.content.parts) { + if (part.functionResponse) { + funcResponses.push(part.functionResponse); + } + } + } + return funcResponses; +} + +/** + * Extract metrics from LLM response chunks (streamed from callLlmAsync). + * The response events from ADK use the same Gemini usageMetadata format. + */ +function extractLlmMetrics( + chunks: any[], + startTime: number, +): Record { + const end = getCurrentUnixTimestamp(); + const metrics: Record = { + duration: end - startTime, + }; + + // Find the last chunk with usageMetadata + for (const chunk of chunks) { + if (chunk?.usageMetadata) { + Object.assign(metrics, extractUsageMetrics(chunk.usageMetadata)); + } + } + + return metrics; +} + +/** + * Extract metrics from a non-streaming generateContent response. + * Reuses the same Gemini usageMetadata format. + */ +function extractGenerateContentMetrics( + response: any, + startTime?: number, +): Record { + const metrics: Record = {}; + + if (startTime) { + const end = getCurrentUnixTimestamp(); + metrics.duration = end - startTime; + } + + if (response?.usageMetadata) { + Object.assign(metrics, extractUsageMetrics(response.usageMetadata)); + } + + return metrics; +} + +/** + * Extract standard token metrics from Gemini usageMetadata. + */ +function extractUsageMetrics(usageMetadata: any): Record { + const metrics: Record = {}; + + if (usageMetadata.promptTokenCount !== undefined) { + metrics.prompt_tokens = usageMetadata.promptTokenCount; + } + if (usageMetadata.candidatesTokenCount !== undefined) { + metrics.completion_tokens = usageMetadata.candidatesTokenCount; + } + if (usageMetadata.totalTokenCount !== undefined) { + metrics.tokens = usageMetadata.totalTokenCount; + } + if (usageMetadata.cachedContentTokenCount !== undefined) { + metrics.prompt_cached_tokens = usageMetadata.cachedContentTokenCount; + } + if (usageMetadata.thoughtsTokenCount !== undefined) { + metrics.completion_reasoning_tokens = usageMetadata.thoughtsTokenCount; + } + + return metrics; +} diff --git a/js/src/wrappers/vitest/wrapper.ts b/js/src/wrappers/vitest/wrapper.ts index 5460e3de4..c1c4b2f5d 100644 --- a/js/src/wrappers/vitest/wrapper.ts +++ b/js/src/wrappers/vitest/wrapper.ts @@ -389,11 +389,11 @@ export function wrapDescribe( }; const wrappedDescribe = wrapBare(originalDescribe); - // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + wrappedDescribe.skip = wrapBare(originalDescribe.skip); - // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + wrappedDescribe.only = wrapBare(originalDescribe.only); - // eslint-disable-next-line @typescript-eslint/consistent-type-assertions + wrappedDescribe.concurrent = wrapBare(originalDescribe.concurrent); if (originalDescribe.todo) wrappedDescribe.todo = originalDescribe.todo; if (originalDescribe.each) diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index dc240469c..1734a672d 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -253,8 +253,8 @@ importers: specifier: ^1.1.3 version: 1.1.3 '@apm-js-collab/code-transformer': - specifier: ^0.8.2 - version: 0.8.2 + specifier: ^0.9.0 + version: 0.9.0 '@next/env': specifier: ^14.2.3 version: 14.2.3 @@ -650,8 +650,8 @@ packages: peerDependencies: openapi-types: '>=7' - '@apm-js-collab/code-transformer@0.8.2': - resolution: {integrity: sha512-YRjJjNq5KFSjDUoqu5pFUWrrsvGOxl6c3bu+uMFc9HNNptZ2rNU/TI2nLw4jnhQNtka972Ee2m3uqbvDQtPeCA==} + '@apm-js-collab/code-transformer@0.9.0': + resolution: {integrity: sha512-cfHtufVUBKJz6se/tQBJqizgoot5AOxhVy5B9Cuo493p6b7hXKBEIi6tcZEbr4XwN1VnV9JZOu52MMTCOCSBMw==} '@babel/code-frame@7.24.7': resolution: {integrity: sha512-BcYH1CVJBO9tvyIZ2jVeXgSIMvGZ2FDRvDdOIVQyuklNKSsx+eppDEBq/g47Ayw+RqNFE+URvOShmf+f/qwAlA==} @@ -5967,7 +5967,7 @@ snapshots: call-me-maybe: 1.0.2 openapi-types: 12.1.3 - '@apm-js-collab/code-transformer@0.8.2': {} + '@apm-js-collab/code-transformer@0.9.0': {} '@babel/code-frame@7.24.7': dependencies: