pi-langfuse 1.5.7 → 1.5.9

package/README.md

@@ -77,6 +77,14 @@ export LANGFUSE_BASE_URL="https://cloud.langfuse.com"  # optional; LANGFUSE_HOST
 
 Saved config takes precedence. Environment variables are only used when `~/.pi/agent/pi-langfuse/config.json` is missing or incomplete.
 
+For short-lived SDK hosts, set the bounded final score-delivery attempt during shutdown:
+
+```bash
+export PI_LANGFUSE_SCORE_SHUTDOWN_TIMEOUT=2  # seconds; defaults to 2 seconds
+```
+
+The extension attempts queued trace-level scores before other shutdown telemetry work. This value cannot extend the overall shutdown deadline.
+
 Privacy controls can also be set through environment variables:
 
 ```bash

package/README_CN.md

@@ -77,6 +77,14 @@ export LANGFUSE_BASE_URL="https://cloud.langfuse.com"  # 可选；也支持 LANG
 
 保存的配置优先级更高。只有当 `~/.pi/agent/pi-langfuse/config.json` 缺失或不完整时，扩展才会使用环境变量。
 
+对于短生命周期的 SDK 宿主，可设置关闭时最终分数发送尝试的上限：
+
+```bash
+export PI_LANGFUSE_SCORE_SHUTDOWN_TIMEOUT=2  # 单位为秒；默认 2 秒
+```
+
+扩展会在其他关闭遥测工作之前尝试发送已排队的 trace 级分数。该值不会延长总关闭超时。
+
 隐私采集策略也可以通过环境变量设置：
 
 ```bash

package/index.ts

