diff --git a/src/index.test.ts b/src/index.test.ts index b7a048e..ad81d61 100644 --- a/src/index.test.ts +++ b/src/index.test.ts @@ -2,6 +2,9 @@ import { describe, it, expect, beforeEach, afterEach, mock } from "bun:test"; import { LangfusePlugin } from "./index"; const mockForceFlush = mock(() => Promise.resolve()); +const mockProcessorShutdown = mock(() => Promise.resolve()); +const mockOnStart = mock(() => {}); +const mockOnEnd = mock(() => {}); const mockStart = mock(() => {}); const mockShutdown = mock(() => Promise.resolve()); @@ -9,7 +12,10 @@ let capturedNodeSDKOptions: Record = {}; mock.module("@langfuse/otel", () => ({ LangfuseSpanProcessor: mock(() => ({ + onStart: mockOnStart, + onEnd: mockOnEnd, forceFlush: mockForceFlush, + shutdown: mockProcessorShutdown, })), })); @@ -46,6 +52,9 @@ describe("LangfusePlugin", () => { beforeEach(() => { mockForceFlush.mockClear(); + mockProcessorShutdown.mockClear(); + mockOnStart.mockClear(); + mockOnEnd.mockClear(); mockStart.mockClear(); mockShutdown.mockClear(); mockLog.mockClear(); @@ -176,6 +185,46 @@ describe("LangfusePlugin", () => { }); }); + describe("payload scrubbing", () => { + it("removes AI SDK prompt and response payload attributes before export", async () => { + setupEnv(); + await LangfusePlugin(mockPluginInput()); + + const [processor] = capturedNodeSDKOptions.spanProcessors as any[]; + const span = { + name: "ai.streamText.doStream", + attributes: { + "ai.prompt": '{"prompt":"large"}', + "ai.prompt.messages": '[{"role":"user","content":"large"}]', + "ai.prompt.tools": '[{"name":"bash"}]', + "ai.prompt.toolChoice": '{"type":"auto"}', + "ai.response.text": "large response", + "ai.response.toolCalls": '[{"toolName":"bash"}]', + "ai.response.object": '{"result":"large"}', + "ai.usage.totalTokens": 42, + "gen_ai.usage.input_tokens": 21, + "gen_ai.response.model": "gemini-3-flash-preview", + }, + }; + + processor.onEnd(span); + + expect(span.attributes["ai.prompt"]).toBeUndefined(); + expect(span.attributes["ai.prompt.messages"]).toBeUndefined(); + expect(span.attributes["ai.prompt.tools"]).toBeUndefined(); + expect(span.attributes["ai.prompt.toolChoice"]).toBeUndefined(); + expect(span.attributes["ai.response.text"]).toBeUndefined(); + expect(span.attributes["ai.response.toolCalls"]).toBeUndefined(); + expect(span.attributes["ai.response.object"]).toBeUndefined(); + expect(span.attributes["ai.usage.totalTokens"]).toBe(42); + expect(span.attributes["gen_ai.usage.input_tokens"]).toBe(21); + expect(span.attributes["gen_ai.response.model"]).toBe( + "gemini-3-flash-preview" + ); + expect(mockOnEnd).toHaveBeenCalledWith(span); + }); + }); + describe("environment configuration", () => { it("uses default baseUrl when not provided", async () => { setupEnv(); diff --git a/src/index.ts b/src/index.ts index 5ad70a9..e508785 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,7 +1,6 @@ import { LangfuseSpanProcessor } from "@langfuse/otel"; import type { Plugin } from "@opencode-ai/plugin"; import { - context, trace, ROOT_CONTEXT, type Context, @@ -9,7 +8,12 @@ import { } from "@opentelemetry/api"; import { AsyncLocalStorageContextManager } from "@opentelemetry/context-async-hooks"; import { NodeSDK } from "@opentelemetry/sdk-node"; -import { RandomIdGenerator } from "@opentelemetry/sdk-trace-base"; +import { + RandomIdGenerator, + type ReadableSpan, + type Span, + type SpanProcessor, +} from "@opentelemetry/sdk-trace-base"; /** * Custom ID generator that reuses a parent trace ID from the environment. @@ -60,6 +64,42 @@ class ParentAwareContextManager extends AsyncLocalStorageContextManager { } } +const AI_SDK_PAYLOAD_ATTRIBUTES = [ + "ai.prompt", + "ai.prompt.messages", + "ai.prompt.tools", + "ai.prompt.toolChoice", + "ai.response.text", + "ai.response.toolCalls", + "ai.response.object", +]; + +class PayloadScrubbingSpanProcessor implements SpanProcessor { + constructor(private readonly delegate: SpanProcessor) {} + + onStart(span: Span, parentContext: Context): void { + this.delegate.onStart(span, parentContext); + } + + onEnd(span: ReadableSpan): void { + if (span.name.startsWith("ai.")) { + const attributes = span.attributes as Record; + for (const attribute of AI_SDK_PAYLOAD_ATTRIBUTES) { + delete attributes[attribute]; + } + } + this.delegate.onEnd(span); + } + + forceFlush(): Promise { + return this.delegate.forceFlush(); + } + + shutdown(): Promise { + return this.delegate.shutdown(); + } +} + export const LangfusePlugin: Plugin = async ({ client }) => { const publicKey = process.env.LANGFUSE_PUBLIC_KEY; const secretKey = process.env.LANGFUSE_SECRET_KEY; @@ -80,12 +120,13 @@ export const LangfusePlugin: Plugin = async ({ client }) => { return {}; } - const processor = new LangfuseSpanProcessor({ + const langfuseProcessor = new LangfuseSpanProcessor({ publicKey, secretKey, baseUrl, environment, }); + const processor = new PayloadScrubbingSpanProcessor(langfuseProcessor); const idGenerator = new ParentAwareIdGenerator(); const parentTraceId = process.env.LANGFUSE_TRACE_ID;