pi-langfuse 1.4.1 → 1.4.3

package/index.ts

@@ -10,7 +10,7 @@
 import { basename } from "node:path";
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 
-import { state, resetRunState } from "./src/state.js";
+import { state, resetRunState, runWithSession, setCurrentSession } from "./src/state.js";
 import { ensureConfig, promptForConfig, loadConfig } from "./src/config.js";
 import { shutdownRuntime } from "./src/langfuse.js";
 import { getMessageFromEvent, extractAssistantOutput } from "./src/utils.js";
@@ -51,72 +51,79 @@ export default async function (pi: ExtensionAPI) {
     },
   });
 
-  pi.on("session_start", async (_event, ctx) => {
+  const getSessionId = (ctx?: any) => {
+    try {
+      const sessionFile = ctx?.sessionManager?.getSessionFile?.();
+      return sessionFile ? basename(sessionFile, ".jsonl") : undefined;
+    } catch {
+      return undefined;
+    }
+  };
+
+  const withSession = <T>(ctx: any, fn: () => T): T => runWithSession(getSessionId(ctx) ?? state.currentSessionId, fn);
+
+  pi.on("session_start", async (_event, ctx) => withSession(ctx, async () => {
     state.setupAttemptedThisSession = false;
     await ensureConfig(ctx);
-    const sessionFile = ctx.sessionManager.getSessionFile();
-    if (sessionFile) {
-      state.currentSessionId = basename(sessionFile, ".jsonl");
-    }
     resetRunState();
-  });
+  }));
 
-  pi.on("model_select", async (event) => {
+  pi.on("model_select", async (event, ctx) => withSession(ctx, async () => {
     state.currentModel = event.model?.id || "";
     state.currentProvider = event.model?.provider || "";
-  });
+  }));
 
-  pi.on("before_agent_start", async (event, ctx) => {
+  pi.on("before_agent_start", async (event, ctx) => withSession(ctx, async () => {
     await startAgentRun(event, ctx);
-  });
+  }));
 
-  pi.on("agent_start", async (event, ctx) => {
+  pi.on("agent_start", async (event, ctx) => withSession(ctx, async () => {
     if (!state.agentState?.root) {
       await startAgentRun(event, ctx);
     }
-  });
+  }));
 
-  pi.on("turn_start", async (event) => {
+  pi.on("turn_start", async (event, ctx) => withSession(ctx, async () => {
     await startTurnObservation(event);
-  });
+  }));
 
-  pi.on("before_provider_request", async (event) => {
+  pi.on("before_provider_request", async (event, ctx) => withSession(ctx, async () => {
     await startGeneration(event);
-  });
+  }));
 
-  pi.on("after_provider_response", async (event) => {
+  pi.on("after_provider_response", async (event, ctx) => withSession(ctx, async () => {
     updateGenerationMetadata(event);
-  });
+  }));
 
-  pi.on("message_update", async (event) => {
+  pi.on("message_update", async (event, ctx) => withSession(ctx, async () => {
     recordTTFT(event);
     const message = getMessageFromEvent(event);
     if (message?.role === "assistant" && state.agentState) {
       state.agentState.latestAssistantOutput = extractAssistantOutput(message);
     }
-  });
+  }));
 
-  pi.on("message_end", async (event) => {
+  pi.on("message_end", async (event, ctx) => withSession(ctx, async () => {
     await finishGenerationFromMessage(event);
-  });
+  }));
 
-  pi.on("tool_execution_start", async (event) => {
+  pi.on("tool_execution_start", async (event, ctx) => withSession(ctx, async () => {
     await startToolObservation(event);
-  });
+  }));
 
-  pi.on("tool_call", async (event) => {
+  pi.on("tool_call", async (event, ctx) => withSession(ctx, async () => {
     await startToolObservation(event);
-  });
+  }));
 
-  pi.on("tool_result", async (event) => {
+  pi.on("tool_result", async (event, ctx) => withSession(ctx, async () => {
     await finishToolObservation(event);
-  });
+  }));
 
-  pi.on("tool_execution_end", async (event) => {
+  pi.on("tool_execution_end", async (event, ctx) => withSession(ctx, async () => {
     await finishToolObservation(event);
-  });
+  }));
 