@@ -159,11 +159,11 @@ export default async function (pi: ExtensionAPI) {
   pi.on("agent_end", async (event, ctx) => withSession(ctx, async () => {
     await finishAgentRun(event);
     const sessionId = state.currentSessionId;
-    setTimeout(() => {
-      shutdownRuntime(sessionId).catch((error) => {
-        console.warn("📊 Langfuse: Deferred shutdown failed", error);
-      });
-    }, 0);
+    try {
+      await shutdownRuntime(sessionId);
+    } catch (error) {
+      console.warn("📊 Langfuse: Shutdown failed", error);
+    }
   }));
 
   const handleSessionInterruption = (reason: string) => {

package/package.json

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

package/src/langfuse.ts

@@ -1,4 +1,4 @@
-import type { LangfuseRuntime, LangfuseScoreClient } from "./types.js";
+import type { LangfuseRuntime, LangfuseScoreClient, PendingScore } from "./types.js";
 import { state } from "./state.js";
 import { randomUUID } from "node:crypto";
 
@@ -60,6 +60,12 @@ interface RestFallbackStore {
 const OTEL_VISIBILITY_TIMEOUT_MS = 1_500;
 const OTEL_VISIBILITY_POLL_INTERVAL_MS = 200;
 const DEFAULT_SHUTDOWN_STEP_TIMEOUT_MS = 2_000;
+const DEFAULT_SCORE_SHUTDOWN_TIMEOUT_MS = 2_000;
+const DEFAULT_LANGFUSE_REQUEST_TIMEOUT_SECONDS = 5;
+const DEFAULT_SCORE_FLUSH_AT = 10;
+const DEFAULT_SCORE_FLUSH_INTERVAL_MS = 1_000;
+const MAX_SCORE_QUEUE_SIZE = 100_000;
+const MAX_SCORE_BATCH_SIZE = 100;
 
 let shutdownStepTimeoutMs = DEFAULT_SHUTDOWN_STEP_TIMEOUT_MS;
 
@@ -67,8 +73,34 @@ function nowIso() {
   return new Date().toISOString();
 }
 
-function delay(ms: number) {
-  return new Promise((resolve) => setTimeout(resolve, ms));
+function resolvePositiveEnvNumber(name: string, fallback: number, integer = false): number {
+  const parsed = Number(process.env[name]);
+  if (!Number.isFinite(parsed) || parsed <= 0) {
+    return fallback;
+  }
+  return integer ? Math.floor(parsed) : parsed;
+}
+
+function getScoreShutdownTimeoutMs(): number {
+  return resolvePositiveEnvNumber(
+    "PI_LANGFUSE_SCORE_SHUTDOWN_TIMEOUT",
+    DEFAULT_SCORE_SHUTDOWN_TIMEOUT_MS / 1_000,
+  ) * 1_000;
+}
+
+function delay(ms: number, signal?: AbortSignal) {
+  return new Promise<void>((resolve, reject) => {
+    if (signal?.aborted) {
+      reject(signal.reason);
+      return;
+    }
+
+    const timeout = setTimeout(resolve, ms);
+    signal?.addEventListener("abort", () => {
+      clearTimeout(timeout);
+      reject(signal.reason);
+    }, { once: true });
+  });
 }
 
 function debugLog(message: string) {
@@ -107,7 +139,14 @@ export function getLastRuntimeError(): { scope: string; message: string; timesta
   return lastRuntimeError;
 }
 
-async function withTimeout<T>(label: string, operation: Promise<T> | undefined): Promise<T | undefined> {
+async function withShutdownDeadline<T>(label: string, startOperation: () => Promise<T> | undefined, deadline: number): Promise<T | undefined> {
+  const remainingMs = deadline - Date.now();
+  if (remainingMs <= 0) {
+    debugLog(`📊 Langfuse: Skipped ${label}; shutdown deadline elapsed`);
+    return undefined;
+  }
+
+  const operation = startOperation();
   if (!operation) {
     return undefined;
   }
@@ -118,9 +157,9 @@ async function withTimeout<T>(label: string, operation: Promise<T> | undefined):
       operation,
       new Promise<undefined>((resolve) => {
         timeout = setTimeout(() => {
-          debugLog(`📊 Langfuse: ${label} timed out after ${shutdownStepTimeoutMs}ms`);
+          debugLog(`📊 Langfuse: ${label} timed out; shutdown deadline elapsed`);
           resolve(undefined);
-        }, shutdownStepTimeoutMs);
+        }, remainingMs);
       }),
     ]);
   } finally {
@@ -130,6 +169,142 @@ async function withTimeout<T>(label: string, operation: Promise<T> | undefined):
   }
 }
 
+function getRuntimeConfig(rt: LangfuseRuntime) {
+  return rt.runtimeConfig ?? state.config;
+}
+
+function ingestionHeaders(rt: LangfuseRuntime): Record<string, string> {
+  const config = getRuntimeConfig(rt);
+  if (!config) {
+    throw new Error("Langfuse runtime config is unavailable");
+  }
+
+  const auth = Buffer.from(`${config.publicKey}:${config.secretKey}`).toString("base64");
+  return {
+    Authorization: `Basic ${auth}`,
+    "Content-Type": "application/json",
+  };
+}
+
+async function ingestBatch(rt: LangfuseRuntime, batch: unknown[], signal: AbortSignal): Promise<unknown[]> {
+  const config = getRuntimeConfig(rt);
+  if (!config) {
+    throw new Error("Langfuse runtime config is unavailable");
+  }
+
+  const response = await fetch(`${config.host.replace(/\/$/, "")}/api/public/ingestion`, {
+    method: "POST",
+    headers: ingestionHeaders(rt),
+    body: JSON.stringify({ batch }),
+    signal,
+  });
+  if (!response.ok) {
+    throw new Error(`Langfuse ingestion failed with HTTP ${response.status}`);
+  }
+
+  const text = await response.text();
+  if (!text) {
+    return [];
+  }
+
+  const responseBody = JSON.parse(text) as { errors?: unknown[] };
+  return Array.isArray(responseBody.errors) ? responseBody.errors : [];
+}
+
+async function flushPendingScores(rt: LangfuseRuntime, signal: AbortSignal): Promise<void> {
+  const pendingScores = rt.pendingScores;
+  if (!pendingScores || pendingScores.length === 0) {
+    return;
+  }
+
+  while (pendingScores.length > 0) {
+    const scores = pendingScores.slice(0, MAX_SCORE_BATCH_SIZE);
+    try {
+      const errors = await ingestBatch(
+        rt,
+        scores.map((score) => ({
+          type: "score-create",
+          id: randomUUID(),
+          timestamp: nowIso(),
+          body: score,
+        })),
+        signal,
+      );
+      pendingScores.splice(0, scores.length);
+      if (errors.length > 0) {
+        rememberRuntimeError("score ingestion", new Error(JSON.stringify(errors)));
+        console.warn("📊 Langfuse: Score ingestion reported errors", errors);
+      }
+    } catch (error) {
+      if ((error as { name?: string }).name !== "AbortError") {
+        rememberRuntimeError("score ingestion", error);
+        console.warn("📊 Langfuse: Failed to flush scores", error);
+      }
+      return;
+    }
+  }
+}
+
+function clearScoreFlushTimer(rt: LangfuseRuntime) {
+  if (rt.scoreFlushTimer) {
+    clearTimeout(rt.scoreFlushTimer);
+    rt.scoreFlushTimer = undefined;
+  }
+}
+
+function scheduleScoreFlush(rt: LangfuseRuntime) {
+  if (
+    rt.scoreFlushStopped
+    || rt.scoreFlushTimer
+    || rt.scoreFlushPromise
+    || !rt.pendingScores?.length
+  ) {
+    return;
+  }
+
+  rt.scoreFlushTimer = setTimeout(() => {
+    rt.scoreFlushTimer = undefined;
+    void startScoreFlush(rt);
+  }, rt.scoreFlushIntervalMs ?? DEFAULT_SCORE_FLUSH_INTERVAL_MS);
+  rt.scoreFlushTimer.unref?.();
+}
+
+function startScoreFlush(rt: LangfuseRuntime): Promise<void> {
+  if (rt.scoreFlushPromise) {
+    return rt.scoreFlushPromise;
+  }
+
+  clearScoreFlushTimer(rt);
+  const controller = new AbortController();
+  const timeout = setTimeout(
+    () => controller.abort(new DOMException("Langfuse score request timed out", "AbortError")),
+    rt.scoreRequestTimeoutMs ?? DEFAULT_LANGFUSE_REQUEST_TIMEOUT_SECONDS * 1_000,
+  );
+  timeout.unref?.();
+  rt.scoreFlushController = controller;
+
+  const promise = flushPendingScores(rt, controller.signal).finally(() => {
+    clearTimeout(timeout);
+    if (rt.scoreFlushPromise === promise) {
+      rt.scoreFlushPromise = undefined;
+    }
+    if (rt.scoreFlushController === controller) {
+      rt.scoreFlushController = undefined;
+    }
+    scheduleScoreFlush(rt);
+  });
+  rt.scoreFlushPromise = promise;
+  return promise;
+}
+
+function stopScoreFlush(rt: LangfuseRuntime) {
+  rt.scoreFlushStopped = true;
+  clearScoreFlushTimer(rt);
+  rt.scoreFlushController?.abort(
+    new DOMException("Langfuse score flushing stopped", "AbortError"),
+  );
+}
+
 function toIso(value: unknown): string | undefined {
   if (!value) {
     return undefined;
@@ -258,27 +433,40 @@ function wrapObservation(
   };
 }
 
-async function traceExists(rt: LangfuseRuntime, traceId: string): Promise<boolean> {
+async function traceExists(rt: LangfuseRuntime, traceId: string, signal: AbortSignal): Promise<boolean> {
+  const config = getRuntimeConfig(rt);
+  if (!config) {
+    return false;
+  }
+
   try {
-    const traceApi = rt.scoreClient.api?.trace;
-    if (!traceApi?.get) {
+    const response = await fetch(
+      `${config.host.replace(/\/$/, "")}/api/public/traces/${encodeURIComponent(traceId)}`,
+      {
+        headers: ingestionHeaders(rt),
+        signal,
+      },
+    );
+    if (response.status === 404) {
       return false;
     }
-    const trace = await withTimeout("Trace visibility check", traceApi.get(traceId));
-    if (!trace) {
-      return false;
+    if (!response.ok) {
+      throw new Error(`Langfuse trace visibility check failed with HTTP ${response.status}`);
     }
     return true;
-  } catch {
+  } catch (error) {
+    if (signal.aborted) {
+      throw error;
+    }
     return false;
   }
 }
 
-async function waitForTraceVisibility(rt: LangfuseRuntime, traceId: string): Promise<boolean> {
+async function waitForTraceVisibility(rt: LangfuseRuntime, traceId: string, signal: AbortSignal): Promise<boolean> {
   const deadline = Date.now() + OTEL_VISIBILITY_TIMEOUT_MS;
 
   while (true) {
-    if (await traceExists(rt, traceId)) {
+    if (await traceExists(rt, traceId, signal)) {
       return true;
     }
 
@@ -287,7 +475,7 @@ async function waitForTraceVisibility(rt: LangfuseRuntime, traceId: string): Pro
       return false;
     }
 
-    await delay(Math.min(OTEL_VISIBILITY_POLL_INTERVAL_MS, remainingMs));
+    await delay(Math.min(OTEL_VISIBILITY_POLL_INTERVAL_MS, remainingMs), signal);
   }
 }
 
@@ -295,14 +483,14 @@ function eventTimestamp(record: { endTime?: string; startTime?: string; timestam
   return record.endTime ?? record.startTime ?? record.timestamp ?? nowIso();
 }
 
-async function fallbackToRestIngestion(rt: LangfuseRuntime) {
+async function fallbackToRestIngestion(rt: LangfuseRuntime, signal: AbortSignal) {
   const store = rt.restFallback as RestFallbackStore | undefined;
   if (!store?.trace || store.attempted) {
     return;
   }
   store.attempted = true;
 
-  if (await waitForTraceVisibility(rt, store.trace.id)) {
+  if (await waitForTraceVisibility(rt, store.trace.id, signal)) {
     return;
   }
 
@@ -355,31 +543,7 @@ async function fallbackToRestIngestion(rt: LangfuseRuntime) {
     });
   }
 
-  const ingestionApi = rt.scoreClient.api?.ingestion;
-  if (!ingestionApi?.batch) {
-    debugLog("📊 Langfuse: REST fallback ingestion is unavailable");
-    return;
-  }
-
-  const response = await withTimeout(
-    "REST fallback ingestion",
-    ingestionApi.batch({
-      batch,
-      metadata: {
-        source: "pi-langfuse",
-        fallback: "rest-ingestion",
-        reason: "otel-trace-not-visible-after-flush",
-      },
-    }),
-  );
-
-  if (!response) {
-    return;
-  }
-
-  const responseBody = response as { errors?: unknown[] } | undefined;
-  const responseErrors = responseBody?.errors;
-  const errors = Array.isArray(responseErrors) ? responseErrors : [];
+  const errors = await ingestBatch(rt, batch, signal);
   if (errors.length > 0) {
     rememberRuntimeError("REST fallback ingestion", new Error(JSON.stringify(errors)));
     console.warn("📊 Langfuse: REST fallback ingestion reported errors", errors);
@@ -426,6 +590,11 @@ export async function getRuntime(): Promise<LangfuseRuntime> {
 
     try {
       ensureOtelContextManager(context, AsyncHooksContextManager);
+      const scoreFlushAt = resolvePositiveEnvNumber("LANGFUSE_FLUSH_AT", DEFAULT_SCORE_FLUSH_AT, true);
+      const scoreFlushIntervalMs =
+        resolvePositiveEnvNumber("LANGFUSE_FLUSH_INTERVAL", DEFAULT_SCORE_FLUSH_INTERVAL_MS / 1_000) * 1_000;
+      const scoreRequestTimeoutMs =
+        resolvePositiveEnvNumber("LANGFUSE_TIMEOUT", DEFAULT_LANGFUSE_REQUEST_TIMEOUT_SECONDS) * 1_000;
       const spanProcessor = new LangfuseSpanProcessor({
         publicKey: state.config.publicKey,
         secretKey: state.config.secretKey,
@@ -449,6 +618,16 @@ export async function getRuntime(): Promise<LangfuseRuntime> {
         tracerProvider,
         clearTracerProvider: () => tracing.setLangfuseTracerProvider(null),
         restFallback,
+        pendingScores: [],
+        scoreFlushAt,
+        scoreFlushIntervalMs,
+        scoreRequestTimeoutMs,
+        scoreFlushStopped: false,
+        runtimeConfig: {
+          publicKey: state.config.publicKey,
+          secretKey: state.config.secretKey,
+          host: state.config.host,
+        },
       };
       lastRuntimeError = null;
     } catch (e) {
@@ -468,17 +647,44 @@ function doShutdownRuntime(): Promise<void> {
 
     const rt = runtime;
     runtime = null;
+    const deadline = Date.now() + shutdownStepTimeoutMs;
+    const controller = new AbortController();
+    const abortTimeout = setTimeout(() => controller.abort(), shutdownStepTimeoutMs);
+    const scoreController = new AbortController();
+    const scoreAbortTimeout = setTimeout(
+      () => scoreController.abort(),
+      Math.min(getScoreShutdownTimeoutMs(), shutdownStepTimeoutMs),
+    );
+    scoreAbortTimeout.unref?.();
+    stopScoreFlush(rt);
 
     try {
-      await withTimeout("OTel force flush", rt.tracerProvider?.forceFlush?.());
-      await fallbackToRestIngestion(rt);
-      await withTimeout("Langfuse score flush", rt.scoreClient.flush?.());
-      await withTimeout("Langfuse client shutdown", rt.scoreClient.shutdown?.());
-      await withTimeout("OTel tracer shutdown", rt.tracerProvider?.shutdown?.());
+      await withShutdownDeadline(
+        "Active score flush",
+        () => rt.scoreFlushPromise,
+        deadline,
+      );
+      await flushPendingScores(rt, scoreController.signal);
+      await withShutdownDeadline("OTel force flush", () => rt.tracerProvider?.forceFlush?.(), deadline);
+      await withShutdownDeadline(
+        "REST fallback ingestion",
+        () => fallbackToRestIngestion(rt, controller.signal),
+        deadline,
+      );
+      await withShutdownDeadline("Langfuse score flush", () => rt.scoreClient.flush?.(), deadline);
+      await withShutdownDeadline("Langfuse client shutdown", () => rt.scoreClient.shutdown?.(), deadline);
+      await withShutdownDeadline("OTel tracer shutdown", () => rt.tracerProvider?.shutdown?.(), deadline);
     } catch (e) {
       rememberRuntimeError("runtime shutdown", e);
       console.warn("📊 Langfuse: Failed to flush/shutdown cleanly", e);
     } finally {
+      clearTimeout(abortTimeout);
+      clearTimeout(scoreAbortTimeout);
+      scoreController.abort();
+      clearScoreFlushTimer(rt);
+      rt.scoreFlushController?.abort();
+      rt.scoreFlushController = undefined;
+      rt.scoreFlushPromise = undefined;
       if (!runtime) {
         rt.clearTracerProvider?.();
       }
@@ -516,7 +722,13 @@ export async function forceShutdownRuntime(): Promise<void> {
 }
 
 export function __setRuntimeForTest(rt: LangfuseRuntime | null, timeoutMs = DEFAULT_SHUTDOWN_STEP_TIMEOUT_MS): void {
+  if (runtime && runtime !== rt) {
+    stopScoreFlush(runtime);
+  }
   runtime = rt;
+  if (rt) {
+    rt.scoreFlushStopped = false;
+  }
   shutdownStepTimeoutMs = timeoutMs;
   activeSessions.clear();
 }
@@ -524,14 +736,35 @@ export function __setRuntimeForTest(rt: LangfuseRuntime | null, timeoutMs = DEFA
 export async function sendScore(name: string, value: number, options: { traceId?: string; observationId?: string } = {}) {
   try {
     const rt = await getRuntime();
-    rt.scoreClient.score?.create({
+    const score: PendingScore = {
       name,
       value,
       dataType: name === "session_had_errors" || name === "tool_is_error" ? "BOOLEAN" : "NUMERIC",
       traceId: options.traceId,
       observationId: options.observationId,
       sessionId: options.traceId ? undefined : state.currentSessionId || undefined,
-    });
+      ...(process.env.LANGFUSE_TRACING_ENVIRONMENT
+        ? { environment: process.env.LANGFUSE_TRACING_ENVIRONMENT }
+        : {}),
+    };
+    if (!rt.pendingScores) {
+      return;
+    }
+    if (rt.pendingScores.length >= MAX_SCORE_QUEUE_SIZE) {
+      const error = new Error(
+        `Langfuse score queue is full (${MAX_SCORE_QUEUE_SIZE}); dropping score`,
+      );
+      rememberRuntimeError("score queue", error);
+      console.warn(`📊 Langfuse: ${error.message}`);
+      return;
+    }
+
+    rt.pendingScores.push(score);
+    if (rt.pendingScores.length >= (rt.scoreFlushAt ?? DEFAULT_SCORE_FLUSH_AT)) {
+      void startScoreFlush(rt);
+    } else {
+      scheduleScoreFlush(rt);
+    }
   } catch (e) {
     rememberRuntimeError(`score ${name}`, e);
     console.warn(`📊 Langfuse: Failed to send score ${name}`, e);

package/src/types.ts

@@ -59,6 +59,22 @@ export interface LangfuseScoreClient {
   shutdown?: () => Promise<void>;
 }
 
+export interface PendingScore {
+  traceId?: string;
+  sessionId?: string;
+  observationId?: string;
+  name: string;
+  value: number;
+  dataType?: "NUMERIC" | "BOOLEAN";
+  environment?: string;
+}
+
+export interface LangfuseRuntimeConfig {
+  publicKey: string;
+  secretKey: string;
+  host: string;
+}
+
 export interface LangfuseRuntime {
   startObservation: (
     name: string,
@@ -79,6 +95,15 @@ export interface LangfuseRuntime {
   tracerProvider?: { forceFlush?: () => Promise<void>; shutdown?: () => Promise<void> };
   clearTracerProvider?: () => void;
   restFallback?: unknown;
+  pendingScores?: PendingScore[];
+  scoreFlushAt?: number;
+  scoreFlushIntervalMs?: number;
+  scoreRequestTimeoutMs?: number;
+  scoreFlushTimer?: NodeJS.Timeout;
+  scoreFlushPromise?: Promise<void>;
+  scoreFlushController?: AbortController;
+  scoreFlushStopped?: boolean;
+  runtimeConfig?: LangfuseRuntimeConfig;
 }
 
 export interface GenerationState {