diff --git a/src/langfuse.ts b/src/langfuse.ts index 641cb77..e1fc62f 100644 --- a/src/langfuse.ts +++ b/src/langfuse.ts @@ -42,6 +42,7 @@ export class LangfuseClient { this.traceState.toolMessageIdsByCallId.clear(); this.traceState.generationParentSpans.clear(); this.traceState.generationInputsBySession.clear(); + this.traceState.generationInputSnapshotsBySession.clear(); this.traceState.toolResultSourceMessageIdsBySession.clear(); this.traceState.turnObservationsByMessageId.clear(); this.traceState.latestTurnObservationsBySession.clear(); @@ -101,6 +102,7 @@ export class LangfuseClient { this.traceState.activeGenerationSteps.delete(sessionID); this.traceState.generationParentSpans.delete(sessionID); this.traceState.generationInputsBySession.delete(sessionID); + this.traceState.generationInputSnapshotsBySession.delete(sessionID); this.traceState.toolResultSourceMessageIdsBySession.delete(sessionID); this.traceState.latestTurnObservationsBySession.delete(sessionID); this.traceState.sessionHistories.delete(sessionID); @@ -511,6 +513,10 @@ export class LangfuseClient { ); } + setGenerationInputSnapshot(sessionID: string, input: unknown) { + this.traceState.generationInputSnapshotsBySession.set(sessionID, input); + } + private traceFormattedUserMessage(input: { sessionID: string; messageID?: string; @@ -1260,6 +1266,15 @@ export class LangfuseClient { sessionID: string, assistantMessageID?: string, ) { + const snapshot = + this.traceState.generationInputSnapshotsBySession.get(sessionID); + this.traceState.generationInputSnapshotsBySession.delete(sessionID); + if (snapshot !== undefined) { + this.traceState.generationInputsBySession.delete(sessionID); + this.traceState.toolResultSourceMessageIdsBySession.delete(sessionID); + return snapshot; + } + const pending = this.traceState.generationInputsBySession.get(sessionID); const sourceMessageID = this.traceState.toolResultSourceMessageIdsBySession.get(sessionID); @@ -1382,6 +1397,7 @@ export type LangfuseTraceState = { activeGenerationSteps: Map; generationParentSpans: Map; generationInputsBySession: Map; + generationInputSnapshotsBySession: Map; toolResultSourceMessageIdsBySession: Map; sessionParentIds: Map; sessionHistories: Map; @@ -1606,7 +1622,9 @@ type SessionError = Extract< { type: "session.error" } >["properties"]["error"]; -export type SessionErrorInfo = NonNullable; +export type SessionErrorInfo = + | NonNullable + | { name: string; message?: string; data?: { message?: unknown } }; type UserMessageInput = { role: "user"; @@ -1740,6 +1758,7 @@ export const createLangfuseClient = (input: { activeGenerationSteps: new Map(), generationParentSpans: new Map(), generationInputsBySession: new Map(), + generationInputSnapshotsBySession: new Map(), toolResultSourceMessageIdsBySession: new Map(), sessionParentIds: new Map(), sessionHistories: new Map(), diff --git a/src/v2.ts b/src/v2.ts index b7ede1f..27b7f49 100644 --- a/src/v2.ts +++ b/src/v2.ts @@ -27,6 +27,7 @@ const LangfusePlugin = { const generationDetails = new Map< string, { + sessionID: string; agent: string; model: { id: string; providerID: string; variant?: string }; started: number; @@ -67,7 +68,11 @@ const LangfusePlugin = { parameters: tool.input, }), ); - langfuse.setPendingToolDefinitions(input.sessionID, tools); + langfuse.setGenerationInputSnapshot(input.sessionID, { + system: input.system, + messages: input.messages, + tools, + }); }), ); @@ -147,6 +152,7 @@ const LangfusePlugin = { if (event.type === "session.step.started") { generationDetails.set(event.data.assistantMessageID, { + sessionID: event.data.sessionID, agent: event.data.agent, model: event.data.model, started: event.created, @@ -240,14 +246,34 @@ const LangfusePlugin = { }); } + if (event.type === "session.execution.failed") { + langfuse.traceSessionError({ + sessionID: event.data.sessionID, + error: { + name: event.data.error.type, + message: event.data.error.message, + }, + }); + for (const [messageID, details] of generationDetails) { + if (details.sessionID === event.data.sessionID) { + generationDetails.delete(messageID); + } + } + await Effect.runPromise(langfuse.forceFlush); + } + if ( event.type === "session.execution.succeeded" || - event.type === "session.execution.failed" || event.type === "session.execution.interrupted" ) { langfuse.endActiveToolObservations(event.data.sessionID); langfuse.endActiveGenerationSteps(event.data.sessionID); langfuse.endActiveTurnObservations(event.data.sessionID); + for (const [messageID, details] of generationDetails) { + if (details.sessionID === event.data.sessionID) { + generationDetails.delete(messageID); + } + } await Effect.runPromise(langfuse.forceFlush); } diff --git a/test/integration/v2.test.ts b/test/integration/v2.test.ts index e60f68c..89e9b72 100644 --- a/test/integration/v2.test.ts +++ b/test/integration/v2.test.ts @@ -10,7 +10,7 @@ import SourcePlugin from "../../src/v2.js"; const runtime = vi.hoisted(() => ({ createLangfuseRuntime: vi.fn(), traceUserPrompt: vi.fn(), - setPendingToolDefinitions: vi.fn(), + setGenerationInputSnapshot: vi.fn(), rememberToolCall: vi.fn(), traceToolStart: vi.fn(), traceToolError: vi.fn(), @@ -19,6 +19,7 @@ const runtime = vi.hoisted(() => ({ startActiveGenerationStep: vi.fn(), traceGeneration: vi.fn(), traceFailedGenerationStep: vi.fn(), + traceSessionError: vi.fn(), traceEvent: vi.fn(), endActiveToolObservations: vi.fn(), endActiveGenerationSteps: vi.fn(), @@ -79,6 +80,17 @@ describe("OpenCode 2 package entrypoint", () => { subscribe: () => ({ async *[Symbol.asyncIterator]() { await Promise.resolve(); + yield { + type: "session.step.started", + created: 100, + data: { + sessionID: "session-1", + assistantMessageID: "assistant-1", + agent: "build", + model: { id: "model-1", providerID: "provider-1" }, + snapshot: "snapshot-1", + }, + }; yield { type: "session.execution.failed", data: { @@ -86,6 +98,22 @@ describe("OpenCode 2 package entrypoint", () => { error: { type: "TestError", message: "failed" }, }, }; + yield { + type: "session.step.ended", + created: 200, + data: { + sessionID: "session-1", + assistantMessageID: "assistant-1", + finish: "stop", + cost: 0, + tokens: { + input: 0, + output: 0, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + }, + }; }, }), }, @@ -104,9 +132,11 @@ describe("OpenCode 2 package entrypoint", () => { expect(runtime.createLangfuseRuntime).toHaveBeenCalledWith({ opencodeVersion: "2.0.4", }); - expect(runtime.endActiveToolObservations).toHaveBeenCalledWith("session-1"); - expect(runtime.endActiveGenerationSteps).toHaveBeenCalledWith("session-1"); - expect(runtime.endActiveTurnObservations).toHaveBeenCalledWith("session-1"); + expect(runtime.traceSessionError).toHaveBeenCalledWith({ + sessionID: "session-1", + error: { name: "TestError", message: "failed" }, + }); + expect(runtime.traceGeneration).not.toHaveBeenCalled(); }); test("traces a complete session with prompt, text, reasoning, and tools", async () => { @@ -262,4 +292,84 @@ describe("OpenCode 2 package entrypoint", () => { }); await cleanup?.(); }); + + test("captures the complete model input from each context hook", async () => { + let context: + | ((input: { + sessionID: string; + system: unknown[]; + messages: unknown[]; + tools: Record< + string, + { description: string; input: Record } + >; + }) => void) + | undefined; + const registration = { dispose: vi.fn(() => Promise.resolve()) }; + const contextInput: unknown = { + app: { version: "2.0.4" }, + session: { + hook: vi.fn((name: string, handler: typeof context) => { + if (name === "context") { + context = handler; + } + return Promise.resolve(registration); + }), + }, + tool: { hook: vi.fn(() => Promise.resolve(registration)) }, + event: { + subscribe: () => ({ + async *[Symbol.asyncIterator]() { + await Promise.resolve(); + yield* []; + }, + }), + }, + }; + const pluginContext = Schema.decodeUnknownSync( + Schema.declare( + (input): input is Parameters[0] => + typeof input === "object" && input !== null, + ), + )(contextInput); + + const cleanup = await SourcePlugin.setup(pluginContext); + await vi.waitFor(() => { + expect(context).toBeTypeOf("function"); + }); + + const system = [{ type: "text", text: "System instructions" }]; + const messages = [ + { role: "user", content: [{ type: "text", text: "Earlier message" }] }, + { role: "assistant", content: [{ type: "text", text: "Earlier reply" }] }, + ]; + context?.({ + sessionID: "session-1", + system, + messages, + tools: { + read: { + description: "Read a file", + input: { type: "object", properties: {} }, + }, + }, + }); + + expect(runtime.setGenerationInputSnapshot).toHaveBeenCalledWith( + "session-1", + { + system, + messages, + tools: [ + { + name: "read", + description: "Read a file", + parameters: { type: "object", properties: {} }, + }, + ], + }, + ); + + await cleanup?.(); + }); });