diff --git a/src/v2.ts b/src/v2.ts index 99d4d92..c16b683 100644 --- a/src/v2.ts +++ b/src/v2.ts @@ -5,6 +5,32 @@ import { Effect } from "effect"; import type { ToolDefinition } from "./langfuse.js"; import { createLangfuseRuntime } from "./runtime.js"; +function formatMessages(messages: SessionHooks["context"]["messages"]) { + return messages.map((message) => ({ + role: message.role, + content: message.content.map((part) => { + if (part.type === "tool-call") { + const { id, name, ...rest } = part; + return { ...rest, toolCallId: id, toolName: name }; + } + if (part.type === "tool-result") { + const { id, name, result, ...rest } = part; + return { + ...rest, + toolCallId: id, + toolName: name, + output: + result.type === "error" + ? { ...result, type: "error-text" } + : result, + ...(result.type === "error" ? { isError: true } : {}), + }; + } + return part; + }), + })); +} + const LangfusePlugin = { id: "langfuse.observability", async setup(ctx) { @@ -33,6 +59,7 @@ const LangfusePlugin = { model: { id: string; providerID: string; variant?: string }; started: number; content: SessionHooks["context"]["messages"][number]["content"][number][]; + toolNames: Map; } >(); @@ -69,7 +96,7 @@ const LangfusePlugin = { ); langfuse.setGenerationInputSnapshot(input.sessionID, { system: input.system, - messages: input.messages, + messages: formatMessages(input.messages), tools, }); }), @@ -89,6 +116,8 @@ const LangfusePlugin = { id: input.id, name: input.tool, input: input.input, + // Replayed calls carry false; omission changes transcript message identity. + providerExecuted: false, }); } langfuse.rememberToolCall({ @@ -165,6 +194,7 @@ const LangfusePlugin = { model: event.data.model, started: event.created, content: [], + toolNames: new Map(), }); langfuse.startActiveGenerationStep({ sessionID: event.data.sessionID, @@ -188,6 +218,38 @@ const LangfusePlugin = { ?.content.push({ type: "reasoning", text: event.data.text }); } + if (event.type === "session.tool.input.started") { + generationDetails + .get(event.data.assistantMessageID) + ?.toolNames.set(event.data.id, event.data.name); + } + + if (event.type === "session.tool.called") { + const details = generationDetails.get( + event.data.assistantMessageID, + ); + const name = details?.toolNames.get(event.data.id); + if (details && name !== undefined) { + const call = { + type: "tool-call" as const, + id: event.data.id, + name, + input: event.data.input, + providerExecuted: event.data.executed, + }; + const index = details.content.findIndex( + (part) => + part.type === "tool-call" && part.id === event.data.id, + ); + // Execution hooks can run after the step has already been finalized. + if (index < 0) { + details.content.push(call); + } else { + details.content[index] = call; + } + } + } + if (event.type === "session.step.ended") { const details = generationDetails.get( event.data.assistantMessageID, @@ -209,12 +271,12 @@ const LangfusePlugin = { finish: event.data.finish, cost: event.data.cost, tokens: event.data.tokens, - output: [ + output: formatMessages([ { role: "assistant", content: details.content, }, - ], + ]), }); generationDetails.delete(event.data.assistantMessageID); } diff --git a/test/integration/v2-exports.test.ts b/test/integration/v2-exports.test.ts index 3578903..7f4f871 100644 --- a/test/integration/v2-exports.test.ts +++ b/test/integration/v2-exports.test.ts @@ -61,7 +61,7 @@ const getJsonAttribute = (span: typeof SpanSchema.Type, key: string) => { return Schema.decodeUnknownSync(Schema.parseJson(Schema.Unknown))(value); }; -test("exports V2 parallel tool responses with explicit IDs and no duplicate context results", async () => { +test("keeps V2 tool calls and observations with their calling generations despite late and duplicate hooks", async () => { const spans: (typeof SpanSchema.Type)[] = []; const collectorErrors: unknown[] = []; const server = createServer((request, response) => { @@ -159,6 +159,9 @@ test("exports V2 parallel tool responses with explicit IDs and no duplicate cont Schema.decodeUnknownSync(SessionIDSchema)("v2-parallel-tools"); const assistantMessageID = Schema.decodeUnknownSync(MessageIDSchema)("v2-tool-generation"); + const followupMessageID = + Schema.decodeUnknownSync(MessageIDSchema)("v2-followup"); + const followupCallID = CallID.make("v2-followup-call"); const agent = Agent.ID.make("build"); const model = { id: Model.ID.make("test-model"), @@ -212,6 +215,7 @@ test("exports V2 parallel tool responses with explicit IDs and no duplicate cont id: call.id, name: "read", input: { path: call.path }, + providerExecuted: false, })), ], }, @@ -238,6 +242,37 @@ test("exports V2 parallel tool responses with explicit IDs and no duplicate cont ], })), ] satisfies SessionHookInputs["context"]["messages"]; + const exportedMessages = [ + messages[0], + { + role: "assistant", + content: [ + ...messages[1].content.slice(0, 3), + ...calls.map((call) => ({ + type: "tool-call", + toolCallId: call.id, + toolName: "read", + input: { path: call.path }, + providerExecuted: false, + })), + ], + }, + ...calls.map((call) => ({ + role: "tool", + content: [ + { + type: "tool-result", + toolCallId: call.id, + toolName: "read", + output: { + type: call.status === "error" ? "error-text" : "text", + value: call.content, + }, + ...(call.status === "error" ? { isError: true } : {}), + }, + ], + })), + ]; const unexpectedApi = () => { throw new Error("The plugin used an unsupported mock host capability"); }; @@ -312,6 +347,42 @@ test("exports V2 parallel tool responses with explicit IDs and no duplicate cont }; } for (const call of calls) { + yield { + id: `v2-input-started-${call.id}`, + durable: { + aggregateID: sessionID, + seq: 4, + version: 1 as const, + }, + type: "session.tool.input.started", + created: Date.now(), + data: { + sessionID, + assistantMessageID, + id: call.id, + name: "read", + }, + }; + yield { + id: `v2-called-${call.id}`, + durable: { + aggregateID: sessionID, + seq: 5, + version: 1 as const, + }, + type: "session.tool.called", + created: Date.now(), + data: { + sessionID, + assistantMessageID, + id: call.id, + input: { path: call.path }, + executed: false, + }, + }; + if (call.status === "error") { + continue; + } await toolHooks.run("execute.before", { sessionID, messageID: assistantMessageID, @@ -323,6 +394,9 @@ test("exports V2 parallel tool responses with explicit IDs and no duplicate cont } // Same names and reversed completions must still match by call ID. for (const call of [...calls].reverse()) { + if (call.status === "error") { + continue; + } const input = { sessionID, messageID: assistantMessageID, @@ -331,30 +405,19 @@ test("exports V2 parallel tool responses with explicit IDs and no duplicate cont id: call.id, input: { path: call.path }, }; - const resultsByStatus = { - error: { - ...input, - status: "error" as const, - error: new ToolError({ message: call.content }), - }, - completed: { - ...input, - status: "completed" as const, - result: { - content: - call.id === "v2-call-second" - ? [ - { type: "text" as const, text: "Second" }, - { type: "text" as const, text: "response" }, - ] - : call.content, - }, + const result = { + ...input, + status: "completed", + result: { + content: + call.id === "v2-call-second" + ? [ + { type: "text", text: "Second" }, + { type: "text", text: "response" }, + ] + : call.content, }, - } satisfies Record< - (typeof calls)[number]["status"], - Parameters[0] - >; - const result = resultsByStatus[call.status]; + } satisfies Parameters[0]; await toolHooks.run("execute.after", result); await toolHooks.run("execute.after", result); } @@ -387,11 +450,81 @@ test("exports V2 parallel tool responses with explicit IDs and no duplicate cont created: Date.now(), data: { sessionID, - assistantMessageID: "v2-followup", + assistantMessageID: followupMessageID, agent, model, }, }; + // A late first-generation hook must not attach to the active followup. + for (const call of calls.filter( + (call) => call.status === "error", + )) { + const input = { + sessionID, + messageID: assistantMessageID, + agent, + tool: "read", + id: call.id, + input: { path: call.path }, + }; + await toolHooks.run("execute.before", input); + await toolHooks.run("execute.before", input); + const result = { + ...input, + status: "error" as const, + error: new ToolError({ message: call.content }), + }; + await toolHooks.run("execute.after", result); + await toolHooks.run("execute.after", result); + yield { + id: `v2-late-called-${call.id}`, + durable: { + aggregateID: sessionID, + seq: 6, + version: 1 as const, + }, + type: "session.tool.called", + created: Date.now(), + data: { + sessionID, + assistantMessageID, + id: call.id, + input: { path: call.path }, + executed: false, + }, + }; + } + yield { + id: "v2-followup-input-started", + durable: { aggregateID: sessionID, seq: 7, version: 1 as const }, + type: "session.tool.input.started", + created: Date.now(), + data: { + sessionID, + assistantMessageID: followupMessageID, + id: followupCallID, + name: "bash", + }, + }; + for (let repeat = 0; repeat < 2; repeat++) { + yield { + id: "v2-followup-called", + durable: { + aggregateID: sessionID, + seq: 8, + version: 1 as const, + }, + type: "session.tool.called", + created: Date.now(), + data: { + sessionID, + assistantMessageID: followupMessageID, + id: followupCallID, + input: { command: "pwd" }, + executed: false, + }, + }; + } yield { id: "v2-followup-ended", durable: { aggregateID: sessionID, seq: 4, version: 1 as const }, @@ -399,12 +532,29 @@ test("exports V2 parallel tool responses with explicit IDs and no duplicate cont created: Date.now(), data: { sessionID, - assistantMessageID: "v2-followup", - finish: "stop", + assistantMessageID: followupMessageID, + finish: "tool-calls", cost: 0, tokens, }, }; + const followupToolInput = { + sessionID, + messageID: followupMessageID, + agent, + tool: "bash", + id: followupCallID, + input: { command: "pwd" }, + }; + await toolHooks.run("execute.before", followupToolInput); + await toolHooks.run("execute.before", followupToolInput); + const followupToolResult = { + ...followupToolInput, + status: "completed" as const, + result: { content: "/project" }, + }; + await toolHooks.run("execute.after", followupToolResult); + await toolHooks.run("execute.after", followupToolResult); yield { id: "v2-execution-succeeded", durable: { aggregateID: sessionID, seq: 5, version: 1 as const }, @@ -497,16 +647,50 @@ test("exports V2 parallel tool responses with explicit IDs and no duplicate cont (span) => span.name === "opencode.generation", ); expect(generations).toHaveLength(2); - const generation = generations.find( - (span) => span.spanId === tools[0].parentSpanId, - ); + const generation = generations.find((span) => { + const metadata = getJsonAttribute(span, "langfuse.observation.metadata"); + return ( + typeof metadata === "object" && + metadata !== null && + "messageID" in metadata && + metadata.messageID === assistantMessageID + ); + }); const followup = generations.find((span) => span !== generation); if (!generation || !followup) { throw new Error("Expected both generations"); } expect(getJsonAttribute(generation, "langfuse.observation.output")).toEqual([ - messages[1], + exportedMessages[1], + ]); + expect( + getJsonAttribute(followup, "langfuse.observation.metadata"), + ).toMatchObject({ messageID: followupMessageID }); + expect(getJsonAttribute(followup, "langfuse.observation.output")).toEqual([ + { + role: "assistant", + content: [ + { + type: "tool-call", + toolCallId: followupCallID, + toolName: "bash", + input: { command: "pwd" }, + providerExecuted: false, + }, + ], + }, ]); + const followupTools = spans.filter((span) => span.name === "bash"); + expect(followupTools).toHaveLength(1); + expect(followupTools[0].parentSpanId).toBe(followup.spanId); + expect( + getJsonAttribute(followupTools[0], "langfuse.observation.output"), + ).toEqual({ + role: "tool", + name: "bash", + tool_call_id: followupCallID, + content: "/project", + }); for (const [index, call] of calls.entries()) { const tool = tools.find((span) => { const output = getJsonAttribute(span, "langfuse.observation.output"); @@ -534,7 +718,7 @@ test("exports V2 parallel tool responses with explicit IDs and no duplicate cont } expect(getJsonAttribute(followup, "langfuse.observation.input")).toEqual({ system: [], - messages, + messages: exportedMessages, tools: [], }); }); diff --git a/test/integration/v2.test.ts b/test/integration/v2.test.ts index 549919e..f31cfdb 100644 --- a/test/integration/v2.test.ts +++ b/test/integration/v2.test.ts @@ -314,9 +314,10 @@ describe("OpenCode 2 package entrypoint", () => { { type: "reasoning", text: "Thinking" }, { type: "tool-call", - id: "call-1", - name: "read", + toolCallId: "call-1", + toolName: "read", input: { path: "README.md" }, + providerExecuted: false, }, ], }), @@ -327,6 +328,122 @@ describe("OpenCode 2 package entrypoint", () => { await cleanup?.(); }); + test.each([false, true])( + "captures tool output before step completion with providerExecuted=%s and late execution hooks", + async (executed) => { + let executeBefore: + | ((input: { + id: string; + messageID: string; + sessionID: string; + tool: string; + input: Record; + }) => void) + | undefined; + const registration = { dispose: vi.fn(() => Promise.resolve()) }; + const contextInput: unknown = { + app: { version: "2.0.4" }, + session: { hook: vi.fn(() => Promise.resolve(registration)) }, + tool: { + hook: vi.fn((name: string, handler: typeof executeBefore) => { + if (name === "execute.before") { + executeBefore = handler; + } + return Promise.resolve(registration); + }), + }, + event: { + 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" }, + }, + }; + yield { + type: "session.tool.input.started", + data: { + assistantMessageID: "assistant-1", + id: "call-1", + name: "read", + }, + }; + for (let repeat = 0; repeat < 2; repeat++) { + yield { + type: "session.tool.called", + data: { + assistantMessageID: "assistant-1", + id: "call-1", + input: { path: "README.md" }, + executed, + }, + }; + } + yield { + type: "session.step.ended", + created: 200, + data: { + sessionID: "session-1", + assistantMessageID: "assistant-1", + finish: "tool-calls", + cost: 0, + tokens: { + input: 10, + output: 5, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + }, + }; + if (!executed) { + executeBefore?.({ + sessionID: "session-1", + messageID: "assistant-1", + id: "call-1", + tool: "read", + input: { path: "README.md" }, + }); + } + }, + }), + }, + }; + const context = Schema.decodeUnknownSync( + Schema.declare( + (input): input is Parameters[0] => + typeof input === "object" && input !== null, + ), + )(contextInput); + const cleanup = await SourcePlugin.setup(context); + await cleanup?.(); + expect(runtime.traceGeneration).toHaveBeenCalledTimes(1); + expect(runtime.traceGeneration).toHaveBeenCalledWith( + expect.objectContaining({ + output: [ + { + role: "assistant", + content: [ + { + type: "tool-call", + toolCallId: "call-1", + toolName: "read", + input: { path: "README.md" }, + providerExecuted: executed, + }, + ], + }, + ], + }), + ); + }, + ); + test("captures the complete model input from each context hook", async () => { let context: | ((input: {