diff --git a/src/langfuse.ts b/src/langfuse.ts index c843acd..3e83e6f 100644 --- a/src/langfuse.ts +++ b/src/langfuse.ts @@ -44,6 +44,8 @@ export class LangfuseClient { this.traceState.turnObservationsByMessageId.clear(); this.traceState.latestTurnObservationsBySession.clear(); this.traceState.finalizedToolCallIds.clear(); + // messageHistoriesBySession is intentionally kept: it must survive the + // session.idle flush between turns so later generations keep prior turns. } endActiveToolObservations(sessionID?: string, error?: SessionErrorInfo) { @@ -337,7 +339,7 @@ export class LangfuseClient { return; } - const generationInput = this.consumeGenerationInput(input.sessionID); + const generationInput = this.collectGenerationInput(input.sessionID); this.withTurnParent(input.sessionID, undefined, () => { const span = this.traceState.tracer.startSpan("opencode.generation", { @@ -650,6 +652,10 @@ export class LangfuseClient { this.traceState.activeGenerationSteps.delete(input.sessionID); } + if (input.mode !== "compaction") { + this.appendAssistantMessageToHistory(input.sessionID, input.messageID); + } + return; } @@ -657,7 +663,11 @@ export class LangfuseClient { return; } - const generationInput = this.consumeGenerationInput(input.sessionID); + const generationInput = this.collectGenerationInput(input.sessionID); + + if (input.mode !== "compaction") { + this.appendAssistantMessageToHistory(input.sessionID, input.messageID); + } this.withTurnParent(input.sessionID, input.parentID, () => { const span = this.traceState.tracer.startSpan("opencode.generation", { @@ -1012,7 +1022,7 @@ export class LangfuseClient { } this.withTurnParent(sessionID, undefined, () => { - const generationInput = this.consumeGenerationInput(sessionID); + const generationInput = this.collectGenerationInput(sessionID); const span = this.traceState.tracer.startSpan("opencode.generation", { attributes: { "langfuse.observation.type": "generation", @@ -1116,21 +1126,53 @@ export class LangfuseClient { ]; } - private consumeGenerationInput(sessionID: string) { - const input = this.traceState.generationInputsBySession.get(sessionID); + private collectGenerationInput(sessionID: string) { + const pending = 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; + if (sourceMessageID) { + this.appendAssistantMessageToHistory(sessionID, sourceMessageID); } - return [ - ...(this.getAssistantMessage(sourceMessageID) ?? []), - ...(input ?? []), - ]; + const history = this.getMessageHistory(sessionID); + + if (pending) { + history.push(...pending); + } + + return history.length ? [...history] : undefined; + } + + private getMessageHistory(sessionID: string) { + let history = this.traceState.messageHistoriesBySession.get(sessionID); + + if (!history) { + history = []; + this.traceState.messageHistoriesBySession.set(sessionID, history); + } + + return history; + } + + private appendAssistantMessageToHistory( + sessionID: string, + messageID: string, + ) { + if (this.traceState.historyAssistantMessageIds.has(messageID)) { + return; + } + + const message = this.getAssistantMessage(messageID); + + if (!message) { + return; + } + + this.traceState.historyAssistantMessageIds.add(messageID); + this.getMessageHistory(sessionID).push(...message); } private rememberToolResult(input: { @@ -1200,6 +1242,8 @@ export type LangfuseTraceState = { generationParentSpans: Map; generationInputsBySession: Map; toolResultSourceMessageIdsBySession: Map; + messageHistoriesBySession: Map; + historyAssistantMessageIds: Set; }; export type MessagePart = Extract< @@ -1356,6 +1400,8 @@ export const createLangfuseClient = (input: { generationParentSpans: new Map(), generationInputsBySession: new Map(), toolResultSourceMessageIdsBySession: new Map(), + messageHistoriesBySession: new Map(), + historyAssistantMessageIds: new Set(), }; const processor = new LangfuseSpanProcessor({ diff --git a/test/integration/plugin.test.ts b/test/integration/plugin.test.ts index 3203b04..d7ad49d 100644 --- a/test/integration/plugin.test.ts +++ b/test/integration/plugin.test.ts @@ -660,6 +660,103 @@ describe.sequential("built plugin", () => { variant: "high", snapshot: "snapshot-1", }); + + const secondGeneration = spans + .filter((span) => span.name === "opencode.generation") + .find((span) => { + const metadata = getJsonAttribute( + span, + "langfuse.observation.metadata", + ); + return ( + typeof metadata === "object" && + metadata !== null && + "messageID" in metadata && + metadata.messageID === secondAssistantMessageID + ); + }); + if (!secondGeneration) { + throw new Error("Expected the second generation span"); + } + + expect( + getJsonAttribute(secondGeneration, "langfuse.observation.input"), + ).toEqual([ + { + role: "user", + content: [{ type: "text", text: "Inspect the repository" }], + tools: expect.any(Array), + }, + { role: "assistant", content: "Repository inspected" }, + { + role: "user", + content: [{ type: "text", text: "Summarize it" }], + tools: expect.any(Array), + }, + ]); + }); + + test("carries prior turns into later generation inputs across idle flushes", async () => { + const sessionID = "history-session"; + const started = startedAt; + + await sendUserMessage({ + sessionID, + messageID: "history-user-1", + text: "Read the README", + started, + }); + await startGeneration({ + id: "history-step-1", + sessionID, + started: started + 100, + }); + await completeGeneration({ + sessionID, + userMessageID: "history-user-1", + assistantMessageID: "history-assistant-1", + started: started + 100, + completed: started + 500, + text: "The README describes the project", + }); + await flushSession(sessionID); + + await sendUserMessage({ + sessionID, + messageID: "history-user-2", + text: "Now summarize it", + started: started + 1_000, + }); + await startGeneration({ + id: "history-step-2", + sessionID, + started: started + 1_100, + }); + await completeGeneration({ + sessionID, + userMessageID: "history-user-2", + assistantMessageID: "history-assistant-2", + started: started + 1_100, + completed: started + 1_500, + text: "A summary", + }); + + const { spans } = await flushSession(sessionID); + const generation = getSpan(spans, "opencode.generation"); + + expect(getJsonAttribute(generation, "langfuse.observation.input")).toEqual([ + { + role: "user", + content: [{ type: "text", text: "Read the README" }], + tools: expect.any(Array), + }, + { role: "assistant", content: "The README describes the project" }, + { + role: "user", + content: [{ type: "text", text: "Now summarize it" }], + tools: expect.any(Array), + }, + ]); }); test("exports reasoning, tool, retry, and compaction spans", async () => { @@ -1043,6 +1140,11 @@ describe.sequential("built plugin", () => { expect( getJsonAttribute(generationSpans[1], "langfuse.observation.input"), ).toEqual([ + { + role: "user", + content: [{ type: "text", text: "Run three tool batches" }], + tools: expect.any(Array), + }, { role: "assistant", tool_calls: [ @@ -1202,6 +1304,11 @@ describe.sequential("built plugin", () => { expect( getJsonAttribute(secondGeneration!, "langfuse.observation.input"), ).toEqual([ + { + role: "user", + content: [{ type: "text", text: "Show recent commits" }], + tools: expect.any(Array), + }, { role: "assistant", tool_calls: [