From 410a9fd8b9e5f882f9606d33f8583f96b3a75946 Mon Sep 17 00:00:00 2001 From: Ben Bachem <10088265+bezbac@users.noreply.github.com> Date: Thu, 6 Aug 2026 15:34:28 +0200 Subject: [PATCH 1/2] Add tool parenting reproduction --- test/integration/plugin.test.ts | 112 ++++++++++++++++++++++++++++++++ 1 file changed, 112 insertions(+) diff --git a/test/integration/plugin.test.ts b/test/integration/plugin.test.ts index 72616a8..aa47bfc 100644 --- a/test/integration/plugin.test.ts +++ b/test/integration/plugin.test.ts @@ -724,6 +724,118 @@ describe.sequential("built plugin", () => { } }); + test.fails( + "parents each tool to the generation that requested it when lifecycle events arrive out of order", + async () => { + const sessionID = "out-of-order-tool-parenting-session"; + const userMessageID = "out-of-order-tool-parenting-user"; + const started = startedAt; + const generations = [1, 2, 3].map((index) => ({ + assistantMessageID: `out-of-order-tool-parenting-assistant-${index}`, + stepID: `out-of-order-tool-parenting-step-${index}`, + tool: `out-of-order-tool-${index}`, + callID: `out-of-order-tool-parenting-call-${index}`, + })); + + const executeTool = async (generation: (typeof generations)[number]) => { + await hooks["tool.execute.before"]?.( + { sessionID, callID: generation.callID, tool: generation.tool }, + { args: { generation: generation.assistantMessageID } }, + ); + await hooks["tool.execute.after"]?.( + { + sessionID, + callID: generation.callID, + tool: generation.tool, + args: { generation: generation.assistantMessageID }, + }, + { title: generation.tool, output: "ok", metadata: {} }, + ); + }; + + await sendUserMessage({ + sessionID, + messageID: userMessageID, + text: "Run three tool batches", + started, + }); + + await startGeneration({ + id: generations[0].stepID, + sessionID, + started: started + 100, + }); + await executeTool(generations[0]); + + // The next generation and its tool begin before the previous generation's + // completed message event reaches the plugin. + await startGeneration({ + id: generations[1].stepID, + sessionID, + started: started + 200, + }); + await executeTool(generations[1]); + await completeGeneration({ + sessionID, + userMessageID, + assistantMessageID: generations[0].assistantMessageID, + started: started + 100, + completed: started + 190, + }); + + await startGeneration({ + id: generations[2].stepID, + sessionID, + started: started + 300, + }); + await executeTool(generations[2]); + await completeGeneration({ + sessionID, + userMessageID, + assistantMessageID: generations[1].assistantMessageID, + started: started + 200, + completed: started + 290, + }); + await completeGeneration({ + sessionID, + userMessageID, + assistantMessageID: generations[2].assistantMessageID, + started: started + 300, + completed: started + 390, + }); + + const { spans } = await flushSession(sessionID); + const generationSpans = generations.map((generation) => { + const span = spans.find((candidate) => { + if (candidate.name !== "opencode.generation") { + return false; + } + + const metadata = getJsonAttribute( + candidate, + "langfuse.observation.metadata", + ); + return ( + typeof metadata === "object" && + metadata !== null && + "messageID" in metadata && + metadata.messageID === generation.assistantMessageID + ); + }); + + expect(span).toBeDefined(); + return span!; + }); + const toolSpans = generations.map((generation) => + getSpan(spans, generation.tool), + ); + + expect(toolSpans.map((span) => span.parentSpanId)).toEqual( + generationSpans.map((span) => span.spanId), + ); + }, + ); + test("exports a failed generation as an error span", async () => { const sessionID = "failed-session"; const started = startedAt; From e564d198aa00b736529e58622823ad0d6ae1e756 Mon Sep 17 00:00:00 2001 From: Ben Bachem <10088265+bezbac@users.noreply.github.com> Date: Thu, 6 Aug 2026 16:45:03 +0200 Subject: [PATCH 2/2] Fix tool observation parenting --- src/index.ts | 44 ++++- src/langfuse.ts | 235 +++++++++++++++++++---- test/integration/plugin.test.ts | 321 ++++++++++++++++++++++---------- 3 files changed, 459 insertions(+), 141 deletions(-) diff --git a/src/index.ts b/src/index.ts index b7bf569..e1deda1 100644 --- a/src/index.ts +++ b/src/index.ts @@ -23,6 +23,7 @@ type SessionNextEvent = properties: { sessionID: string; timestamp: number; + assistantMessageID?: string; agent: string; model: NonNullable; snapshot?: string; @@ -39,9 +40,23 @@ type SessionNextEvent = properties: { sessionID: string; timestamp: number; + assistantMessageID?: string; error: { message: string }; }; } + | { + id: string; + type: "session.next.tool.called"; + properties: { + sessionID: string; + timestamp: number; + assistantMessageID?: string; + callID: string; + tool: string; + input: Record; + provider: { executed: boolean; metadata?: unknown }; + }; + } | { id: string; type: "session.next.retried"; @@ -183,6 +198,7 @@ const eventHook = (event: OpencodeEvent, shutdown?: () => Promise) => if (event.type === "session.next.step.started") { langfuse.startActiveGenerationStep({ sessionID: event.properties.sessionID, + assistantMessageID: event.properties.assistantMessageID, agent: event.properties.agent, model: event.properties.model, started: event.properties.timestamp, @@ -194,11 +210,22 @@ const eventHook = (event: OpencodeEvent, shutdown?: () => Promise) => langfuse.traceFailedGenerationStep({ id: event.id, sessionID: event.properties.sessionID, + assistantMessageID: event.properties.assistantMessageID, completed: event.properties.timestamp, error: event.properties.error, }); } + if ( + event.type === "session.next.tool.called" && + event.properties.assistantMessageID + ) { + langfuse.rememberToolCall({ + callID: event.properties.callID, + messageID: event.properties.assistantMessageID, + }); + } + if (event.type === "session.next.retried") { langfuse.traceEvent({ id: event.id, @@ -239,7 +266,22 @@ const eventHook = (event: OpencodeEvent, shutdown?: () => Promise) => if (event.type === "message.updated") { const message = event.properties.info; - if (message.role !== "assistant" || !message.time.completed) { + if (message.role !== "assistant") { + return; + } + + langfuse.startActiveGenerationStep({ + sessionID: message.sessionID, + assistantMessageID: message.id, + agent: message.mode, + model: { + id: message.modelID, + providerID: message.providerID, + }, + started: message.time.created, + }); + + if (!message.time.completed) { return; } diff --git a/src/langfuse.ts b/src/langfuse.ts index 7eb2e0f..916d910 100644 --- a/src/langfuse.ts +++ b/src/langfuse.ts @@ -37,6 +37,8 @@ export class LangfuseClient { 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.turnObservationsByMessageId.clear(); this.traceState.latestTurnObservationsBySession.clear(); @@ -63,13 +65,18 @@ export class LangfuseClient { observation.span.end(); this.traceState.activeToolObservations.delete(callID); this.traceState.finalizedToolCallIds.add(callID); + this.traceState.toolMessageIdsByCallId.delete(callID); } } endActiveGenerationSteps(sessionID?: string, error?: SessionErrorInfo) { - for (const [activeSessionID, step] of this.traceState - .activeGenerationSteps) { - if (sessionID && activeSessionID !== sessionID) { + const activeSteps = new Set([ + ...this.traceState.activeGenerationSteps.values(), + ...this.traceState.activeGenerationStepsByMessageId.values(), + ]); + + for (const step of activeSteps) { + if (sessionID && step.sessionID !== sessionID) { continue; } @@ -84,8 +91,21 @@ export class LangfuseClient { } step.span.end(); - this.traceState.activeGenerationSteps.delete(activeSessionID); - this.traceState.generationParentSpans.delete(activeSessionID); + } + + for (const [activeSessionID, step] of this.traceState + .activeGenerationSteps) { + if (!sessionID || step.sessionID === sessionID) { + this.traceState.activeGenerationSteps.delete(activeSessionID); + this.traceState.generationParentSpans.delete(activeSessionID); + } + } + + for (const [messageID, step] of this.traceState + .activeGenerationStepsByMessageId) { + if (!sessionID || step.sessionID === sessionID) { + this.traceState.activeGenerationStepsByMessageId.delete(messageID); + } } } @@ -230,16 +250,79 @@ export class LangfuseClient { startActiveGenerationStep(input: { sessionID: string; + assistantMessageID?: string; agent: string; model: NonNullable; started: number; snapshot?: string; }) { + const messageID = input.assistantMessageID; + const existingMessageStep = messageID + ? this.traceState.activeGenerationStepsByMessageId.get(messageID) + : undefined; const existingStep = this.traceState.activeGenerationSteps.get( input.sessionID, ); - if (existingStep && !existingStep.model) { + if ( + messageID && + !existingMessageStep && + this.traceState.generationSpansByMessageId.has(messageID) + ) { + return; + } + + if (existingMessageStep && messageID) { + const updatedStep = { + ...existingMessageStep, + agent: input.agent, + model: { + ...input.model, + variant: input.model.variant ?? existingMessageStep.model?.variant, + }, + started: input.started, + snapshot: input.snapshot ?? existingMessageStep.snapshot, + }; + + existingMessageStep.span.setAttribute( + "langfuse.observation.model.name", + input.model.id, + ); + existingMessageStep.span.setAttribute( + "langfuse.observation.metadata", + JSON.stringify({ + agent: updatedStep.agent, + providerID: updatedStep.model?.providerID, + variant: updatedStep.model?.variant, + snapshot: updatedStep.snapshot, + }), + ); + this.traceState.activeGenerationStepsByMessageId.set( + messageID, + updatedStep, + ); + + if (!existingStep || existingStep.messageID === messageID) { + this.traceState.activeGenerationSteps.set(input.sessionID, updatedStep); + } + + return; + } + + if (existingStep && !existingStep.messageID && messageID) { + const updatedStep = { + ...existingStep, + sessionID: input.sessionID, + messageID, + agent: input.agent, + model: { + ...input.model, + variant: input.model.variant ?? existingStep.model?.variant, + }, + started: input.started, + snapshot: input.snapshot ?? existingStep.snapshot, + }; + existingStep.span.setAttribute( "langfuse.observation.model.name", input.model.id, @@ -253,18 +336,22 @@ export class LangfuseClient { snapshot: input.snapshot, }), ); - this.traceState.activeGenerationSteps.set(input.sessionID, { - ...existingStep, - agent: input.agent, - model: input.model, - started: input.started, - snapshot: input.snapshot, - }); + this.traceState.activeGenerationSteps.set(input.sessionID, updatedStep); + this.traceState.activeGenerationStepsByMessageId.set( + messageID, + updatedStep, + ); + this.traceState.generationSpansByMessageId.set( + messageID, + existingStep.span, + ); return; } - existingStep?.span.end(new Date(input.started)); + if (!messageID && existingStep) { + return; + } if (!this.getTurnObservation(input.sessionID, undefined)) { return; @@ -287,12 +374,21 @@ export class LangfuseClient { }); this.traceState.activeGenerationSteps.set(input.sessionID, { + sessionID: input.sessionID, + messageID, agent: input.agent, model: input.model, span, started: input.started, snapshot: input.snapshot, }); + if (messageID) { + const step = this.traceState.activeGenerationSteps.get(input.sessionID); + if (step) { + this.traceState.activeGenerationStepsByMessageId.set(messageID, step); + this.traceState.generationSpansByMessageId.set(messageID, span); + } + } this.traceState.generationParentSpans.set(input.sessionID, span); }); } @@ -430,6 +526,17 @@ export class LangfuseClient { parts.set(part.id, part); this.traceState.assistantParts.set(part.messageID, parts); + + if (part.type === "tool") { + this.rememberToolCall({ + callID: part.callID, + messageID: part.messageID, + }); + } + } + + rememberToolCall(input: { callID: string; messageID: string }) { + this.traceState.toolMessageIdsByCallId.set(input.callID, input.messageID); } traceGeneration(input: { @@ -472,7 +579,14 @@ export class LangfuseClient { JSON.stringify(output), ); } - const step = this.traceState.activeGenerationSteps.get(input.sessionID); + const activeStep = this.traceState.activeGenerationSteps.get( + input.sessionID, + ); + const step = + this.traceState.activeGenerationStepsByMessageId.get(input.messageID) ?? + (activeStep?.messageID === input.messageID || !activeStep?.messageID + ? activeStep + : undefined); if (step) { step.span.setAttribute("langfuse.observation.model.name", input.modelID); @@ -518,7 +632,11 @@ export class LangfuseClient { this.flushPendingReasoning(input.messageID, step.span); step.span.end(new Date(input.completed)); - this.traceState.activeGenerationSteps.delete(input.sessionID); + this.traceState.activeGenerationStepsByMessageId.delete(input.messageID); + + if (activeStep === step) { + this.traceState.activeGenerationSteps.delete(input.sessionID); + } return; } @@ -594,6 +712,7 @@ export class LangfuseClient { traceFailedGenerationStep(input: { id: string; sessionID: string; + assistantMessageID?: string; completed: number; error: { message: string }; }) { @@ -603,7 +722,18 @@ export class LangfuseClient { this.traceState.tracedGenerationIds.add(input.id); - const step = this.traceState.activeGenerationSteps.get(input.sessionID); + const activeStep = this.traceState.activeGenerationSteps.get( + input.sessionID, + ); + const step = input.assistantMessageID + ? (this.traceState.activeGenerationStepsByMessageId.get( + input.assistantMessageID, + ) ?? + (activeStep?.messageID === input.assistantMessageID || + !activeStep?.messageID + ? activeStep + : undefined)) + : activeStep; if (step) { step.span.setAttribute( @@ -625,7 +755,14 @@ export class LangfuseClient { }); step.span.recordException(input.error); step.span.end(new Date(input.completed)); - this.traceState.activeGenerationSteps.delete(input.sessionID); + const messageID = input.assistantMessageID ?? step.messageID; + if (messageID) { + this.traceState.activeGenerationStepsByMessageId.delete(messageID); + } + + if (activeStep === step) { + this.traceState.activeGenerationSteps.delete(input.sessionID); + } return; } @@ -712,25 +849,29 @@ export class LangfuseClient { this.traceState.finalizedToolCallIds.delete(input.callID); this.ensureGenerationParent(input.sessionID); - this.withObservationParent(input.sessionID, () => { - const span = this.traceState.tracer.startSpan(input.tool, { - attributes: { - "langfuse.observation.type": "tool", - "session.id": input.sessionID, - "langfuse.observation.input": JSON.stringify(input.args), - "langfuse.observation.metadata": JSON.stringify({ - callID: input.callID, - tool: input.tool, - }), - }, - }); + this.withObservationParent( + input.sessionID, + () => { + const span = this.traceState.tracer.startSpan(input.tool, { + attributes: { + "langfuse.observation.type": "tool", + "session.id": input.sessionID, + "langfuse.observation.input": JSON.stringify(input.args), + "langfuse.observation.metadata": JSON.stringify({ + callID: input.callID, + tool: input.tool, + }), + }, + }); - this.traceState.activeToolObservations.set(input.callID, { - span, - sessionID: input.sessionID, - tool: input.tool, - }); - }); + this.traceState.activeToolObservations.set(input.callID, { + span, + sessionID: input.sessionID, + tool: input.tool, + }); + }, + this.traceState.toolMessageIdsByCallId.get(input.callID), + ); } traceToolEnd(input: { @@ -775,6 +916,7 @@ export class LangfuseClient { span.end(); 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 }) { @@ -800,6 +942,7 @@ export class LangfuseClient { span.end(new Date(input.completed)); this.traceState.activeToolObservations.delete(input.callID); this.traceState.finalizedToolCallIds.add(input.callID); + this.traceState.toolMessageIdsByCallId.delete(input.callID); } private ensureGenerationParent(sessionID: string) { @@ -822,7 +965,10 @@ export class LangfuseClient { }, }); - this.traceState.activeGenerationSteps.set(sessionID, { span }); + this.traceState.activeGenerationSteps.set(sessionID, { + sessionID, + span, + }); this.traceState.generationParentSpans.set(sessionID, span); }); } @@ -848,8 +994,15 @@ export class LangfuseClient { ); } - private withObservationParent(sessionID: string, fn: () => T) { + private withObservationParent( + sessionID: string, + fn: () => T, + messageID?: string, + ) { const parentSpan = + (messageID + ? this.traceState.generationSpansByMessageId.get(messageID) + : undefined) ?? this.traceState.activeGenerationSteps.get(sessionID)?.span ?? this.traceState.generationParentSpans.get(sessionID); @@ -902,6 +1055,8 @@ export type LangfuseTraceState = { Map >; generationSpansByMessageId: Map; + activeGenerationStepsByMessageId: Map; + toolMessageIdsByCallId: Map; assistantParts: Map>; turnObservationsByMessageId: Map; latestTurnObservationsBySession: Map; @@ -984,6 +1139,8 @@ export type ToolObservation = { }; export type ActiveGenerationStep = { + sessionID: string; + messageID?: string; agent?: string; model?: { id: string; @@ -1059,6 +1216,8 @@ export const createLangfuseClient = (input: { Map >(), generationSpansByMessageId: new Map(), + activeGenerationStepsByMessageId: new Map(), + toolMessageIdsByCallId: new Map(), assistantParts: new Map>(), turnObservationsByMessageId: new Map(), latestTurnObservationsBySession: new Map(), diff --git a/test/integration/plugin.test.ts b/test/integration/plugin.test.ts index aa47bfc..133b5bf 100644 --- a/test/integration/plugin.test.ts +++ b/test/integration/plugin.test.ts @@ -52,6 +52,7 @@ type SessionNextEvent = properties: { sessionID: string; timestamp: number; + assistantMessageID?: string; agent: string; model: { id: string; @@ -70,6 +71,19 @@ type SessionNextEvent = error: { message: string }; }; } + | { + id: string; + type: "session.next.tool.called"; + properties: { + sessionID: string; + timestamp: number; + assistantMessageID?: string; + callID: string; + tool: string; + input: Record; + provider: { executed: boolean; metadata?: unknown }; + }; + } | { id: string; type: "session.next.retried"; @@ -235,6 +249,7 @@ const sendUserMessage = async (input: { const startGeneration = async (input: { id: string; sessionID: string; + assistantMessageID?: string; started: number; snapshot?: string; }) => { @@ -244,6 +259,7 @@ const startGeneration = async (input: { properties: { sessionID: input.sessionID, timestamp: input.started, + assistantMessageID: input.assistantMessageID, agent: "build", model: { id: "test-model", @@ -255,6 +271,37 @@ const startGeneration = async (input: { }); }; +const startAssistantMessage = async (input: { + sessionID: string; + userMessageID: string; + assistantMessageID: string; + started: number; +}) => { + await emitEvent({ + type: "message.updated", + properties: { + info: { + id: input.assistantMessageID, + sessionID: input.sessionID, + parentID: input.userMessageID, + role: "assistant", + mode: "build", + modelID: "test-model", + providerID: "test-provider", + path: { cwd: "/test", root: "/test" }, + cost: 0, + tokens: { + input: 0, + output: 0, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + time: { created: input.started }, + }, + }, + }); +}; + const completeGeneration = async (input: { sessionID: string; userMessageID: string; @@ -724,117 +771,187 @@ describe.sequential("built plugin", () => { } }); - test.fails( - "parents each tool to the generation that requested it when lifecycle events arrive out of order", - async () => { - const sessionID = "out-of-order-tool-parenting-session"; - const userMessageID = "out-of-order-tool-parenting-user"; - const started = startedAt; - const generations = [1, 2, 3].map((index) => ({ - assistantMessageID: `out-of-order-tool-parenting-assistant-${index}`, - stepID: `out-of-order-tool-parenting-step-${index}`, - tool: `out-of-order-tool-${index}`, - callID: `out-of-order-tool-parenting-call-${index}`, - })); - - const executeTool = async (generation: (typeof generations)[number]) => { - await hooks["tool.execute.before"]?.( - { sessionID, callID: generation.callID, tool: generation.tool }, - { args: { generation: generation.assistantMessageID } }, + test("parents each tool to the generation that requested it when lifecycle events arrive out of order", async () => { + const sessionID = "out-of-order-tool-parenting-session"; + const userMessageID = "out-of-order-tool-parenting-user"; + const started = startedAt; + const generations = [1, 2, 3].map((index) => ({ + assistantMessageID: `out-of-order-tool-parenting-assistant-${index}`, + stepID: `out-of-order-tool-parenting-step-${index}`, + tool: `out-of-order-tool-${index}`, + callID: `out-of-order-tool-parenting-call-${index}`, + })); + + const executeTool = async (generation: (typeof generations)[number]) => { + await emitEvent({ + id: `${generation.callID}-called`, + type: "session.next.tool.called", + properties: { + sessionID, + timestamp: started, + assistantMessageID: generation.assistantMessageID, + callID: generation.callID, + tool: generation.tool, + input: {}, + provider: { executed: false }, + }, + }); + await hooks["tool.execute.before"]?.( + { sessionID, callID: generation.callID, tool: generation.tool }, + { args: { generation: generation.assistantMessageID } }, + ); + await hooks["tool.execute.after"]?.( + { + sessionID, + callID: generation.callID, + tool: generation.tool, + args: { generation: generation.assistantMessageID }, + }, + { title: generation.tool, output: "ok", metadata: {} }, + ); + }; + + await sendUserMessage({ + sessionID, + messageID: userMessageID, + text: "Run three tool batches", + started, + }); + + await startAssistantMessage({ + sessionID, + userMessageID, + assistantMessageID: generations[0].assistantMessageID, + started: started + 100, + }); + await startGeneration({ + id: generations[0].stepID, + sessionID, + assistantMessageID: generations[0].assistantMessageID, + started: started + 100, + }); + await executeTool(generations[0]); + + // The next generation and its tool begin before the previous generation's + // completed message event reaches the plugin. + await startAssistantMessage({ + sessionID, + userMessageID, + assistantMessageID: generations[1].assistantMessageID, + started: started + 200, + }); + await startGeneration({ + id: generations[1].stepID, + sessionID, + assistantMessageID: generations[1].assistantMessageID, + started: started + 200, + }); + await executeTool(generations[1]); + await completeGeneration({ + sessionID, + userMessageID, + assistantMessageID: generations[0].assistantMessageID, + started: started + 100, + completed: started + 190, + }); + + await startAssistantMessage({ + sessionID, + userMessageID, + assistantMessageID: generations[2].assistantMessageID, + started: started + 300, + }); + await startGeneration({ + id: generations[2].stepID, + sessionID, + assistantMessageID: generations[2].assistantMessageID, + started: started + 300, + }); + await executeTool(generations[2]); + await completeGeneration({ + sessionID, + userMessageID, + assistantMessageID: generations[1].assistantMessageID, + started: started + 200, + completed: started + 290, + }); + await completeGeneration({ + sessionID, + userMessageID, + assistantMessageID: generations[2].assistantMessageID, + started: started + 300, + completed: started + 390, + }); + + const { spans } = await flushSession(sessionID); + const generationSpans = generations.map((generation) => { + const span = spans.find((candidate) => { + if (candidate.name !== "opencode.generation") { + return false; + } + + const metadata = getJsonAttribute( + candidate, + "langfuse.observation.metadata", ); - await hooks["tool.execute.after"]?.( - { - sessionID, - callID: generation.callID, - tool: generation.tool, - args: { generation: generation.assistantMessageID }, - }, - { title: generation.tool, output: "ok", metadata: {} }, + return ( + typeof metadata === "object" && + metadata !== null && + "messageID" in metadata && + metadata.messageID === generation.assistantMessageID ); - }; - - await sendUserMessage({ - sessionID, - messageID: userMessageID, - text: "Run three tool batches", - started, }); - await startGeneration({ - id: generations[0].stepID, - sessionID, - started: started + 100, - }); - await executeTool(generations[0]); + expect(span).toBeDefined(); + return span!; + }); + const toolSpans = generations.map((generation) => + getSpan(spans, generation.tool), + ); - // The next generation and its tool begin before the previous generation's - // completed message event reaches the plugin. - await startGeneration({ - id: generations[1].stepID, - sessionID, - started: started + 200, - }); - await executeTool(generations[1]); - await completeGeneration({ - sessionID, - userMessageID, - assistantMessageID: generations[0].assistantMessageID, - started: started + 100, - completed: started + 190, - }); + expect(toolSpans.map((span) => span.parentSpanId)).toEqual( + generationSpans.map((span) => span.spanId), + ); + }); - await startGeneration({ - id: generations[2].stepID, - sessionID, - started: started + 300, - }); - await executeTool(generations[2]); - await completeGeneration({ - sessionID, - userMessageID, - assistantMessageID: generations[1].assistantMessageID, - started: started + 200, - completed: started + 290, - }); - await completeGeneration({ - sessionID, - userMessageID, - assistantMessageID: generations[2].assistantMessageID, - started: started + 300, - completed: started + 390, - }); + test("preserves step metadata when the assistant message arrives later", async () => { + const sessionID = "late-assistant-message-session"; + const userMessageID = "late-assistant-message-user"; + const assistantMessageID = "late-assistant-message-assistant"; + const started = startedAt; - const { spans } = await flushSession(sessionID); - const generationSpans = generations.map((generation) => { - const span = spans.find((candidate) => { - if (candidate.name !== "opencode.generation") { - return false; - } - - const metadata = getJsonAttribute( - candidate, - "langfuse.observation.metadata", - ); - return ( - typeof metadata === "object" && - metadata !== null && - "messageID" in metadata && - metadata.messageID === generation.assistantMessageID - ); - }); + await sendUserMessage({ + sessionID, + messageID: userMessageID, + text: "Keep the generation metadata", + started, + }); + await startGeneration({ + id: "late-assistant-message-step", + sessionID, + assistantMessageID, + started: started + 100, + snapshot: "snapshot-1", + }); + await startAssistantMessage({ + sessionID, + userMessageID, + assistantMessageID, + started: started + 100, + }); - expect(span).toBeDefined(); - return span!; - }); - const toolSpans = generations.map((generation) => - getSpan(spans, generation.tool), - ); + const { spans } = await flushSession(sessionID); + const generation = getSpan(spans, "opencode.generation"); - expect(toolSpans.map((span) => span.parentSpanId)).toEqual( - generationSpans.map((span) => span.spanId), - ); - }, - ); + expect( + getJsonAttribute(generation, "langfuse.observation.metadata"), + ).toEqual({ + agent: "build", + providerID: "test-provider", + variant: "high", + snapshot: "snapshot-1", + }); + }); test("exports a failed generation as an error span", async () => { const sessionID = "failed-session";