pi-langfuse 1.4.4 → 1.4.7

package/index.ts

@@ -13,6 +13,7 @@ import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import { state, resetRunState, runWithSession, setCurrentSession } from "./src/state.js";
 import { ensureConfig, promptForConfig, loadConfig } from "./src/config.js";
 import { shutdownRuntime } from "./src/langfuse.js";
+import { handleLangfusePrivacyCommand, handleLangfuseTestCommand } from "./src/commands.js";
 import { getMessageFromEvent, extractAssistantOutput, getCapturePolicy } from "./src/utils.js";
 import { applyCapturePolicy } from "./src/capture-policy.js";
 import { startAgentRun, finishAgentRun } from "./src/handlers/agent.js";
@@ -52,6 +53,20 @@ export default async function (pi: ExtensionAPI) {
     },
   });
 
+  pi.registerCommand("langfuse-test", {
+    description: "Send a test trace to Langfuse to verify configuration",
+    handler: async (args, ctx) => {
+      await handleLangfuseTestCommand(String(args ?? ""), ctx);
+    },
+  });
+
+  pi.registerCommand("langfuse-privacy", {
+    description: "View or set Langfuse telemetry privacy preset",
+    handler: async (args, ctx) => {
+      await handleLangfusePrivacyCommand(String(args ?? ""), ctx);
+    },
+  });
+
   const getSessionId = (ctx?: any) => {
     try {
       const sessionFile = ctx?.sessionManager?.getSessionFile?.();

package/package.json

@@ -1,6 +1,6 @@
 {
 	"name": "pi-langfuse",
-	"version": "1.4.4",
+	"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

@@ -0,0 +1,284 @@
+import { existsSync, readFileSync } from "node:fs";
+
+import { CONFIG_PATH } from "./constants.js";
+import { loadConfig, saveConfig, ensureConfig } from "./config.js";
+import { createCapturePolicy, type PrivacyPreset, type CapturePolicy } from "./capture-policy.js";
+import { getRuntime, forceShutdownRuntime as shutdownLangfuseRuntime } from "./langfuse.js";
+import { state } from "./state.js";
+import type { LangfuseRuntime } from "./types.js";
+
+const PRIVACY_PRESETS = ["metadata-only", "prompts-only", "conversations", "full-debug"] as const;
+
+export interface CommandContextLike {
+  hasUI?: boolean;
+  ui?: {
+    notify?: (message: string, level?: "info" | "warning" | "error") => void;
+    select?: (title: string, options: string[]) => Promise<string | undefined>;
+  };
+}
+
+interface CommandDeps {
+  configPath?: string;
+  getRuntime?: () => Promise<LangfuseRuntime>;
+  forceShutdownRuntime?: () => Promise<void>;
+}
+
+function notify(ctx: CommandContextLike, message: string, level: "info" | "warning" | "error" = "info") {
+  if (ctx.hasUI && ctx.ui?.notify) {
+    ctx.ui.notify(message, level);
+    return;
+  }
+  const prefix = level === "error" ? "❌" : level === "warning" ? "⚠️" : "📊";
+  console.log(`${prefix} Langfuse: ${message}`);
+}
+
+function parseCommandArgs(args: string): { values: Record<string, string>; positional: string[]; malformed: string[] } {
+  const values: Record<string, string> = {};
+  const positional: string[] = [];
+  const malformed: string[] = [];
+
+  for (const part of args.trim().split(/\s+/)) {
+    if (!part) {
+      continue;
+    }
+    const eq = part.indexOf("=");
+    if (eq === -1) {
+      positional.push(part);
+      continue;
+    }
+    if (eq === 0) {
+      malformed.push(part);
+      continue;
+    }
+    values[part.slice(0, eq)] = part.slice(eq + 1);
+  }
+
+  return { values, positional, malformed };
+}
+
+function isPrivacyPreset(value: string | undefined): value is PrivacyPreset {
+  return PRIVACY_PRESETS.includes(value as PrivacyPreset);
+}
+
+function inferPreset(policy: CapturePolicy): PrivacyPreset | "custom" {
+  const entries: Array<[PrivacyPreset, CapturePolicy]> = [
+    [
+      "metadata-only",
+      {
+        captureInputs: false,
+        captureOutputs: false,
+        captureToolIo: false,
+        captureSystemPrompt: false,
+        captureCwd: false,
+      },
+    ],
+    [
+      "prompts-only",
+      {
+        captureInputs: true,
+        captureOutputs: false,
+        captureToolIo: false,
+        captureSystemPrompt: false,
+        captureCwd: false,
+      },
+    ],
+    [
+      "conversations",
+      {
+        captureInputs: true,
+        captureOutputs: true,
+        captureToolIo: false,
+        captureSystemPrompt: false,
+        captureCwd: false,
+      },
+    ],
+    [
+      "full-debug",
+      {
+        captureInputs: true,
+        captureOutputs: true,
+        captureToolIo: true,
+        captureSystemPrompt: true,
+        captureCwd: true,
+      },
+    ],
+  ];
+
+  for (const [preset, presetPolicy] of entries) {
+    if (
+      policy.captureInputs === presetPolicy.captureInputs &&
+      policy.captureOutputs === presetPolicy.captureOutputs &&
+      policy.captureToolIo === presetPolicy.captureToolIo &&
+      policy.captureSystemPrompt === presetPolicy.captureSystemPrompt &&
+      policy.captureCwd === presetPolicy.captureCwd
+    ) {
+      return preset;
+    }
+  }
+  return "custom";
+}
+
+function describePolicy(policy: CapturePolicy) {
+  return [
+    `captureInputs: ${policy.captureInputs}`,
+    `captureOutputs: ${policy.captureOutputs}`,
+    `captureToolIo: ${policy.captureToolIo}`,
+    `captureSystemPrompt: ${policy.captureSystemPrompt}`,
+    `captureCwd: ${policy.captureCwd}`,
+  ].join("\n");
+}
+
+function readPersistedConfig(path: string) {
+  if (!existsSync(path)) {
+    return {};
+  }
+  try {
+    return JSON.parse(readFileSync(path, "utf-8")) as Record<string, unknown>;
+  } catch {
+    return {};
+  }
+}
+
+function hasActiveAgentObservation() {
+  for (const sessionState of state.sessionStates.values()) {
+    if (sessionState.agentState?.root) {
+      return true;
+    }
+  }
+  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,
+  deps: CommandDeps = {},
+): Promise<boolean> {
+  const configPath = deps.configPath ?? CONFIG_PATH;
+  const parsed = parseCommandArgs(args);
+  if (parsed.malformed.length > 0) {
+    notify(ctx, `Couldn't understand '${parsed.malformed[0]}'. Use /langfuse-privacy preset=metadata-only.`, "warning");
+    return false;
+  }
+
+  const requestedPreset = parsed.values.preset ?? parsed.positional[0];
+  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;
+  }
+
+  if (!isPrivacyPreset(requestedPreset)) {
+    notify(
+      ctx,
+      `Unknown privacy preset '${requestedPreset}'. Use one of: ${PRIVACY_PRESETS.join(", ")}.`,
+      "warning",
+    );
+    return false;
+  }
+
+  return savePrivacyPreset(requestedPreset, ctx, configPath);
+}
+
+export async function handleLangfuseTestCommand(
+  _args: string,
+  ctx: CommandContextLike,
+  deps: CommandDeps = {},
+): Promise<boolean> {
+  if (!state.config && !(await ensureConfig(ctx))) {
+    notify(ctx, "Langfuse is not configured yet. Run /langfuse-setup first.", "warning");
+    return false;
+  }
+
+  if (hasActiveAgentObservation()) {
+    notify(ctx, "Langfuse test skipped because an agent run is active. Try again after the run finishes.", "warning");
+    return false;
+  }
+
+  let runtimeInitialized = false;
+  try {
+    const rt = await (deps.getRuntime ?? getRuntime)();
+    runtimeInitialized = true;
+    rt.propagateAttributes(
+      {
+        traceName: "pi-langfuse-test",
+        metadata: {
+          source: "pi-langfuse",
+          command: "langfuse-test",
+        },
+      },
+      () => {
+        const observation = rt.startObservation(
+          "pi-langfuse-test",
+          {
+            input: { command: "/langfuse-test" },
+            output: "ok",
+            metadata: {
+              source: "pi-langfuse",
+              command: "langfuse-test",
+            },
+          },
+          { asType: "span" },
+        );
+        observation.end();
+        return observation;
+      },
+    );
+    notify(ctx, `Langfuse test succeeded. Test trace sent to ${state.config?.host}.`);
+    return true;
+  } catch (error) {
+    const message = error instanceof Error ? error.message : String(error);
+    notify(ctx, `Langfuse test failed: ${message}`, "error");
+    return false;
+  } finally {
+    if (runtimeInitialized) {
+      await (deps.forceShutdownRuntime ?? shutdownLangfuseRuntime)();
+    }
+  }
+}

package/src/config.ts

@@ -1,6 +1,7 @@
 import { mkdirSync, readFileSync, existsSync, writeFileSync } from "node:fs";
+import { dirname } from "node:path";
 import type { Config } from "./types.js";
-import { CONFIG_DIR, CONFIG_PATH, DEFAULT_LANGFUSE_HOST } from "./constants.js";
+import { CONFIG_PATH, DEFAULT_LANGFUSE_HOST } from "./constants.js";
 import { state } from "./state.js";
 import { forceShutdownRuntime } from "./langfuse.js";
 import { createCapturePolicy, type EnvLike } from "./capture-policy.js";
@@ -50,9 +51,9 @@ export function loadConfig(env: EnvLike = process.env as EnvLike, path = CONFIG_
   return loadConfigFromFile(path, env) || loadConfigFromEnv(env);
 }
 
-export function saveConfig(config: Config) {
-  mkdirSync(CONFIG_DIR, { recursive: true });
-  writeFileSync(CONFIG_PATH, `${JSON.stringify(config, null, 2)}\n`, "utf-8");
+export function saveConfig(config: Config, path = CONFIG_PATH) {
+  mkdirSync(dirname(path), { recursive: true });
+  writeFileSync(path, `${JSON.stringify(config, null, 2)}\n`, "utf-8");
 }
 
 async function collectConfigFromUI(ctx: any, reason: string): Promise<Config | null> {

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();