From 8e6f86a904e3a9bc1d5ce15d10d30147e2b9b780 Mon Sep 17 00:00:00 2001 From: idityage Date: Sun, 9 Aug 2026 22:33:47 +0530 Subject: [PATCH] Link child sessions to parent traces --- src/index.ts | 30 ++++- src/langfuse.ts | 206 +++++++++++++++++++++++++------- test/integration/plugin.test.ts | 110 +++++++++++++++++ 3 files changed, 298 insertions(+), 48 deletions(-) diff --git a/src/index.ts b/src/index.ts index a104475..7893e0b 100644 --- a/src/index.ts +++ b/src/index.ts @@ -150,16 +150,21 @@ const eventHook = (event: OpencodeEvent, shutdown?: () => Promise) => Effect.gen(function* () { const langfuse = yield* LangfuseClientService; - const finalizeSessionTracing = () => { - langfuse.endActiveToolObservations(); - langfuse.endActiveGenerationSteps(); - langfuse.endActiveTurnObservations(); - langfuse.clearTraceState(); + const finalizeSessionTracing = (sessionID?: string) => { + langfuse.endActiveToolObservations(sessionID); + langfuse.endActiveGenerationSteps(sessionID); + langfuse.endActiveTurnObservations(sessionID); + + if (sessionID) { + langfuse.clearSessionTraceState(sessionID); + } else { + langfuse.clearTraceState(); + } }; if (event.type === "session.idle") { yield* log("info", "Flushing spans"); - finalizeSessionTracing(); + finalizeSessionTracing(event.properties.sessionID); yield* langfuse.forceFlush; } @@ -175,6 +180,19 @@ const eventHook = (event: OpencodeEvent, shutdown?: () => Promise) => } } + if (event.type === "session.created" || event.type === "session.updated") { + langfuse.rememberSessionParent({ + sessionID: event.properties.info.id, + parentSessionID: event.properties.info.parentID, + }); + } + + if (event.type === "session.deleted") { + langfuse.rememberSessionParent({ + sessionID: event.properties.info.id, + }); + } + if (event.type === "session.error" && event.properties.sessionID) { langfuse.traceSessionError({ sessionID: event.properties.sessionID, diff --git a/src/langfuse.ts b/src/langfuse.ts index c843acd..687b142 100644 --- a/src/langfuse.ts +++ b/src/langfuse.ts @@ -44,6 +44,61 @@ export class LangfuseClient { this.traceState.turnObservationsByMessageId.clear(); this.traceState.latestTurnObservationsBySession.clear(); this.traceState.finalizedToolCallIds.clear(); + this.traceState.sessionParentIds.clear(); + } + + clearSessionTraceState(sessionID: string) { + const sessionMessageIds = new Set(); + + for (const [messageID, parts] of this.traceState.assistantParts) { + if ( + Array.from(parts.values()).some((part) => part.sessionID === sessionID) + ) { + sessionMessageIds.add(messageID); + this.traceState.assistantParts.delete(messageID); + } + } + + for (const [messageID, step] of this.traceState + .activeGenerationStepsByMessageId) { + if (step.sessionID === sessionID) { + sessionMessageIds.add(messageID); + this.traceState.activeGenerationStepsByMessageId.delete(messageID); + } + } + + for (const [messageID, observation] of this.traceState + .turnObservationsByMessageId) { + if (observation.sessionID === sessionID) { + sessionMessageIds.add(messageID); + this.traceState.turnObservationsByMessageId.delete(messageID); + } + } + + for (const messageID of sessionMessageIds) { + this.traceState.tracedMessageIds.delete(messageID); + this.traceState.tracedGenerationIds.delete(messageID); + this.traceState.generationSpansByMessageId.delete(messageID); + } + + for (const [callID, messageID] of this.traceState.toolMessageIdsByCallId) { + if (sessionMessageIds.has(messageID)) { + this.traceState.toolMessageIdsByCallId.delete(callID); + } + } + + for (const reasoningID of this.traceState.tracedReasoningIds) { + if (reasoningID.startsWith(`${sessionID}:`)) { + this.traceState.tracedReasoningIds.delete(reasoningID); + } + } + + this.traceState.abortedSessions.delete(sessionID); + this.traceState.activeGenerationSteps.delete(sessionID); + this.traceState.generationParentSpans.delete(sessionID); + this.traceState.generationInputsBySession.delete(sessionID); + this.traceState.toolResultSourceMessageIdsBySession.delete(sessionID); + this.traceState.latestTurnObservationsBySession.delete(sessionID); } endActiveToolObservations(sessionID?: string, error?: SessionErrorInfo) { @@ -110,15 +165,47 @@ export class LangfuseClient { } } - endActiveTurnObservations() { - for (const observation of new Set( - this.traceState.latestTurnObservationsBySession.values(), - )) { + endActiveTurnObservations(sessionID?: string) { + const observations = new Set([ + ...this.traceState.latestTurnObservationsBySession.values(), + ...this.traceState.turnObservationsByMessageId.values(), + ]); + + for (const observation of observations) { + if (sessionID && observation.sessionID !== sessionID) { + continue; + } + observation.span.end(); } - this.traceState.turnObservationsByMessageId.clear(); - this.traceState.latestTurnObservationsBySession.clear(); + for (const [messageID, observation] of this.traceState + .turnObservationsByMessageId) { + if (!sessionID || observation.sessionID === sessionID) { + this.traceState.turnObservationsByMessageId.delete(messageID); + } + } + + for (const [activeSessionID, observation] of this.traceState + .latestTurnObservationsBySession) { + if (!sessionID || observation.sessionID === sessionID) { + this.traceState.latestTurnObservationsBySession.delete(activeSessionID); + } + } + } + + rememberSessionParent(input: { + sessionID: string; + parentSessionID?: string; + }) { + if (input.parentSessionID) { + this.traceState.sessionParentIds.set( + input.sessionID, + input.parentSessionID, + ); + } else { + this.traceState.sessionParentIds.delete(input.sessionID); + } } traceEvent(input: { @@ -461,43 +548,15 @@ export class LangfuseClient { this.traceState.generationParentSpans.delete(input.sessionID); - const span = this.traceState.tracer.startSpan("opencode.turn", { - attributes: { - "langfuse.observation.type": "agent", - "langfuse.internal.is_app_root": true, - "session.id": input.sessionID, - "langfuse.observation.input": JSON.stringify([formattedMessage]), - "langfuse.observation.metadata": JSON.stringify({ - messageID: input.messageID, - agent: input.agent, - providerID: input.model?.providerID, - modelID: input.model?.modelID, - }), - }, - }); - - const observation = { - span, - sessionID: input.sessionID, - messageID: input.messageID, - } satisfies TurnObservation; - - if (input.messageID) { - this.traceState.turnObservationsByMessageId.set( - input.messageID, - observation, - ); - } - - this.traceState.latestTurnObservationsBySession.set( + const parentSessionID = this.traceState.sessionParentIds.get( input.sessionID, - observation, ); - - context.with(trace.setSpan(context.active(), span), () => { - const event = this.traceState.tracer.startSpan("opencode.message.user", { + const parentSpan = this.getSessionParentSpan(input.sessionID); + const startTurn = () => { + const span = this.traceState.tracer.startSpan("opencode.turn", { attributes: { - "langfuse.observation.type": "event", + "langfuse.observation.type": "agent", + "langfuse.internal.is_app_root": !parentSpan, "session.id": input.sessionID, "langfuse.observation.input": JSON.stringify([formattedMessage]), "langfuse.observation.metadata": JSON.stringify({ @@ -505,12 +564,59 @@ export class LangfuseClient { agent: input.agent, providerID: input.model?.providerID, modelID: input.model?.modelID, + parentSessionID, }), }, }); - event.end(); - }); + if (parentSpan) { + span.setAttribute("langfuse.internal.is_app_root", false); + } + + const observation = { + span, + sessionID: input.sessionID, + messageID: input.messageID, + } satisfies TurnObservation; + + if (input.messageID) { + this.traceState.turnObservationsByMessageId.set( + input.messageID, + observation, + ); + } + + this.traceState.latestTurnObservationsBySession.set( + input.sessionID, + observation, + ); + + context.with(trace.setSpan(context.active(), span), () => { + const event = this.traceState.tracer.startSpan( + "opencode.message.user", + { + attributes: { + "langfuse.observation.type": "event", + "session.id": input.sessionID, + "langfuse.observation.input": JSON.stringify([formattedMessage]), + "langfuse.observation.metadata": JSON.stringify({ + messageID: input.messageID, + agent: input.agent, + providerID: input.model?.providerID, + modelID: input.model?.modelID, + parentSessionID, + }), + }, + }, + ); + + event.end(); + }); + }; + + parentSpan + ? context.with(trace.setSpan(context.active(), parentSpan), startTurn) + : startTurn(); } rememberAssistantPart(part: MessagePart) { @@ -1054,6 +1160,20 @@ export class LangfuseClient { ); } + private getSessionParentSpan(sessionID: string) { + const parentSessionID = this.traceState.sessionParentIds.get(sessionID); + + if (!parentSessionID) { + return undefined; + } + + return ( + this.traceState.activeGenerationSteps.get(parentSessionID)?.span ?? + this.traceState.generationParentSpans.get(parentSessionID) ?? + this.traceState.latestTurnObservationsBySession.get(parentSessionID)?.span + ); + } + private withObservationParent( sessionID: string, fn: () => T, @@ -1200,6 +1320,7 @@ export type LangfuseTraceState = { generationParentSpans: Map; generationInputsBySession: Map; toolResultSourceMessageIdsBySession: Map; + sessionParentIds: Map; }; export type MessagePart = Extract< @@ -1356,6 +1477,7 @@ export const createLangfuseClient = (input: { generationParentSpans: new Map(), generationInputsBySession: new Map(), toolResultSourceMessageIdsBySession: new Map(), + sessionParentIds: new Map(), }; const processor = new LangfuseSpanProcessor({ diff --git a/test/integration/plugin.test.ts b/test/integration/plugin.test.ts index 3203b04..3ea0508 100644 --- a/test/integration/plugin.test.ts +++ b/test/integration/plugin.test.ts @@ -193,6 +193,20 @@ const getSpan = (spans: OtlpSpan[], name: string) => { return span; }; +const getSessionSpan = (spans: OtlpSpan[], name: string, sessionID: string) => { + const span = spans.find( + (candidate) => + candidate.name === name && + getAttributes(candidate)["session.id"] === sessionID, + ); + + if (!span) { + throw new Error(`Expected a ${name} span for session ${sessionID}`); + } + + return span; +}; + const emitEvent = async (event: TestPluginEvent) => { // These events are emitted by current OpenCode versions but are not yet part of // the PluginEvent union exported by our pinned @opencode-ai/plugin version. @@ -899,6 +913,102 @@ describe.sequential("built plugin", () => { ).toHaveLength(0); }); + test("links child agent sessions to the parent trace", async () => { + const parentSessionID = "parent-agent-session"; + const childSessionID = "child-agent-session"; + const parentUserMessageID = "parent-agent-user"; + const parentAssistantMessageID = "parent-agent-assistant"; + const childUserMessageID = "child-agent-user"; + const childAssistantMessageID = "child-agent-assistant"; + const started = startedAt; + + await sendUserMessage({ + sessionID: parentSessionID, + messageID: parentUserMessageID, + text: "Delegate this investigation", + started, + }); + await startGeneration({ + id: "parent-agent-step", + sessionID: parentSessionID, + started: started + 100, + }); + await emitEvent({ + type: "session.created", + properties: { + info: { + id: childSessionID, + projectID: "test-project", + directory: "/test", + parentID: parentSessionID, + title: "Child investigation", + version: "1.15.5", + time: { created: started + 150, updated: started + 150 }, + }, + }, + }); + + await sendUserMessage({ + sessionID: childSessionID, + messageID: childUserMessageID, + text: "Inspect the child scope", + started: started + 200, + }); + await startGeneration({ + id: "child-agent-step", + sessionID: childSessionID, + started: started + 300, + }); + await completeGeneration({ + sessionID: childSessionID, + userMessageID: childUserMessageID, + assistantMessageID: childAssistantMessageID, + started: started + 300, + completed: started + 600, + text: "Child result", + }); + const childFlush = await flushSession(childSessionID); + + await completeGeneration({ + sessionID: parentSessionID, + userMessageID: parentUserMessageID, + assistantMessageID: parentAssistantMessageID, + started: started + 100, + completed: started + 900, + text: "Parent result", + }); + const parentFlush = await flushSession(parentSessionID); + const spans = [...childFlush.spans, ...parentFlush.spans]; + + const parentGeneration = getSessionSpan( + spans, + "opencode.generation", + parentSessionID, + ); + const childTurn = getSessionSpan(spans, "opencode.turn", childSessionID); + const childGeneration = getSessionSpan( + spans, + "opencode.generation", + childSessionID, + ); + + expect(childTurn.traceId).toBe(parentGeneration.traceId); + expect(childTurn.parentSpanId).toBe(parentGeneration.spanId); + expect(getAttributes(childTurn)).toMatchObject({ + "langfuse.internal.is_app_root": false, + }); + expect( + getJsonAttribute(childTurn, "langfuse.observation.metadata"), + ).toMatchObject({ + parentSessionID, + }); + expect(childGeneration.traceId).toBe(parentGeneration.traceId); + expect(childGeneration.parentSpanId).toBe(childTurn.spanId); + expect( + getJsonAttribute(parentGeneration, "langfuse.observation.output"), + ).toEqual([{ role: "assistant", content: "Parent result" }]); + }); + 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";