From 1ddb7ffc698a1b343cfd474f91f1ca5a8c788b22 Mon Sep 17 00:00:00 2001 From: Omer Wilmowski Date: Thu, 16 Apr 2026 19:40:59 +0300 Subject: [PATCH] feat: support trace stitching via LANGFUSE_TRACE_ID env var MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When LANGFUSE_TRACE_ID is set to a valid 32-char hex trace ID, all OTEL spans created by the Vercel AI SDK will reuse that trace ID instead of generating a new one. This stitches OpenCode's tool and LLM spans under the caller's Langfuse trace. Uses a custom IdGenerator that extends RandomIdGenerator, overriding only generateTraceId() — span IDs remain random as usual. Co-Authored-By: Claude Opus 4.6 (1M context) --- package.json | 2 +- src/index.test.ts | 56 ++++++++++++++++++++++++++++++++++++++++++++--- src/index.ts | 43 +++++++++++++++++++++++++++++++++--- 3 files changed, 94 insertions(+), 7 deletions(-) diff --git a/package.json b/package.json index 82cede9..d6d3e16 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "opencode-plugin-langfuse", - "version": "0.1.8", + "version": "0.2.0", "description": "Langfuse observability plugin for OpenCode - LLM tracing, prompt versioning, and cost tracking", "main": "dist/index.js", "types": "dist/index.d.ts", diff --git a/src/index.test.ts b/src/index.test.ts index c08b77c..b7a048e 100644 --- a/src/index.test.ts +++ b/src/index.test.ts @@ -3,6 +3,9 @@ import { LangfusePlugin } from "./index"; const mockForceFlush = mock(() => Promise.resolve()); const mockStart = mock(() => {}); +const mockShutdown = mock(() => Promise.resolve()); + +let capturedNodeSDKOptions: Record = {}; mock.module("@langfuse/otel", () => ({ LangfuseSpanProcessor: mock(() => ({ @@ -11,9 +14,13 @@ mock.module("@langfuse/otel", () => ({ })); mock.module("@opentelemetry/sdk-node", () => ({ - NodeSDK: mock(() => ({ - start: mockStart, - })), + NodeSDK: mock((options: Record) => { + capturedNodeSDKOptions = options; + return { + start: mockStart, + shutdown: mockShutdown, + }; + }), })); const mockLog = mock(() => {}); @@ -40,7 +47,9 @@ describe("LangfusePlugin", () => { beforeEach(() => { mockForceFlush.mockClear(); mockStart.mockClear(); + mockShutdown.mockClear(); mockLog.mockClear(); + capturedNodeSDKOptions = {}; }); afterEach(() => { @@ -197,4 +206,45 @@ describe("LangfusePlugin", () => { }); }); }); + + describe("trace stitching", () => { + it("passes idGenerator to NodeSDK", async () => { + setupEnv(); + await LangfusePlugin(mockPluginInput()); + + expect(capturedNodeSDKOptions.idGenerator).toBeDefined(); + }); + + it("logs parent trace ID when LANGFUSE_TRACE_ID is set", async () => { + setupEnv({ + LANGFUSE_TRACE_ID: "abcdef1234567890abcdef1234567890", + }); + + await LangfusePlugin(mockPluginInput()); + + expect(mockLog).toHaveBeenCalledWith({ + body: { + service: "langfuse-otel", + level: "info", + message: + "OTEL tracing initialized → https://cloud.langfuse.com (stitching to parent trace abcdef12…)", + }, + }); + }); + + it("does not log parent trace when LANGFUSE_TRACE_ID is not set", async () => { + setupEnv(); + delete process.env.LANGFUSE_TRACE_ID; + + await LangfusePlugin(mockPluginInput()); + + expect(mockLog).toHaveBeenCalledWith({ + body: { + service: "langfuse-otel", + level: "info", + message: "OTEL tracing initialized → https://cloud.langfuse.com", + }, + }); + }); + }); }); diff --git a/src/index.ts b/src/index.ts index e2cbb45..065dd0a 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,6 +1,34 @@ import { LangfuseSpanProcessor } from "@langfuse/otel"; import type { Plugin } from "@opencode-ai/plugin"; import { NodeSDK } from "@opentelemetry/sdk-node"; +import { RandomIdGenerator } from "@opentelemetry/sdk-trace-base"; + +/** + * Custom ID generator that reuses a parent trace ID from the environment. + * + * When LANGFUSE_TRACE_ID is set, all root spans created by the Vercel AI SDK + * will share that trace ID — which is the snakur-2 parent trace. This stitches + * OpenCode's tool/LLM spans under the caller's Langfuse trace instead of + * creating standalone traces. + * + * OTEL trace IDs are 32-char hex strings. Langfuse observation IDs (span IDs) + * are 16-char hex. The random generator handles span IDs as usual. + */ +class ParentAwareIdGenerator extends RandomIdGenerator { + private readonly parentTraceId: string | undefined; + + constructor() { + super(); + const raw = process.env.LANGFUSE_TRACE_ID; + if (raw && /^[0-9a-f]{32}$/i.test(raw)) { + this.parentTraceId = raw.toLowerCase(); + } + } + + override generateTraceId(): string { + return this.parentTraceId ?? super.generateTraceId(); + } +} export const LangfusePlugin: Plugin = async ({ client }) => { const publicKey = process.env.LANGFUSE_PUBLIC_KEY; @@ -29,12 +57,21 @@ export const LangfusePlugin: Plugin = async ({ client }) => { environment, }); + const idGenerator = new ParentAwareIdGenerator(); + const parentTraceId = process.env.LANGFUSE_TRACE_ID; + const sdk = new NodeSDK({ spanProcessors: [processor], + idGenerator, }); sdk.start(); - log("info", `OTEL tracing initialized → ${baseUrl}`); + + if (parentTraceId) { + log("info", `OTEL tracing initialized → ${baseUrl} (stitching to parent trace ${parentTraceId.slice(0, 8)}…)`); + } else { + log("info", `OTEL tracing initialized → ${baseUrl}`); + } return { config: async (config) => { @@ -48,10 +85,10 @@ export const LangfusePlugin: Plugin = async ({ client }) => { event: async ({ event }) => { if (event.type === "session.idle") { log("info", "Flushing OTEL spans before idle"); - await processor.forceFlush(); // Flushes the trace to Langfuse + await processor.forceFlush(); } - if (event.type === "server.instance.disposed") await sdk.shutdown(); // Flushes the trace to Langfuse + if (event.type === "server.instance.disposed") await sdk.shutdown(); }, }; };