-  pi.on("turn_end", async (event) => {
+  pi.on("turn_end", async (event, ctx) => withSession(ctx, async () => {
     state.turnCount++;
     const message = getMessageFromEvent(event);
     if (message?.role === "assistant") {
@@ -124,16 +131,16 @@ export default async function (pi: ExtensionAPI) {
       await finishGenerationFromMessage(event);
     }
     finishTurnObservation(event);
-  });
+  }));
 
-  pi.on("agent_end", async (event) => {
+  pi.on("agent_end", async (event, ctx) => withSession(ctx, async () => {
     await finishAgentRun(event);
     setTimeout(() => {
       shutdownRuntime().catch((error) => {
         console.warn("📊 Langfuse: Deferred shutdown failed", error);
       });
     }, 0);
-  });
+  }));
 
   const handleSessionInterruption = (reason: string) => {
     if (state.agentState?.root) {
@@ -143,15 +150,21 @@ export default async function (pi: ExtensionAPI) {
     resetRunState();
   };
 
-  pi.on("session_before_switch", async () => {
-    handleSessionInterruption("Session switched");
+  pi.on("session_before_switch", async (_event, ctx) => {
+    const sessionId = getSessionId(ctx);
+    if (sessionId) {
+      setCurrentSession(sessionId);
+    }
   });
 
-  pi.on("session_before_fork", async () => {
-    handleSessionInterruption("Session forked");
+  pi.on("session_before_fork", async (_event, ctx) => {
+    const sessionId = getSessionId(ctx);
+    if (sessionId) {
+      setCurrentSession(sessionId);
+    }
   });
 
-  pi.on("session_compact", async (event) => {
+  pi.on("session_compact", async (event, ctx) => withSession(ctx, async () => {
     if (state.agentState?.root) {
       const parent = state.agentState.activeTurn ?? state.agentState.root;
       try {
@@ -169,10 +182,10 @@ export default async function (pi: ExtensionAPI) {
         // ignore
       }
     }
-  });
+  }));
 
-  pi.on("session_shutdown", async () => {
+  pi.on("session_shutdown", async (_event, ctx) => withSession(ctx, async () => {
     handleSessionInterruption("Session shutdown before agent completed");
     await shutdownRuntime();
-  });
+  }));
 }

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "pi-langfuse",
-	"version": "1.4.1",
+	"version": "1.4.3",
 	"description": "Langfuse extension for Pi coding agent",
 	"repository": {
 		"type": "git",
@@ -12,6 +12,16 @@
 	"homepage": "https://github.com/gooyoung/pi-langfuse#readme",
 	"type": "module",
 	"main": "index.ts",
+	"files": [
+		"index.ts",
+		"src/",
+		"types/",
+		"README.md",
+		"README_CN.md",
+		"image.png",
+		"skills-lock.json",
+		"tsconfig.json"
+	],
 	"scripts": {
 		"typecheck": "tsc --noEmit"
 	},

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

@@ -1,38 +1,150 @@
+import { AsyncLocalStorage } from "node:async_hooks";
 import type { Config, AgentState } from "./types.js";
 
+export interface SessionRunState {
+  currentModel: string;
+  currentProvider: string;
+  agentState: AgentState | null;
+  toolCallCount: number;
+  errorCount: number;
+  turnCount: number;
+  tracingDisabled: boolean;
+}
+
+const DEFAULT_SESSION_ID = "__pi_langfuse_default_session__";
+
+let activeSessionId = DEFAULT_SESSION_ID;
+const sessionScope = new AsyncLocalStorage<string>();
+
+function createSessionRunState(): SessionRunState {
+  return {
+    currentModel: "",
+    currentProvider: "",
+    agentState: null,
+    toolCallCount: 0,
+    errorCount: 0,
+    turnCount: 0,
+    tracingDisabled: false,
+  };
+}
+
+function normalizeSessionId(sessionId?: string) {
+  return sessionId || DEFAULT_SESSION_ID;
+}
+
+function getActiveSessionId() {
+  return sessionScope.getStore() ?? activeSessionId;
+}
+
+export function getSessionRunState(sessionId = getActiveSessionId()): SessionRunState {
+  const normalizedSessionId = normalizeSessionId(sessionId);
+  let sessionState = state.sessionStates.get(normalizedSessionId);
+  if (!sessionState) {
+    sessionState = createSessionRunState();
+    state.sessionStates.set(normalizedSessionId, sessionState);
+  }
+  return sessionState;
+}
+
+export function setCurrentSession(sessionId?: string) {
+  activeSessionId = normalizeSessionId(sessionId);
+  getSessionRunState(activeSessionId);
+}
+
+export function runWithSession<T>(sessionId: string | undefined, fn: () => T): T {
+  const normalizedSessionId = normalizeSessionId(sessionId);
+  setCurrentSession(normalizedSessionId);
+  return sessionScope.run(normalizedSessionId, fn);
+}
+
 export const state = {
   config: null as Config | null,
   setupAttemptedThisSession: false,
-  
-  currentSessionId: "",
-  currentModel: "",
-  currentProvider: "",
-  agentState: null as AgentState | null,
-  
-  // Evaluation tracking state
-  toolCallCount: 0,
-  errorCount: 0,
-  turnCount: 0,
+  sessionStates: new Map<string, SessionRunState>(),
+
+  get currentSessionId() {
+    const sessionId = getActiveSessionId();
+    return sessionId === DEFAULT_SESSION_ID ? "" : sessionId;
+  },
+  set currentSessionId(sessionId: string) {
+    setCurrentSession(sessionId);
+  },
+
+  get currentModel() {
+    return getSessionRunState().currentModel;
+  },
+  set currentModel(model: string) {
+    getSessionRunState().currentModel = model;
+  },
+
+  get currentProvider() {
+    return getSessionRunState().currentProvider;
+  },
+  set currentProvider(provider: string) {
+    getSessionRunState().currentProvider = provider;
+  },
+
+  get agentState() {
+    return getSessionRunState().agentState;
+  },
+  set agentState(agentState: AgentState | null) {
+    getSessionRunState().agentState = agentState;
+  },
+
+  get toolCallCount() {
+    return getSessionRunState().toolCallCount;
+  },
+  set toolCallCount(toolCallCount: number) {
+    getSessionRunState().toolCallCount = toolCallCount;
+  },
+
+  get errorCount() {
+    return getSessionRunState().errorCount;
+  },
+  set errorCount(errorCount: number) {
+    getSessionRunState().errorCount = errorCount;
+  },
+
+  get turnCount() {
+    return getSessionRunState().turnCount;
+  },
+  set turnCount(turnCount: number) {
+    getSessionRunState().turnCount = turnCount;
+  },
+
+  get isTracingDisabled() {
+    return getSessionRunState().tracingDisabled;
+  },
+  set isTracingDisabled(disabled: boolean) {
+    getSessionRunState().tracingDisabled = disabled;
+  },
 };
 
-export function resetRunState() {
-  state.agentState = null;
-  state.toolCallCount = 0;
-  state.errorCount = 0;
-  state.turnCount = 0;
-  state.currentModel = "";
-  state.currentProvider = "";
+export function resetRunState(sessionId = getActiveSessionId()) {
+  state.sessionStates.set(normalizeSessionId(sessionId), createSessionRunState());
 }
 
-export function computeEvaluationScores() {
-  const toolSuccessRate = state.toolCallCount > 0 ? (state.toolCallCount - state.errorCount) / state.toolCallCount : 1;
-  const sessionHadErrors = state.errorCount > 0;
+export function clearAllSessionStates() {
+  state.sessionStates.clear();
+  activeSessionId = DEFAULT_SESSION_ID;
+  getSessionRunState();
+}
+
+export function computeEvaluationScores(sessionId = getActiveSessionId()) {
+  const sessionState = getSessionRunState(sessionId);
+  const toolSuccessRate =
+    sessionState.toolCallCount > 0
+      ? (sessionState.toolCallCount - sessionState.errorCount) / sessionState.toolCallCount
+      : 1;
+  const sessionHadErrors = sessionState.errorCount > 0;
 
   return {
-    tool_call_count: state.toolCallCount,
-    turn_count: state.turnCount,
-    total_tool_errors: state.errorCount,
+    tool_call_count: sessionState.toolCallCount,
+    turn_count: sessionState.turnCount,
+    total_tool_errors: sessionState.errorCount,
     tool_success_rate: toolSuccessRate,
     session_had_errors: sessionHadErrors ? 1 : 0,
   };
 }
+
+getSessionRunState();

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