From 4429d4d23dbc030cec2512abbbcb63fa2c0a8780 Mon Sep 17 00:00:00 2001 From: Ben Bachem <10088265+bezbac@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:27:08 +0200 Subject: [PATCH 1/2] Trace ChatML conversations --- src/index.ts | 92 ++++++- src/langfuse.ts | 348 ++++++++++++++++++------- src/opencode.ts | 4 +- test/integration/plugin.test.ts | 444 ++++++++++++++++++++++++++++++-- 4 files changed, 755 insertions(+), 133 deletions(-) diff --git a/src/index.ts b/src/index.ts index e1deda1..c24c3e1 100644 --- a/src/index.ts +++ b/src/index.ts @@ -10,6 +10,7 @@ import { createLangfuseClient, type ActiveGenerationStep, type LangfuseClient, + type ToolDefinition, } from "./langfuse.js"; import { OpencodeClientService } from "./opencode.js"; import { log } from "./utils.js"; @@ -186,10 +187,40 @@ const eventHook = (event: OpencodeEvent, shutdown?: () => Promise) => langfuse.rememberAssistantPart(part); langfuse.traceReasoningPart(part); + if (part.type === "tool" && part.state.status === "running") { + langfuse.traceToolStart({ + sessionID: part.sessionID, + messageID: part.messageID, + callID: part.callID, + tool: part.tool, + args: part.state.input, + started: part.state.time.start, + }); + } + + if (part.type === "tool" && part.state.status === "completed") { + langfuse.traceToolEnd({ + sessionID: part.sessionID, + messageID: part.messageID, + callID: part.callID, + tool: part.tool, + args: part.state.input, + title: part.state.title, + output: part.state.output, + started: part.state.time.start, + completed: part.state.time.end, + }); + } + if (part.type === "tool" && part.state.status === "error") { langfuse.traceToolError({ + sessionID: part.sessionID, + messageID: part.messageID, callID: part.callID, + tool: part.tool, + args: part.state.input, error: part.state.error, + started: part.state.time.start, completed: part.state.time.end, }); } @@ -223,6 +254,9 @@ const eventHook = (event: OpencodeEvent, shutdown?: () => Promise) => langfuse.rememberToolCall({ callID: event.properties.callID, messageID: event.properties.assistantMessageID, + sessionID: event.properties.sessionID, + tool: event.properties.tool, + args: event.properties.input, }); } @@ -240,13 +274,12 @@ const eventHook = (event: OpencodeEvent, shutdown?: () => Promise) => } if (event.type === "session.next.reasoning.ended") { - langfuse.traceReasoning({ + langfuse.rememberReasoning({ reasoningID: event.properties.reasoningID, sessionID: event.properties.sessionID, timestamp: event.properties.timestamp, text: event.properties.text, messageID: event.properties.assistantMessageID, - source: "session.next.reasoning.ended", }); } @@ -378,6 +411,7 @@ const main = Effect.gen(function* () { langfuse.clearTraceState(); }); const shutdownOnce = createShutdownOnce(langfuse); + const toolDefinitions = new Map>(); const runHook = ( hookName: string, @@ -438,16 +472,64 @@ const main = Effect.gen(function* () { "chat.message": (input, output) => runHook( "chat.message", - Effect.try({ - try: () => + Effect.gen(function* () { + let tools: ToolDefinition[] | undefined; + + if (input.model) { + const enabledTools = output.message.tools; + const cacheKey = JSON.stringify([ + input.model.providerID, + input.model.modelID, + enabledTools, + ]); + let pendingTools = toolDefinitions.get(cacheKey); + + if (!pendingTools) { + pendingTools = opencode.tool + .list({ + query: { + provider: input.model.providerID, + model: input.model.modelID, + }, + }) + .then(({ data }) => + (data ?? []) + .filter((tool) => enabledTools?.[tool.id] !== false) + .map((tool) => ({ + name: tool.id, + description: tool.description, + ...(typeof tool.parameters === "object" && + tool.parameters !== null && + !Array.isArray(tool.parameters) + ? { + parameters: tool.parameters as Record< + string, + unknown + >, + } + : {}), + })), + ) + .catch(() => { + toolDefinitions.delete(cacheKey); + return []; + }); + toolDefinitions.set(cacheKey, pendingTools); + } + + tools = yield* Effect.promise(() => pendingTools); + } + + yield* Effect.sync(() => langfuse.traceUserMessage({ sessionID: input.sessionID, messageID: input.messageID, agent: input.agent, model: input.model, parts: output.parts, + tools, }), - catch: (error) => error, + ); }), ), diff --git a/src/langfuse.ts b/src/langfuse.ts index 916d910..50b3ee8 100644 --- a/src/langfuse.ts +++ b/src/langfuse.ts @@ -35,11 +35,12 @@ export class LangfuseClient { this.traceState.abortedSessions.clear(); this.traceState.tracedEventIds.clear(); this.traceState.tracedReasoningIds.clear(); - this.traceState.pendingReasoningPartsByMessageId.clear(); this.traceState.generationSpansByMessageId.clear(); this.traceState.activeGenerationStepsByMessageId.clear(); this.traceState.toolMessageIdsByCallId.clear(); this.traceState.generationParentSpans.clear(); + this.traceState.generationInputsBySession.clear(); + this.traceState.toolResultSourceMessageIdsBySession.clear(); this.traceState.turnObservationsByMessageId.clear(); this.traceState.latestTurnObservationsBySession.clear(); this.traceState.finalizedToolCallIds.clear(); @@ -166,14 +167,12 @@ export class LangfuseClient { this.withObservationParent(input.sessionID, startEvent); } - traceReasoning(input: { + rememberReasoning(input: { reasoningID: string; sessionID: string; timestamp: number; text: string; messageID?: string; - source: string; - parentSpan?: ApiSpan; }) { if (!input.text.trim()) { return; @@ -187,30 +186,22 @@ export class LangfuseClient { this.traceState.tracedReasoningIds.add(reasoningTraceKey); - const parentSpan = - input.parentSpan ?? - (input.messageID - ? this.traceState.generationSpansByMessageId.get(input.messageID) - : undefined); - - const generationParentSpan = - parentSpan ?? - this.traceState.activeGenerationSteps.get(input.sessionID)?.span ?? - this.traceState.generationParentSpans.get(input.sessionID); + if (!input.messageID) { + return; + } - this.traceEvent({ - id: `reasoning:${reasoningTraceKey}`, + const parts = + this.traceState.assistantParts.get(input.messageID) ?? + new Map(); + parts.set(input.reasoningID, { + id: input.reasoningID, sessionID: input.sessionID, - name: "opencode.generation.reasoning", - timestamp: input.timestamp, - output: { text: input.text }, - metadata: { - reasoningID: input.reasoningID, - messageID: input.messageID, - source: input.source, - }, - parentSpan: generationParentSpan, + messageID: input.messageID, + type: "reasoning", + text: input.text, + time: { start: input.timestamp, end: input.timestamp }, }); + this.traceState.assistantParts.set(input.messageID, parts); } traceReasoningPart(part: MessagePart) { @@ -220,31 +211,12 @@ export class LangfuseClient { return; } - const generationSpan = - this.traceState.generationSpansByMessageId.get(part.messageID) ?? - this.traceState.activeGenerationSteps.get(part.sessionID)?.span ?? - this.traceState.generationParentSpans.get(part.sessionID); - - if (!generationSpan) { - const pending = - this.traceState.pendingReasoningPartsByMessageId.get(part.messageID) ?? - new Map(); - pending.set(part.id, part); - this.traceState.pendingReasoningPartsByMessageId.set( - part.messageID, - pending, - ); - return; - } - - this.traceReasoning({ + this.rememberReasoning({ reasoningID: part.id, sessionID: part.sessionID, timestamp: completed, text: part.text, messageID: part.messageID, - source: "message.part.updated", - parentSpan: generationSpan, }); } @@ -357,12 +329,19 @@ export class LangfuseClient { return; } + const generationInput = this.consumeGenerationInput(input.sessionID); + this.withTurnParent(input.sessionID, undefined, () => { const span = this.traceState.tracer.startSpan("opencode.generation", { attributes: { "langfuse.observation.type": "generation", "session.id": input.sessionID, "langfuse.observation.model.name": input.model.id, + ...(generationInput + ? { + "langfuse.observation.input": JSON.stringify(generationInput), + } + : {}), "langfuse.observation.metadata": JSON.stringify({ agent: input.agent, providerID: input.model.providerID, @@ -399,6 +378,7 @@ export class LangfuseClient { agent?: string; model?: { providerID: string; modelID: string }; parts: MessagePart[]; + tools?: ToolDefinition[]; }) { if ( input.messageID && @@ -409,9 +389,9 @@ export class LangfuseClient { this.traceState.abortedSessions.delete(input.sessionID); - const formattedInput = { + const formattedMessage = { role: "user" as const, - parts: input.parts.map((part) => { + content: input.parts.map((part) => { if (part.type === "text") { return { type: part.type, text: part.text ?? "" }; } @@ -447,6 +427,17 @@ export class LangfuseClient { return { type: part.type }; }), }; + const generationInput = [ + { + ...formattedMessage, + ...(input.tools?.length ? { tools: input.tools } : {}), + }, + ]; + + this.traceState.generationInputsBySession.set( + input.sessionID, + generationInput, + ); if (input.messageID) { this.traceState.tracedMessageIds.add(input.messageID); @@ -468,7 +459,7 @@ export class LangfuseClient { "langfuse.observation.type": "span", "langfuse.internal.is_app_root": true, "session.id": input.sessionID, - "langfuse.observation.input": JSON.stringify(formattedInput), + "langfuse.observation.input": JSON.stringify([formattedMessage]), "langfuse.observation.metadata": JSON.stringify({ messageID: input.messageID, agent: input.agent, @@ -501,7 +492,7 @@ export class LangfuseClient { attributes: { "langfuse.observation.type": "event", "session.id": input.sessionID, - "langfuse.observation.input": JSON.stringify(formattedInput), + "langfuse.observation.input": JSON.stringify([formattedMessage]), "langfuse.observation.metadata": JSON.stringify({ messageID: input.messageID, agent: input.agent, @@ -535,8 +526,32 @@ export class LangfuseClient { } } - rememberToolCall(input: { callID: string; messageID: string }) { + rememberToolCall(input: { + callID: string; + messageID: string; + sessionID?: string; + tool?: string; + args?: Record; + }) { this.traceState.toolMessageIdsByCallId.set(input.callID, input.messageID); + + if (!input.sessionID || !input.tool || !input.args) { + return; + } + + const parts = + this.traceState.assistantParts.get(input.messageID) ?? + new Map(); + parts.set(`tool:${input.callID}`, { + id: `tool:${input.callID}`, + sessionID: input.sessionID, + messageID: input.messageID, + type: "tool", + callID: input.callID, + tool: input.tool, + state: { status: "pending", input: input.args, raw: "" }, + }); + this.traceState.assistantParts.set(input.messageID, parts); } traceGeneration(input: { @@ -569,8 +584,7 @@ export class LangfuseClient { this.traceState.tracedGenerationIds.add(input.messageID); - const text = this.getAssistantText(input.messageID); - const output = text ? { text } : undefined; + const output = this.getAssistantMessage(input.messageID); const turn = this.getTurnObservation(input.sessionID, input.parentID); if (input.mode !== "compaction") { @@ -629,8 +643,6 @@ export class LangfuseClient { input.messageID, step.span, ); - this.flushPendingReasoning(input.messageID, step.span); - step.span.end(new Date(input.completed)); this.traceState.activeGenerationStepsByMessageId.delete(input.messageID); @@ -645,12 +657,19 @@ export class LangfuseClient { return; } + const generationInput = this.consumeGenerationInput(input.sessionID); + this.withTurnParent(input.sessionID, input.parentID, () => { const span = this.traceState.tracer.startSpan("opencode.generation", { attributes: { "langfuse.observation.type": "generation", "session.id": input.sessionID, "langfuse.observation.model.name": input.modelID, + ...(generationInput + ? { + "langfuse.observation.input": JSON.stringify(generationInput), + } + : {}), "langfuse.observation.output": JSON.stringify(output), "langfuse.observation.usage_details": JSON.stringify({ input: input.tokens.input, @@ -679,36 +698,10 @@ export class LangfuseClient { this.traceState.generationParentSpans.set(input.sessionID, span); this.traceState.generationSpansByMessageId.set(input.messageID, span); - this.flushPendingReasoning(input.messageID, span); span.end(new Date(input.completed)); }); } - private flushPendingReasoning(messageID: string, parentSpan: ApiSpan) { - const pending = - this.traceState.pendingReasoningPartsByMessageId.get(messageID) ?? - new Map(); - this.traceState.pendingReasoningPartsByMessageId.delete(messageID); - - for (const part of pending.values()) { - const completed = getCompletedReasoningTimestamp(part); - - if (completed === undefined) { - continue; - } - - this.traceReasoning({ - reasoningID: part.id, - sessionID: part.sessionID, - timestamp: completed, - text: part.text, - messageID: part.messageID, - source: "message.part.updated", - parentSpan, - }); - } - } - traceFailedGenerationStep(input: { id: string; sessionID: string; @@ -844,9 +837,16 @@ export class LangfuseClient { callID: string; tool: string; args: unknown; + messageID?: string; + started?: number; }) { - this.traceState.activeToolObservations.get(input.callID)?.span.end(); - this.traceState.finalizedToolCallIds.delete(input.callID); + if ( + this.traceState.finalizedToolCallIds.has(input.callID) || + this.traceState.activeToolObservations.has(input.callID) + ) { + return; + } + this.ensureGenerationParent(input.sessionID); this.withObservationParent( @@ -862,6 +862,9 @@ export class LangfuseClient { tool: input.tool, }), }, + ...(input.started === undefined + ? {} + : { startTime: new Date(input.started) }), }); this.traceState.activeToolObservations.set(input.callID, { @@ -870,7 +873,8 @@ export class LangfuseClient { tool: input.tool, }); }, - this.traceState.toolMessageIdsByCallId.get(input.callID), + input.messageID ?? + this.traceState.toolMessageIdsByCallId.get(input.callID), ); } @@ -881,6 +885,9 @@ export class LangfuseClient { args: unknown; title: string; output: string; + messageID?: string; + started?: number; + completed?: number; }) { if (this.traceState.finalizedToolCallIds.has(input.callID)) { return; @@ -892,6 +899,8 @@ export class LangfuseClient { callID: input.callID, tool: input.tool, args: input.args, + messageID: input.messageID, + started: input.started, }); } @@ -913,18 +922,54 @@ export class LangfuseClient { }), ); - span.end(); + span.end( + input.completed === undefined ? undefined : new Date(input.completed), + ); + this.rememberToolResult({ + sessionID: input.sessionID, + callID: input.callID, + tool: input.tool, + content: input.output, + messageID: input.messageID, + }); this.traceState.activeToolObservations.delete(input.callID); this.traceState.finalizedToolCallIds.add(input.callID); this.traceState.toolMessageIdsByCallId.delete(input.callID); } - traceToolError(input: { callID: string; error: string; completed: number }) { + traceToolError(input: { + callID: string; + error: string; + completed: number; + sessionID?: string; + tool?: string; + args?: unknown; + messageID?: string; + started?: number; + }) { if (this.traceState.finalizedToolCallIds.has(input.callID)) { return; } - const span = this.traceState.activeToolObservations.get(input.callID)?.span; + if ( + !this.traceState.activeToolObservations.has(input.callID) && + input.sessionID && + input.tool + ) { + this.traceToolStart({ + sessionID: input.sessionID, + callID: input.callID, + tool: input.tool, + args: input.args, + messageID: input.messageID, + started: input.started, + }); + } + + const observation = this.traceState.activeToolObservations.get( + input.callID, + ); + const span = observation?.span; if (!span) { return; @@ -940,6 +985,15 @@ export class LangfuseClient { }); span.recordException({ message: input.error }); span.end(new Date(input.completed)); + if (observation) { + this.rememberToolResult({ + sessionID: observation.sessionID, + callID: input.callID, + tool: input.tool ?? observation.tool, + content: input.error, + messageID: input.messageID, + }); + } this.traceState.activeToolObservations.delete(input.callID); this.traceState.finalizedToolCallIds.add(input.callID); this.traceState.toolMessageIdsByCallId.delete(input.callID); @@ -958,10 +1012,16 @@ export class LangfuseClient { } this.withTurnParent(sessionID, undefined, () => { + const generationInput = this.consumeGenerationInput(sessionID); const span = this.traceState.tracer.startSpan("opencode.generation", { attributes: { "langfuse.observation.type": "generation", "session.id": sessionID, + ...(generationInput + ? { + "langfuse.observation.input": JSON.stringify(generationInput), + } + : {}), }, }); @@ -1011,16 +1071,94 @@ export class LangfuseClient { : fn(); } - private getAssistantText(messageID: string) { - return Array.from( + private getAssistantMessage(messageID: string) { + const parts = Array.from( this.traceState.assistantParts.get(messageID)?.values() ?? [], - ) + ); + const content = parts .filter( (part): part is Extract => part.type === "text" && Boolean(part.text), ) .map((part) => part.text) .join(""); + const thinking = parts + .filter( + (part): part is Extract => + part.type === "reasoning" && Boolean(part.text), + ) + .map((part) => ({ type: "thinking" as const, content: part.text })); + const toolCallsById = new Map( + parts + .filter( + (part): part is Extract => + part.type === "tool", + ) + .map((part) => [part.callID, part] as const), + ); + const toolCalls = Array.from(toolCallsById.values()).map((part) => ({ + id: part.callID, + name: part.tool, + arguments: JSON.stringify(part.state.input), + })); + + if (!content && thinking.length === 0 && toolCalls.length === 0) { + return undefined; + } + + return [ + { + role: "assistant" as const, + ...(content ? { content } : {}), + ...(thinking.length ? { thinking } : {}), + ...(toolCalls.length ? { tool_calls: toolCalls } : {}), + }, + ]; + } + + private consumeGenerationInput(sessionID: string) { + const input = this.traceState.generationInputsBySession.get(sessionID); + const sourceMessageID = + this.traceState.toolResultSourceMessageIdsBySession.get(sessionID); + this.traceState.generationInputsBySession.delete(sessionID); + this.traceState.toolResultSourceMessageIdsBySession.delete(sessionID); + + if (!sourceMessageID) { + return input; + } + + return [ + ...(this.getAssistantMessage(sourceMessageID) ?? []), + ...(input ?? []), + ]; + } + + private rememberToolResult(input: { + sessionID: string; + callID: string; + tool: string; + content: string; + messageID?: string; + }) { + const messageID = + input.messageID ?? + this.traceState.toolMessageIdsByCallId.get(input.callID); + const toolResults = + this.traceState.generationInputsBySession.get(input.sessionID) ?? []; + toolResults.push({ + role: "tool", + name: input.tool, + tool_call_id: input.callID, + content: input.content, + }); + this.traceState.generationInputsBySession.set(input.sessionID, toolResults); + + if (messageID) { + this.traceState.toolResultSourceMessageIdsBySession.set( + input.sessionID, + messageID, + ); + } } private getSessionErrorMessage(error: SessionErrorInfo) { @@ -1050,10 +1188,6 @@ export type LangfuseTraceState = { tracedGenerationIds: Set; tracedEventIds: Set; tracedReasoningIds: Set; - pendingReasoningPartsByMessageId: Map< - string, - Map - >; generationSpansByMessageId: Map; activeGenerationStepsByMessageId: Map; toolMessageIdsByCallId: Map; @@ -1064,6 +1198,8 @@ export type LangfuseTraceState = { finalizedToolCallIds: Set; activeGenerationSteps: Map; generationParentSpans: Map; + generationInputsBySession: Map; + toolResultSourceMessageIdsBySession: Map; }; export type MessagePart = Extract< @@ -1123,9 +1259,26 @@ export type SessionErrorInfo = NonNullable; export type UserMessageInput = { role: "user"; - parts: FormattedMessagePart[]; + content: FormattedMessagePart[]; + tools?: ToolDefinition[]; +}; + +export type ToolDefinition = { + name: string; + description?: string; + parameters?: Record; }; +type ChatMlMessage = + | UserMessageInput + | { + role: "tool"; + name: string; + tool_call_id: string; + content: string; + } + | NonNullable>[number]; + export type TurnObservation = { span: ApiSpan; sessionID: string; @@ -1150,6 +1303,7 @@ export type ActiveGenerationStep = { span: ApiSpan; started?: number; snapshot?: string; + input?: ChatMlMessage[]; }; export class LangfuseClientService extends EffectContext.Tag( @@ -1211,10 +1365,6 @@ export const createLangfuseClient = (input: { tracedGenerationIds: new Set(), tracedEventIds: new Set(), tracedReasoningIds: new Set(), - pendingReasoningPartsByMessageId: new Map< - string, - Map - >(), generationSpansByMessageId: new Map(), activeGenerationStepsByMessageId: new Map(), toolMessageIdsByCallId: new Map(), @@ -1225,6 +1375,8 @@ export const createLangfuseClient = (input: { finalizedToolCallIds: new Set(), activeGenerationSteps: new Map(), generationParentSpans: new Map(), + generationInputsBySession: new Map(), + toolResultSourceMessageIdsBySession: new Map(), }; const processor = new LangfuseSpanProcessor({ diff --git a/src/opencode.ts b/src/opencode.ts index 66c52de..3872522 100644 --- a/src/opencode.ts +++ b/src/opencode.ts @@ -1,9 +1,7 @@ import type { PluginInput } from "@opencode-ai/plugin"; import { Context as EffectContext } from "effect"; -export type OpencodeClient = { - app: PluginInput["client"]["app"]; -}; +export type OpencodeClient = Pick; export class OpencodeClientService extends EffectContext.Tag( "OpencodeClientService", diff --git a/test/integration/plugin.test.ts b/test/integration/plugin.test.ts index 133b5bf..27dbee1 100644 --- a/test/integration/plugin.test.ts +++ b/test/integration/plugin.test.ts @@ -153,6 +153,8 @@ let plugin: Plugin; let collectorStatus = 200; let hooksDisposed = false; let collectorBaseUrl: string; +let toolListCalls = 0; +let toolListShouldFail = false; const startedAt = 1_750_000_000_000; @@ -357,6 +359,38 @@ const createHooks = async (baseUrl: string) => { app: { log: () => Promise.resolve(), }, + tool: { + list: () => { + toolListCalls += 1; + + if (toolListShouldFail) { + return Promise.reject(new Error("Tool discovery unavailable")); + } + + return Promise.resolve({ + data: [ + { + id: "read", + description: "Read a file", + parameters: { + type: "object", + properties: { path: { type: "string" } }, + required: ["path"], + }, + }, + { + id: "webfetch", + description: "Fetch a URL", + parameters: { + type: "object", + properties: { url: { type: "string" } }, + required: ["url"], + }, + }, + ], + }); + }, + }, } as unknown as Parameters[0]["client"]; return plugin({ client } as Parameters[0]); @@ -429,6 +463,8 @@ beforeEach(async () => { collectorErrors.length = 0; collectorStatus = 200; hooksDisposed = false; + toolListCalls = 0; + toolListShouldFail = false; hooks = await createHooks(collectorBaseUrl); }); @@ -531,6 +567,17 @@ describe.sequential("built plugin", () => { "opencode.generation", ].sort(), ); + for (const span of spans.filter( + (span) => + span.name === "opencode.turn" || span.name === "opencode.message.user", + )) { + expect(getJsonAttribute(span, "langfuse.observation.input")).toEqual([ + { + role: "user", + content: expect.any(Array), + }, + ]); + } const firstGeneration = spans .filter((span) => span.name === "opencode.generation") @@ -557,6 +604,35 @@ describe.sequential("built plugin", () => { "langfuse.user.id": "test-user", "session.id": sessionID, }); + expect( + getJsonAttribute(firstGeneration, "langfuse.observation.input"), + ).toEqual([ + { + role: "user", + content: [{ type: "text", text: "Inspect the repository" }], + tools: [ + { + name: "read", + description: "Read a file", + parameters: { + type: "object", + properties: { path: { type: "string" } }, + required: ["path"], + }, + }, + { + name: "webfetch", + description: "Fetch a URL", + parameters: { + type: "object", + properties: { url: { type: "string" } }, + required: ["url"], + }, + }, + ], + }, + ]); + expect(toolListCalls).toBe(1); expect( getJsonAttribute(firstGeneration, "langfuse.observation.usage_details"), ).toEqual({ @@ -622,6 +698,19 @@ describe.sequential("built plugin", () => { text: "The README contains the project overview.", }, }); + await emitEvent({ + id: "nested-observations-read-called", + type: "session.next.tool.called", + properties: { + sessionID, + timestamp: started + 260, + assistantMessageID, + callID: "nested-observations-call", + tool: "read", + input: { path: "README.md" }, + provider: { executed: false }, + }, + }); await hooks["tool.execute.before"]?.( { sessionID, callID: "nested-observations-call", tool: "read" }, { args: { path: "README.md" } }, @@ -635,12 +724,29 @@ describe.sequential("built plugin", () => { }, { title: "README.md", output: "# Project", metadata: {} }, ); + await emitEvent({ + type: "message.part.updated", + properties: { + part: { + id: "nested-observations-read-part", + sessionID, + messageID: assistantMessageID, + type: "tool", + callID: "nested-observations-call", + tool: "read", + state: { + status: "completed", + input: { path: "README.md" }, + title: "README.md", + output: "# Project", + metadata: {}, + time: { start: started + 260, end: started + 280 }, + }, + }, + }, + }); const failedToolStarted = Date.now(); const failedToolEnded = failedToolStarted + 5_000; - await hooks["tool.execute.before"]?.( - { sessionID, callID: "timed-out-webfetch", tool: "webfetch" }, - { args: { url: "https://example.com", timeout: 5 } }, - ); await emitEvent({ type: "message.part.updated", properties: { @@ -704,8 +810,6 @@ describe.sequential("built plugin", () => { "opencode.turn", "opencode.message.user", "opencode.generation", - "opencode.generation.reasoning", - "opencode.generation.reasoning", "read", "webfetch", "opencode.generation.retry", @@ -714,6 +818,39 @@ describe.sequential("built plugin", () => { ); const generation = getSpan(spans, "opencode.generation"); + expect(getJsonAttribute(generation, "langfuse.observation.output")).toEqual( + [ + { + role: "assistant", + content: "Repository inspected", + thinking: [ + { + type: "thinking", + content: "I should inspect the README.", + }, + { + type: "thinking", + content: "The README contains the project overview.", + }, + ], + tool_calls: [ + { + id: "nested-observations-call", + name: "read", + arguments: JSON.stringify({ path: "README.md" }), + }, + { + id: "timed-out-webfetch", + name: "webfetch", + arguments: JSON.stringify({ + url: "https://example.com", + timeout: 5, + }), + }, + ], + }, + ], + ); const tool = getSpan(spans, "read"); expect(getJsonAttribute(tool, "langfuse.observation.input")).toEqual({ @@ -752,23 +889,9 @@ describe.sequential("built plugin", () => { ); expect(compaction.parentSpanId).toBe(generation.spanId); - const reasoningSpans = spans.filter((span) => { - if (span.name !== "opencode.generation.reasoning") { - return false; - } - - const metadata = getJsonAttribute(span, "langfuse.observation.metadata"); - return ( - typeof metadata === "object" && - metadata !== null && - "messageID" in metadata && - metadata.messageID === assistantMessageID - ); - }); - expect(reasoningSpans).toHaveLength(2); - for (const reasoning of reasoningSpans) { - expect(reasoning.parentSpanId).toBe(generation.spanId); - } + expect( + spans.filter((span) => span.name === "opencode.generation.reasoning"), + ).toHaveLength(0); }); test("parents each tool to the generation that requested it when lifecycle events arrive out of order", async () => { @@ -912,6 +1035,263 @@ describe.sequential("built plugin", () => { expect(toolSpans.map((span) => span.parentSpanId)).toEqual( generationSpans.map((span) => span.spanId), ); + expect( + getJsonAttribute(generationSpans[1], "langfuse.observation.input"), + ).toEqual([ + { + role: "assistant", + tool_calls: [ + { + id: generations[0].callID, + name: generations[0].tool, + arguments: "{}", + }, + ], + }, + { + role: "tool", + name: generations[0].tool, + tool_call_id: generations[0].callID, + content: "ok", + }, + ]); + }); + + test("creates a nested tool observation from message parts without execution hooks", async () => { + const sessionID = "message-part-tool-session"; + const userMessageID = "message-part-tool-user"; + const firstAssistantMessageID = "message-part-tool-assistant-1"; + const secondAssistantMessageID = "message-part-tool-assistant-2"; + const callID = "message-part-tool-call"; + const started = startedAt; + const toolStarted = started + 200; + const toolCompleted = started + 350; + + await sendUserMessage({ + sessionID, + messageID: userMessageID, + text: "Show recent commits", + started, + }); + await startAssistantMessage({ + sessionID, + userMessageID, + assistantMessageID: firstAssistantMessageID, + started: started + 100, + }); + await startGeneration({ + id: "message-part-tool-step-1", + sessionID, + assistantMessageID: firstAssistantMessageID, + started: started + 100, + }); + + const runningPart = { + id: "message-part-tool-part", + sessionID, + messageID: firstAssistantMessageID, + type: "tool" as const, + callID, + tool: "bash", + state: { + status: "running" as const, + input: { command: "git log --oneline -10" }, + title: "Recent commits", + time: { start: toolStarted }, + }, + metadata: { providerExecuted: true }, + }; + await emitEvent({ + type: "message.part.updated", + properties: { part: runningPart }, + }); + await emitEvent({ + type: "message.part.updated", + properties: { part: runningPart }, + }); + await emitEvent({ + type: "message.part.updated", + properties: { + part: { + ...runningPart, + state: { + status: "completed" as const, + input: runningPart.state.input, + title: "Recent commits", + output: "07f9a68 Fix tool observation parenting", + metadata: {}, + time: { start: toolStarted, end: toolCompleted }, + }, + }, + }, + }); + await completeGeneration({ + sessionID, + userMessageID, + assistantMessageID: firstAssistantMessageID, + started: started + 100, + completed: started + 400, + }); + + await startAssistantMessage({ + sessionID, + userMessageID, + assistantMessageID: secondAssistantMessageID, + started: started + 500, + }); + await startGeneration({ + id: "message-part-tool-step-2", + sessionID, + assistantMessageID: secondAssistantMessageID, + started: started + 500, + }); + await completeGeneration({ + sessionID, + userMessageID, + assistantMessageID: secondAssistantMessageID, + started: started + 500, + completed: started + 700, + text: "Here are the recent commits.", + }); + + const { spans } = await flushSession(sessionID); + const generations = spans.filter( + (span) => span.name === "opencode.generation", + ); + const findGeneration = (messageID: string) => + generations.find((span) => { + const metadata = getJsonAttribute( + span, + "langfuse.observation.metadata", + ); + return ( + typeof metadata === "object" && + metadata !== null && + "messageID" in metadata && + metadata.messageID === messageID + ); + }); + const firstGeneration = findGeneration(firstAssistantMessageID); + const secondGeneration = findGeneration(secondAssistantMessageID); + expect(firstGeneration).toBeDefined(); + expect(secondGeneration).toBeDefined(); + + const toolSpans = spans.filter((span) => span.name === "bash"); + expect(toolSpans).toHaveLength(1); + expect(toolSpans[0].parentSpanId).toBe(firstGeneration!.spanId); + expect(toolSpans[0].startTimeUnixNano).toBe( + (BigInt(toolStarted) * 1_000_000n).toString(), + ); + expect(toolSpans[0].endTimeUnixNano).toBe( + (BigInt(toolCompleted) * 1_000_000n).toString(), + ); + expect( + getJsonAttribute(toolSpans[0], "langfuse.observation.input"), + ).toEqual({ command: "git log --oneline -10" }); + expect( + getJsonAttribute(toolSpans[0], "langfuse.observation.output"), + ).toEqual({ + title: "Recent commits", + output: "07f9a68 Fix tool observation parenting", + }); + expect( + getJsonAttribute(secondGeneration!, "langfuse.observation.input"), + ).toEqual([ + { + role: "assistant", + tool_calls: [ + { + id: callID, + name: "bash", + arguments: JSON.stringify({ command: "git log --oneline -10" }), + }, + ], + }, + { + role: "tool", + name: "bash", + tool_call_id: callID, + content: "07f9a68 Fix tool observation parenting", + }, + ]); + }); + + test("continues tracing when available tool discovery fails", async () => { + const sessionID = "tool-discovery-failure-session"; + toolListShouldFail = true; + + await sendUserMessage({ + sessionID, + messageID: "tool-discovery-failure-user", + text: "Continue without tool definitions", + started: startedAt, + }); + await startGeneration({ + id: "tool-discovery-failure-step", + sessionID, + started: startedAt + 100, + }); + await completeGeneration({ + sessionID, + userMessageID: "tool-discovery-failure-user", + assistantMessageID: "tool-discovery-failure-assistant", + started: startedAt + 100, + completed: startedAt + 500, + text: "Completed", + }); + + const { spans } = await flushSession(sessionID); + expect( + getJsonAttribute( + getSpan(spans, "opencode.generation"), + "langfuse.observation.input", + ), + ).toEqual([ + { + role: "user", + content: [{ type: "text", text: "Continue without tool definitions" }], + }, + ]); + + toolListShouldFail = false; + const retrySessionID = "tool-discovery-retry-session"; + await sendUserMessage({ + sessionID: retrySessionID, + messageID: "tool-discovery-retry-user", + text: "Retry tool discovery", + started: startedAt + 1_000, + }); + await startGeneration({ + id: "tool-discovery-retry-step", + sessionID: retrySessionID, + started: startedAt + 1_100, + }); + await completeGeneration({ + sessionID: retrySessionID, + userMessageID: "tool-discovery-retry-user", + assistantMessageID: "tool-discovery-retry-assistant", + started: startedAt + 1_100, + completed: startedAt + 1_500, + text: "Completed with tools", + }); + + const retry = await flushSession(retrySessionID); + expect(toolListCalls).toBe(2); + expect( + getJsonAttribute( + getSpan(retry.spans, "opencode.generation"), + "langfuse.observation.input", + ), + ).toEqual([ + { + role: "user", + content: [{ type: "text", text: "Retry tool discovery" }], + tools: expect.arrayContaining([ + expect.objectContaining({ name: "read" }), + expect.objectContaining({ name: "webfetch" }), + ]), + }, + ]); }); test("preserves step metadata when the assistant message arrives later", async () => { @@ -1070,9 +1450,7 @@ describe.sequential("built plugin", () => { const turn = getSpan(spans, "opencode.turn"); expect(getJsonAttribute(generation, "langfuse.observation.output")).toEqual( - { - text: "Fallback output", - }, + [{ role: "assistant", content: "Fallback output" }], ); expect(generation.traceId).toBe(turn.traceId); expect(generation.parentSpanId).toBe(turn.spanId); @@ -1142,9 +1520,21 @@ describe.sequential("built plugin", () => { "opencode.message.user", "opencode.generation", "opencode.generation.retry", - "opencode.generation.reasoning", ].sort(), ); + + expect( + getJsonAttribute( + getSpan(spans, "opencode.generation"), + "langfuse.observation.output", + ), + ).toEqual([ + { + role: "assistant", + content: "One output", + thinking: [{ type: "thinking", content: "Think once" }], + }, + ]); }); test("does not reject hooks when the collector returns an error", async () => { From e14eb6f4cd1c21670edfec50df0098a4c06a9389 Mon Sep 17 00:00:00 2001 From: Ben Bachem <10088265+bezbac@users.noreply.github.com> Date: Thu, 6 Aug 2026 18:35:03 +0200 Subject: [PATCH 2/2] Refine ChatML tracing --- src/index.ts | 10 ++--- src/langfuse.ts | 69 ++++++++++++--------------------- test/integration/plugin.test.ts | 5 +++ 3 files changed, 33 insertions(+), 51 deletions(-) diff --git a/src/index.ts b/src/index.ts index c24c3e1..a104475 100644 --- a/src/index.ts +++ b/src/index.ts @@ -8,7 +8,6 @@ import { Data, Effect, Layer, Schema } from "effect"; import { LangfuseClientService, createLangfuseClient, - type ActiveGenerationStep, type LangfuseClient, type ToolDefinition, } from "./langfuse.js"; @@ -26,7 +25,9 @@ type SessionNextEvent = timestamp: number; assistantMessageID?: string; agent: string; - model: NonNullable; + model: Parameters< + LangfuseClient["startActiveGenerationStep"] + >[0]["model"]; snapshot?: string; }; } @@ -502,10 +503,7 @@ const main = Effect.gen(function* () { tool.parameters !== null && !Array.isArray(tool.parameters) ? { - parameters: tool.parameters as Record< - string, - unknown - >, + parameters: tool.parameters, } : {}), })), diff --git a/src/langfuse.ts b/src/langfuse.ts index 50b3ee8..c843acd 100644 --- a/src/langfuse.ts +++ b/src/langfuse.ts @@ -205,9 +205,19 @@ export class LangfuseClient { } traceReasoningPart(part: MessagePart) { + if ( + part.type !== "reasoning" || + typeof part.id !== "string" || + typeof part.sessionID !== "string" || + typeof part.messageID !== "string" || + typeof part.text !== "string" + ) { + return; + } + const completed = getCompletedReasoningTimestamp(part); - if (!isCompletedReasoningPart(part) || completed === undefined) { + if (completed === undefined) { return; } @@ -252,7 +262,6 @@ export class LangfuseClient { ...input.model, variant: input.model.variant ?? existingMessageStep.model?.variant, }, - started: input.started, snapshot: input.snapshot ?? existingMessageStep.snapshot, }; @@ -291,7 +300,6 @@ export class LangfuseClient { ...input.model, variant: input.model.variant ?? existingStep.model?.variant, }, - started: input.started, snapshot: input.snapshot ?? existingStep.snapshot, }; @@ -358,7 +366,6 @@ export class LangfuseClient { agent: input.agent, model: input.model, span, - started: input.started, snapshot: input.snapshot, }); if (messageID) { @@ -456,7 +463,7 @@ export class LangfuseClient { const span = this.traceState.tracer.startSpan("opencode.turn", { attributes: { - "langfuse.observation.type": "span", + "langfuse.observation.type": "agent", "langfuse.internal.is_app_root": true, "session.id": input.sessionID, "langfuse.observation.input": JSON.stringify([formattedMessage]), @@ -519,26 +526,19 @@ export class LangfuseClient { this.traceState.assistantParts.set(part.messageID, parts); if (part.type === "tool") { - this.rememberToolCall({ - callID: part.callID, - messageID: part.messageID, - }); + this.traceState.toolMessageIdsByCallId.set(part.callID, part.messageID); } } rememberToolCall(input: { callID: string; messageID: string; - sessionID?: string; - tool?: string; - args?: Record; + sessionID: string; + tool: string; + args: Record; }) { this.traceState.toolMessageIdsByCallId.set(input.callID, input.messageID); - if (!input.sessionID || !input.tool || !input.args) { - return; - } - const parts = this.traceState.assistantParts.get(input.messageID) ?? new Map(); @@ -1207,36 +1207,17 @@ export type MessagePart = Extract< { type: "message.part.updated" } >["properties"]["part"]; -type CompletedReasoningPart = MessagePart & { - id: string; - sessionID: string; - text: string; - messageID: string; - time: { completed?: number; end?: number }; -}; - -function isCompletedReasoningPart( - part: MessagePart, -): part is CompletedReasoningPart { - return ( - part.type === "reasoning" && - typeof part.id === "string" && - typeof part.sessionID === "string" && - typeof part.messageID === "string" && - typeof part.text === "string" && - typeof getCompletedReasoningTimestamp(part) === "number" - ); -} - function getCompletedReasoningTimestamp(part: MessagePart) { - const time = (part as { time?: { completed?: unknown; end?: unknown } }).time; + if (!("time" in part) || !part.time || typeof part.time !== "object") { + return undefined; + } - if (typeof time?.completed === "number") { - return time.completed; + if ("completed" in part.time && typeof part.time.completed === "number") { + return part.time.completed; } - if (typeof time?.end === "number") { - return time.end; + if ("end" in part.time && typeof part.time.end === "number") { + return part.time.end; } return undefined; @@ -1266,7 +1247,7 @@ export type UserMessageInput = { export type ToolDefinition = { name: string; description?: string; - parameters?: Record; + parameters?: object; }; type ChatMlMessage = @@ -1301,9 +1282,7 @@ export type ActiveGenerationStep = { variant?: string; }; span: ApiSpan; - started?: number; snapshot?: string; - input?: ChatMlMessage[]; }; export class LangfuseClientService extends EffectContext.Tag( diff --git a/test/integration/plugin.test.ts b/test/integration/plugin.test.ts index 27dbee1..3203b04 100644 --- a/test/integration/plugin.test.ts +++ b/test/integration/plugin.test.ts @@ -578,6 +578,11 @@ describe.sequential("built plugin", () => { }, ]); } + for (const turn of spans.filter((span) => span.name === "opencode.turn")) { + expect(getAttributes(turn)).toMatchObject({ + "langfuse.observation.type": "agent", + }); + } const firstGeneration = spans .filter((span) => span.name === "opencode.generation")