pi-langfuse 1.4.2 → 1.4.3

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "pi-langfuse",
-	"version": "1.4.2",
+	"version": "1.4.3",
 	"description": "Langfuse extension for Pi coding agent",
 	"repository": {
 		"type": "git",

package/src/constants.ts

@@ -10,3 +10,4 @@ export const MAX_TOOL_PAYLOAD_LENGTH = 24_000;
 export const MAX_DEPTH = 6;
 export const MAX_ARRAY_ITEMS = 50;
 export const MAX_OBJECT_KEYS = 80;
+export const MAX_PAYLOAD_NODES = 2_000;

package/src/handlers/agent.ts

@@ -19,6 +19,7 @@ export function updateTraceIO(input?: unknown, output?: unknown) {
 
 export async function startAgentRun(event: Record<string, unknown>, ctx: any) {
   if (!(await ensureConfig(ctx))) {
+    state.isTracingDisabled = true;
     return;
   }
 
@@ -92,6 +93,7 @@ export async function startAgentRun(event: Record<string, unknown>, ctx: any) {
     updateTraceIO(promptInput, undefined);
   } catch (e) {
     console.warn("📊 Langfuse: Failed to create agent observation", e);
+    state.isTracingDisabled = true;
   }
 }
 

package/src/handlers/generation.ts

@@ -13,7 +13,7 @@ import {
 import type { GenerationState, ObservationUpdate } from "../types.js";
 
 export function getOpenGeneration(): GenerationState | undefined {
-  if (!state.agentState) {
+  if (state.isTracingDisabled || !state.agentState) {
     return undefined;
   }
 
@@ -29,7 +29,7 @@ export function getOpenGeneration(): GenerationState | undefined {
 }
 
 export async function startGeneration(event: Record<string, unknown>) {
-  if (!state.agentState?.root) {
+  if (state.isTracingDisabled || !state.agentState?.root) {
     return;
   }
 
@@ -79,7 +79,7 @@ export async function startGeneration(event: Record<string, unknown>) {
 }
 
 export function updateGenerationMetadata(event: Record<string, unknown>) {
-  if (!state.agentState) {
+  if (state.isTracingDisabled || !state.agentState) {
     return;
   }
 
@@ -132,7 +132,7 @@ export function updateGenerationMetadata(event: Record<string, unknown>) {
 }
 
 export function recordTTFT(event: Record<string, unknown>) {
-  if (!state.agentState) {
+  if (state.isTracingDisabled || !state.agentState) {
     return;
   }
 
@@ -150,7 +150,7 @@ export function recordTTFT(event: Record<string, unknown>) {
 }
 
 export async function finishGenerationFromMessage(event: Record<string, unknown>) {
-  if (!state.agentState) {
+  if (state.isTracingDisabled || !state.agentState) {
     return;
   }
 
@@ -190,7 +190,7 @@ export async function finishGenerationFromMessage(event: Record<string, unknown>
 }
 
 export async function createFallbackGenerationFromTurn(event: Record<string, unknown>, message: Record<string, unknown>) {
-  if (!state.agentState?.root || state.agentState.generationOrder.length > 0) {
+  if (state.isTracingDisabled || !state.agentState?.root || state.agentState.generationOrder.length > 0) {
     return;
   }
 

package/src/handlers/tool.ts

@@ -12,7 +12,7 @@ import {
 import { MAX_TOOL_PAYLOAD_LENGTH } from "../constants.js";
 
 export async function startToolObservation(event: Record<string, unknown>) {
-  if (!state.agentState?.root) {
+  if (state.isTracingDisabled || !state.agentState?.root) {
     return;
   }
 
@@ -59,7 +59,7 @@ export async function startToolObservation(event: Record<string, unknown>) {
 }
 
 export async function finishToolObservation(event: Record<string, unknown>) {
-  if (!state.agentState) {
+  if (state.isTracingDisabled || !state.agentState) {
     return;
   }
 
@@ -119,7 +119,7 @@ export async function finishToolObservation(event: Record<string, unknown>) {
 }
 
 export function closeDanglingObservations(statusMessage: string) {
-  if (!state.agentState) {
+  if (state.isTracingDisabled || !state.agentState) {
     return;
   }
 

package/src/handlers/turn.ts

@@ -3,7 +3,7 @@ import { getRuntime } from "../langfuse.js";
 import { shapePayload } from "../utils.js";
 
 export async function startTurnObservation(event: Record<string, unknown>) {
-  if (!state.agentState?.root) {
+  if (state.isTracingDisabled || !state.agentState?.root) {
     return;
   }
 
@@ -40,7 +40,7 @@ export async function startTurnObservation(event: Record<string, unknown>) {
 }
 
 export function finishTurnObservation(event?: Record<string, unknown>) {
-  if (!state.agentState?.activeTurn) {
+  if (state.isTracingDisabled || !state.agentState?.activeTurn) {
     return;
   }
 

package/src/langfuse.ts

@@ -42,7 +42,8 @@ interface RestFallbackStore {
   attempted: boolean;
 }
 
-const OTEL_VISIBILITY_DELAY_MS = 1_500;
+const OTEL_VISIBILITY_TIMEOUT_MS = 1_500;
+const OTEL_VISIBILITY_POLL_INTERVAL_MS = 200;
 
 function nowIso() {
   return new Date().toISOString();
@@ -190,6 +191,23 @@ async function traceExists(rt: LangfuseRuntime, traceId: string): Promise<boolea
   }
 }
 
+async function waitForTraceVisibility(rt: LangfuseRuntime, traceId: string): Promise<boolean> {
+  const deadline = Date.now() + OTEL_VISIBILITY_TIMEOUT_MS;
+
+  while (true) {
+    if (await traceExists(rt, traceId)) {
+      return true;
+    }
+
+    const remainingMs = deadline - Date.now();
+    if (remainingMs <= 0) {
+      return false;
+    }
+
+    await delay(Math.min(OTEL_VISIBILITY_POLL_INTERVAL_MS, remainingMs));
+  }
+}
+
 function eventTimestamp(record: { endTime?: string; startTime?: string; timestamp?: string }) {
   return record.endTime ?? record.startTime ?? record.timestamp ?? nowIso();
 }
@@ -201,8 +219,7 @@ async function fallbackToRestIngestion(rt: LangfuseRuntime) {
   }
   store.attempted = true;
 
-  await delay(OTEL_VISIBILITY_DELAY_MS);
-  if (await traceExists(rt, store.trace.id)) {
+  if (await waitForTraceVisibility(rt, store.trace.id)) {
     return;
   }
 

package/src/state.ts

@@ -8,6 +8,7 @@ export interface SessionRunState {
   toolCallCount: number;
   errorCount: number;
   turnCount: number;
+  tracingDisabled: boolean;
 }
 
 const DEFAULT_SESSION_ID = "__pi_langfuse_default_session__";
@@ -23,6 +24,7 @@ function createSessionRunState(): SessionRunState {
     toolCallCount: 0,
     errorCount: 0,
     turnCount: 0,
+    tracingDisabled: false,
   };
 }
 
@@ -109,6 +111,13 @@ export const state = {
   set turnCount(turnCount: number) {
     getSessionRunState().turnCount = turnCount;
   },
+
+  get isTracingDisabled() {
+    return getSessionRunState().tracingDisabled;
+  },
+  set isTracingDisabled(disabled: boolean) {
+    getSessionRunState().tracingDisabled = disabled;
+  },
 };
 
 export function resetRunState(sessionId = getActiveSessionId()) {

package/src/utils.ts

@@ -2,6 +2,7 @@ import {
   MAX_ARRAY_ITEMS,
   MAX_DEPTH,
   MAX_OBJECT_KEYS,
+  MAX_PAYLOAD_NODES,
   MAX_STRING_LENGTH,
   MAX_TOOL_PAYLOAD_LENGTH,
 } from "./constants.js";
@@ -23,11 +24,25 @@ export function tryParseJson(value: string): unknown {
   }
 }
 
-export function shapePayload(value: unknown, options: { maxString?: number; depth?: number } = {}): unknown {
+const PAYLOAD_TOO_LARGE = "[payload too large]";
+
+export function shapePayload(value: unknown, options: { maxString?: number; depth?: number; maxNodes?: number } = {}): unknown {
   const maxString = options.maxString ?? MAX_STRING_LENGTH;
   const depth = options.depth ?? MAX_DEPTH;
+  const maxNodes = options.maxNodes ?? MAX_PAYLOAD_NODES;
+  const budget = { exhausted: false, nodeCount: 0 };
 
   function visit(item: unknown, remainingDepth: number, seen: WeakSet<object>): unknown {
+    if (budget.exhausted) {
+      return PAYLOAD_TOO_LARGE;
+    }
+
+    budget.nodeCount++;
+    if (budget.nodeCount > maxNodes) {
+      budget.exhausted = true;
+      return PAYLOAD_TOO_LARGE;
+    }
+
     if (typeof item === "string") {
       const truncated = truncate(item, maxString);
       const parsed = tryParseJson(truncated);
@@ -59,7 +74,15 @@ export function shapePayload(value: unknown, options: { maxString?: number; dept
     }
 
     if (Array.isArray(item)) {
-      return item.slice(0, MAX_ARRAY_ITEMS).map((entry) => visit(entry, remainingDepth - 1, seen));
+      const output: unknown[] = [];
+      const limit = Math.min(item.length, MAX_ARRAY_ITEMS);
+      for (let index = 0; index < limit; index++) {
+        output.push(visit(item[index], remainingDepth - 1, seen));
+        if (budget.exhausted) {
+          break;
+        }
+      }
+      return output;
     }
 
     if (item instanceof Error) {
@@ -77,8 +100,16 @@ export function shapePayload(value: unknown, options: { maxString?: number; dept
       seen.add(item);
 
       const output: Record<string, unknown> = {};
-      for (const [key, entry] of Object.entries(item as Record<string, unknown>).slice(0, MAX_OBJECT_KEYS)) {
-        output[key] = visit(entry, remainingDepth - 1, seen);
+      let keyCount = 0;
+      for (const key in item as Record<string, unknown>) {
+        if (!Object.hasOwn(item, key)) {
+          continue;
+        }
+        output[key] = visit((item as Record<string, unknown>)[key], remainingDepth - 1, seen);
+        keyCount++;
+        if (budget.exhausted || keyCount >= MAX_OBJECT_KEYS) {
+          break;
+        }
       }
       return output;
     }

package/types/langfuse-runtime-shims.d.ts

@@ -0,0 +1,49 @@
+declare module "@opentelemetry/sdk-trace-base" {
+  export class BasicTracerProvider {
+    constructor(options?: { spanProcessors?: unknown[] });
+    forceFlush?(): Promise<void>;
+    shutdown?(): Promise<void>;
+  }
+}
+
+declare module "@langfuse/otel" {
+  export class LangfuseSpanProcessor {
+    constructor(options: {
+      publicKey: string;
+      secretKey: string;
+      baseUrl: string;
+    });
+    forceFlush?(): Promise<void>;
+    shutdown?(): Promise<void>;
+  }
+}
+
+declare module "@langfuse/tracing" {
+  export function setLangfuseTracerProvider(provider: unknown): void;
+
+  export function startObservation(
+    name: string,
+    body?: Record<string, unknown>,
+    options?: { asType?: string },
+  ): unknown;
+
+  export function propagateAttributes<T>(
+    params: {
+      sessionId?: string;
+      traceName?: string;
+      metadata?: Record<string, string>;
+      tags?: string[];
+    },
+    fn: () => T,
+  ): T;
+}
+
+declare module "@langfuse/client" {
+  export class LangfuseClient {
+    constructor(options: {
+      publicKey: string;
+      secretKey: string;
+      baseUrl: string;
+    });
+  }
+}