pi-langfuse 1.4.5 → 1.4.7

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "pi-langfuse",
-	"version": "1.4.5",
+	"version": "1.4.7",
 	"description": "Langfuse extension for Pi coding agent",
 	"repository": {
 		"type": "git",
@@ -24,7 +24,8 @@
 		"tsconfig.json"
 	],
 	"scripts": {
-		"typecheck": "tsc --noEmit"
+		"typecheck": "tsc --noEmit",
+		"test": "tsx --test test/*.test.ts"
 	},
 	"keywords": [
 		"pi-package",
@@ -60,6 +61,7 @@
 		"node": ">=22"
 	},
 	"devDependencies": {
+		"tsx": "^4.19.0",
 		"typescript": "^6.0.3"
 	}
 }

package/src/commands.ts

@@ -13,6 +13,7 @@ export interface CommandContextLike {
   hasUI?: boolean;
   ui?: {
     notify?: (message: string, level?: "info" | "warning" | "error") => void;
+    select?: (title: string, options: string[]) => Promise<string | undefined>;
   };
 }
 
@@ -147,6 +148,36 @@ function hasActiveAgentObservation() {
   return false;
 }
 
+function savePrivacyPreset(
+  requestedPreset: PrivacyPreset,
+  ctx: CommandContextLike,
+  configPath: string,
+): boolean {
+  const existing = readPersistedConfig(configPath);
+  const loaded = state.config ?? loadConfig(process.env, configPath);
+
+  const publicKey = existing.publicKey ?? loaded?.publicKey;
+  const secretKey = existing.secretKey ?? loaded?.secretKey;
+  const host = existing.host ?? loaded?.host;
+
+  if (!publicKey || !secretKey || !host) {
+    notify(ctx, "Langfuse is not configured yet. Run /langfuse-setup before changing privacy settings.", "warning");
+    return false;
+  }
+
+  const nextConfig = {
+    publicKey: String(publicKey),
+    secretKey: String(secretKey),
+    host: String(host),
+    privacyPreset: requestedPreset,
+  };
+  saveConfig(nextConfig, configPath);
+  state.config = loadConfig(process.env, configPath);
+
+  notify(ctx, `Langfuse privacy preset saved: ${requestedPreset}\n${describePolicy(state.config?.capturePolicy ?? createCapturePolicy())}`);
+  return true;
+}
+
 export async function handleLangfusePrivacyCommand(
   args: string,
   ctx: CommandContextLike,
@@ -163,6 +194,22 @@ export async function handleLangfusePrivacyCommand(
   if (!requestedPreset) {
     state.config = state.config ?? loadConfig(process.env, configPath);
     const policy = state.config?.capturePolicy ?? createCapturePolicy();
+    if (ctx.hasUI && ctx.ui?.select) {
+      const currentPreset = inferPreset(policy);
+      const selectedPreset = await ctx.ui.select(
+        `Langfuse privacy preset (current: ${currentPreset})`,
+        [...PRIVACY_PRESETS],
+      );
+      if (!selectedPreset) {
+        notify(ctx, `Current Langfuse privacy preset: ${currentPreset}\n${describePolicy(policy)}`);
+        return true;
+      }
+      if (!isPrivacyPreset(selectedPreset)) {
+        notify(ctx, `Unknown privacy preset '${selectedPreset}'. Use one of: ${PRIVACY_PRESETS.join(", ")}.`, "warning");
+        return false;
+      }
+      return savePrivacyPreset(selectedPreset, ctx, configPath);
+    }
     notify(ctx, `Current Langfuse privacy preset: ${inferPreset(policy)}\n${describePolicy(policy)}`);
     return true;
   }
@@ -176,29 +223,7 @@ export async function handleLangfusePrivacyCommand(
     return false;
   }
 
-  const existing = readPersistedConfig(configPath);
-  const loaded = state.config ?? loadConfig(process.env, configPath);
-  
-  const publicKey = existing.publicKey ?? loaded?.publicKey;
-  const secretKey = existing.secretKey ?? loaded?.secretKey;
-  const host = existing.host ?? loaded?.host;
-
-  if (!publicKey || !secretKey || !host) {
-    notify(ctx, "Langfuse is not configured yet. Run /langfuse-setup before changing privacy settings.", "warning");
-    return false;
-  }
-
-  const nextConfig = {
-    publicKey: String(publicKey),
-    secretKey: String(secretKey),
-    host: String(host),
-    privacyPreset: requestedPreset,
-  };
-  saveConfig(nextConfig, configPath);
-  state.config = loadConfig(process.env, configPath);
-
-  notify(ctx, `Langfuse privacy preset saved: ${requestedPreset}\n${describePolicy(state.config?.capturePolicy ?? createCapturePolicy())}`);
-  return true;
+  return savePrivacyPreset(requestedPreset, ctx, configPath);
 }
 
 export async function handleLangfuseTestCommand(
@@ -245,8 +270,6 @@ export async function handleLangfuseTestCommand(
         return observation;
       },
     );
-    await rt.tracerProvider?.forceFlush?.();
-    await rt.scoreClient.flush?.();
     notify(ctx, `Langfuse test succeeded. Test trace sent to ${state.config?.host}.`);
     return true;
   } catch (error) {

package/src/langfuse.ts

@@ -45,6 +45,9 @@ interface RestFallbackStore {
 
 const OTEL_VISIBILITY_TIMEOUT_MS = 1_500;
 const OTEL_VISIBILITY_POLL_INTERVAL_MS = 200;
+const DEFAULT_SHUTDOWN_STEP_TIMEOUT_MS = 2_000;
+
+let shutdownStepTimeoutMs = DEFAULT_SHUTDOWN_STEP_TIMEOUT_MS;
 
 function nowIso() {
   return new Date().toISOString();
@@ -54,6 +57,29 @@ function delay(ms: number) {
   return new Promise((resolve) => setTimeout(resolve, ms));
 }
 
+async function withTimeout<T>(label: string, operation: Promise<T> | undefined): Promise<T | undefined> {
+  if (!operation) {
+    return undefined;
+  }
+
+  let timeout: NodeJS.Timeout | undefined;
+  try {
+    return await Promise.race([
+      operation,
+      new Promise<undefined>((resolve) => {
+        timeout = setTimeout(() => {
+          console.log(`📊 Langfuse: ${label} timed out after ${shutdownStepTimeoutMs}ms`);
+          resolve(undefined);
+        }, shutdownStepTimeoutMs);
+      }),
+    ]);
+  } finally {
+    if (timeout) {
+      clearTimeout(timeout);
+    }
+  }
+}
+
 function toIso(value: unknown): string | undefined {
   if (!value) {
     return undefined;
@@ -185,7 +211,10 @@ async function traceExists(rt: LangfuseRuntime, traceId: string): Promise<boolea
     if (!getTrace) {
       return false;
     }
-    await getTrace(traceId);
+    const trace = await withTimeout("Trace visibility check", getTrace(traceId));
+    if (!trace) {
+      return false;
+    }
     return true;
   } catch {
     return false;
@@ -272,21 +301,34 @@ async function fallbackToRestIngestion(rt: LangfuseRuntime) {
     });
   }
 
-  const response = await rt.scoreClient.api?.ingestion?.batch?.({
-    batch,
-    metadata: {
-      source: "pi-langfuse",
-      fallback: "rest-ingestion",
-      reason: "otel-trace-not-visible-after-flush",
-    },
-  });
+  const ingestionBatch = rt.scoreClient.api?.ingestion?.batch;
+  if (!ingestionBatch) {
+    console.log("📊 Langfuse: REST fallback ingestion is unavailable");
+    return;
+  }
+
+  const response = await withTimeout(
+    "REST fallback ingestion",
+    ingestionBatch({
+      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 errors = Array.isArray(responseBody?.errors) ? responseBody.errors : [];
   if (errors.length > 0) {
     console.warn("📊 Langfuse: REST fallback ingestion reported errors", errors);
   } else {
-    console.warn(`📊 Langfuse: OTel trace ${trace.id} was not visible; wrote fallback trace via REST ingestion`);
+    console.log(`📊 Langfuse: OTel trace ${trace.id} was not visible; wrote fallback trace via REST ingestion`);
   }
 }
 
@@ -356,11 +398,11 @@ function doShutdownRuntime(): Promise<void> {
     runtime = null;
 
     try {
-      await rt.tracerProvider?.forceFlush?.();
+      await withTimeout("OTel force flush", rt.tracerProvider?.forceFlush?.());
       await fallbackToRestIngestion(rt);
-      await rt.scoreClient.flush?.();
-      await rt.scoreClient.shutdown?.();
-      await rt.tracerProvider?.shutdown?.();
+      await withTimeout("Langfuse score flush", rt.scoreClient.flush?.());
+      await withTimeout("Langfuse client shutdown", rt.scoreClient.shutdown?.());
+      await withTimeout("OTel tracer shutdown", rt.tracerProvider?.shutdown?.());
     } catch (e) {
       console.warn("📊 Langfuse: Failed to flush/shutdown cleanly", e);
     } finally {
@@ -400,6 +442,12 @@ export async function forceShutdownRuntime(): Promise<void> {
   await doShutdownRuntime();
 }
 
+export function __setRuntimeForTest(rt: LangfuseRuntime | null, timeoutMs = DEFAULT_SHUTDOWN_STEP_TIMEOUT_MS): void {
+  runtime = rt;
+  shutdownStepTimeoutMs = timeoutMs;
+  activeSessions.clear();
+}
+
 export async function sendScore(name: string, value: number, options: { traceId?: string; observationId?: string } = {}) {
   try {
     const rt = await getRuntime();