Improve V2 generation tracing

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<string, ActiveGenerationStep>;
   generationParentSpans: Map<string, ApiSpan>;
   generationInputsBySession: Map<string, ChatMlMessage[]>;
+  generationInputSnapshotsBySession: Map<string, unknown>;
   toolResultSourceMessageIdsBySession: Map<string, string>;
   sessionParentIds: Map<string, string>;
   sessionHistories: Map<string, SessionHistory>;
@@ -1606,7 +1622,9 @@ type SessionError = Extract<
   { type: "session.error" }
 >["properties"]["error"];
 
-export type SessionErrorInfo = NonNullable<SessionError>;
+export type SessionErrorInfo =
+  | NonNullable<SessionError>
+  | { name: string; message?: string; data?: { message?: unknown } };
 
 type UserMessageInput = {
   role: "user";
@@ -1740,6 +1758,7 @@ export const createLangfuseClient = (input: {
       activeGenerationSteps: new Map<string, ActiveGenerationStep>(),
       generationParentSpans: new Map<string, ApiSpan>(),
       generationInputsBySession: new Map<string, ChatMlMessage[]>(),
+      generationInputSnapshotsBySession: new Map<string, unknown>(),
       toolResultSourceMessageIdsBySession: new Map<string, string>(),
       sessionParentIds: new Map<string, string>(),
       sessionHistories: new Map<string, SessionHistory>(),

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);
           }
 

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,13 +80,40 @@ 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: {
                 sessionID: "session-1",
                 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<string, unknown> }
+          >;
+        }) => 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<typeof SourcePlugin.setup>[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?.();
+  });
 });