Fix incorrect latency after tool timeout (#15)

* Add timeout regression test

* End timed out tool spans

Author: bezbac

src/index.ts

@@ -166,8 +166,18 @@ const eventHook = (event: OpencodeEvent, shutdown?: () => Promise<void>) =>
     }
 
     if (event.type === "message.part.updated") {
-      langfuse.rememberAssistantPart(event.properties.part);
-      langfuse.traceReasoningPart(event.properties.part);
+      const part = event.properties.part;
+
+      langfuse.rememberAssistantPart(part);
+      langfuse.traceReasoningPart(part);
+
+      if (part.type === "tool" && part.state.status === "error") {
+        langfuse.traceToolError({
+          callID: part.callID,
+          error: part.state.error,
+          completed: part.state.time.end,
+        });
+      }
     }
 
     if (event.type === "session.next.step.started") {

src/langfuse.ts

@@ -40,6 +40,7 @@ export class LangfuseClient {
     this.traceState.generationParentSpans.clear();
     this.traceState.turnObservationsByMessageId.clear();
     this.traceState.latestTurnObservationsBySession.clear();
+    this.traceState.finalizedToolCallIds.clear();
   }
 
   endActiveToolObservations(sessionID?: string, error?: SessionErrorInfo) {
@@ -61,6 +62,7 @@ export class LangfuseClient {
 
       observation.span.end();
       this.traceState.activeToolObservations.delete(callID);
+      this.traceState.finalizedToolCallIds.add(callID);
     }
   }
 
@@ -707,6 +709,7 @@ export class LangfuseClient {
     args: unknown;
   }) {
     this.traceState.activeToolObservations.get(input.callID)?.span.end();
+    this.traceState.finalizedToolCallIds.delete(input.callID);
     this.ensureGenerationParent(input.sessionID);
 
     this.withObservationParent(input.sessionID, () => {
@@ -738,6 +741,10 @@ export class LangfuseClient {
     title: string;
     output: string;
   }) {
+    if (this.traceState.finalizedToolCallIds.has(input.callID)) {
+      return;
+    }
+
     if (!this.traceState.activeToolObservations.has(input.callID)) {
       this.traceToolStart({
         sessionID: input.sessionID,
@@ -767,6 +774,32 @@ export class LangfuseClient {
 
     span.end();
     this.traceState.activeToolObservations.delete(input.callID);
+    this.traceState.finalizedToolCallIds.add(input.callID);
+  }
+
+  traceToolError(input: { callID: string; error: string; completed: number }) {
+    if (this.traceState.finalizedToolCallIds.has(input.callID)) {
+      return;
+    }
+
+    const span = this.traceState.activeToolObservations.get(input.callID)?.span;
+
+    if (!span) {
+      return;
+    }
+
+    span.setAttribute(
+      "langfuse.observation.output",
+      JSON.stringify({ error: input.error }),
+    );
+    span.setStatus({
+      code: SpanStatusCode.ERROR,
+      message: input.error,
+    });
+    span.recordException({ message: input.error });
+    span.end(new Date(input.completed));
+    this.traceState.activeToolObservations.delete(input.callID);
+    this.traceState.finalizedToolCallIds.add(input.callID);
   }
 
   private ensureGenerationParent(sessionID: string) {
@@ -873,6 +906,7 @@ export type LangfuseTraceState = {
   turnObservationsByMessageId: Map<string, TurnObservation>;
   latestTurnObservationsBySession: Map<string, TurnObservation>;
   activeToolObservations: Map<string, ToolObservation>;
+  finalizedToolCallIds: Set<string>;
   activeGenerationSteps: Map<string, ActiveGenerationStep>;
   generationParentSpans: Map<string, ApiSpan>;
 };
@@ -1029,6 +1063,7 @@ export const createLangfuseClient = (input: {
       turnObservationsByMessageId: new Map<string, TurnObservation>(),
       latestTurnObservationsBySession: new Map<string, TurnObservation>(),
       activeToolObservations: new Map<string, ToolObservation>(),
+      finalizedToolCallIds: new Set<string>(),
       activeGenerationSteps: new Map<string, ActiveGenerationStep>(),
       generationParentSpans: new Map<string, ApiSpan>(),
     };

test/integration/plugin.test.ts

@@ -24,6 +24,8 @@ interface OtlpSpan {
   traceId: string;
   spanId: string;
   parentSpanId?: string;
+  startTimeUnixNano: string;
+  endTimeUnixNano: string;
   attributes: Array<{ key: string; value: OtlpValue }>;
   status?: { code?: number; message?: string };
   events?: Array<{ name: string }>;
@@ -586,6 +588,40 @@ describe.sequential("built plugin", () => {
       },
       { title: "README.md", output: "# Project", metadata: {} },
     );
+    const failedToolStarted = Date.now();
+    const failedToolEnded = failedToolStarted + 5_000;
+    await hooks["tool.execute.before"]?.(
+      { sessionID, callID: "timed-out-webfetch", tool: "webfetch" },
+      { args: { url: "https://example.com", timeout: 5 } },
+    );
+    await emitEvent({
+      type: "message.part.updated",
+      properties: {
+        part: {
+          id: "timed-out-webfetch-part",
+          sessionID,
+          messageID: assistantMessageID,
+          type: "tool",
+          callID: "timed-out-webfetch",
+          tool: "webfetch",
+          state: {
+            status: "error",
+            input: { url: "https://example.com", timeout: 5 },
+            error: "Tool execution timed out after 5 seconds",
+            time: { start: failedToolStarted, end: failedToolEnded },
+          },
+        },
+      },
+    });
+    await hooks["tool.execute.after"]?.(
+      {
+        sessionID,
+        callID: "timed-out-webfetch",
+        tool: "webfetch",
+        args: { url: "https://example.com", timeout: 5 },
+      },
+      { title: "Web fetch", output: "late output", metadata: {} },
+    );
     await emitEvent({
       id: "nested-observations-retry",
       type: "session.next.retried",
@@ -624,6 +660,7 @@ describe.sequential("built plugin", () => {
         "opencode.generation.reasoning",
         "opencode.generation.reasoning",
         "read",
+        "webfetch",
         "opencode.generation.retry",
         "opencode.generation.compaction",
       ].sort(),
@@ -642,6 +679,18 @@ describe.sequential("built plugin", () => {
     expect(tool.traceId).toBe(generation.traceId);
     expect(tool.parentSpanId).toBe(generation.spanId);
 
+    const failedTool = getSpan(spans, "webfetch");
+    expect(failedTool.endTimeUnixNano).toBe(
+      (BigInt(failedToolEnded) * 1_000_000n).toString(),
+    );
+    expect(failedTool.status).toEqual({
+      code: 2,
+      message: "Tool execution timed out after 5 seconds",
+    });
+    expect(getJsonAttribute(failedTool, "langfuse.observation.output")).toEqual(
+      { error: "Tool execution timed out after 5 seconds" },
+    );
+
     const retry = getSpan(spans, "opencode.generation.retry");
     expect(getJsonAttribute(retry, "langfuse.observation.metadata")).toEqual({
       attempt: 2,