pi-langfuse 1.4.3 → 1.4.4

package/src/config.ts

@@ -2,18 +2,25 @@ import { mkdirSync, readFileSync, existsSync, writeFileSync } from "node:fs";
 import type { Config } from "./types.js";
 import { CONFIG_DIR, CONFIG_PATH, DEFAULT_LANGFUSE_HOST } from "./constants.js";
 import { state } from "./state.js";
-import { shutdownRuntime } from "./langfuse.js";
+import { forceShutdownRuntime } from "./langfuse.js";
+import { createCapturePolicy, type EnvLike } from "./capture-policy.js";
 
-export function loadConfigFromFile(): Config | null {
-  if (existsSync(CONFIG_PATH)) {
+export function loadConfigFromFile(path = CONFIG_PATH, env: EnvLike = process.env as EnvLike): Config | null {
+  if (existsSync(path)) {
     try {
-      const content = readFileSync(CONFIG_PATH, "utf-8");
-      const config = JSON.parse(content) as Config;
+      const content = readFileSync(path, "utf-8");
+      const config = JSON.parse(content) as Config & { capture?: EnvLike; privacyPreset?: string };
       if (config.publicKey && config.secretKey) {
+        const captureSource: EnvLike = {
+          ...(config.capture ?? {}),
+          ...(config.privacyPreset ? { LANGFUSE_PRIVACY_PRESET: config.privacyPreset } : {}),
+          ...env,
+        };
         return {
           publicKey: config.publicKey,
           secretKey: config.secretKey,
           host: config.host || DEFAULT_LANGFUSE_HOST,
+          capturePolicy: createCapturePolicy(captureSource),
         };
       }
     } catch (e) {
@@ -24,9 +31,9 @@ export function loadConfigFromFile(): Config | null {
   return null;
 }
 
-export function loadConfigFromEnv(): Config | null {
-  const publicKey = process.env.LANGFUSE_PUBLIC_KEY || "";
-  const secretKey = process.env.LANGFUSE_SECRET_KEY || "";
+export function loadConfigFromEnv(env: EnvLike = process.env as EnvLike): Config | null {
+  const publicKey = env.LANGFUSE_PUBLIC_KEY || "";
+  const secretKey = env.LANGFUSE_SECRET_KEY || "";
   if (!publicKey || !secretKey) {
     return null;
   }
@@ -34,12 +41,13 @@ export function loadConfigFromEnv(): Config | null {
   return {
     publicKey,
     secretKey,
-    host: process.env.LANGFUSE_BASE_URL || process.env.LANGFUSE_HOST || DEFAULT_LANGFUSE_HOST,
+    host: env.LANGFUSE_BASE_URL || env.LANGFUSE_HOST || DEFAULT_LANGFUSE_HOST,
+    capturePolicy: createCapturePolicy(env),
   };
 }
 
-export function loadConfig(): Config | null {
-  return loadConfigFromFile() || loadConfigFromEnv();
+export function loadConfig(env: EnvLike = process.env as EnvLike, path = CONFIG_PATH): Config | null {
+  return loadConfigFromFile(path, env) || loadConfigFromEnv(env);
 }
 
 export function saveConfig(config: Config) {
@@ -72,6 +80,7 @@ async function collectConfigFromUI(ctx: any, reason: string): Promise<Config | n
     publicKey,
     secretKey,
     host: hostInput || DEFAULT_LANGFUSE_HOST,
+    capturePolicy: createCapturePolicy(),
   };
 }
 
@@ -115,7 +124,7 @@ export async function ensureConfig(ctx: any): Promise<boolean> {
 export async function promptForConfig(ctx: any): Promise<boolean> {
   state.setupAttemptedThisSession = false;
   state.config = null;
-  await shutdownRuntime();
+  await forceShutdownRuntime();
 
   const config = await collectConfigFromUI(ctx, "Manual setup requested");
   if (!config) {

package/src/handlers/agent.ts

@@ -1,8 +1,25 @@
 import { state, resetRunState, computeEvaluationScores } from "../state.js";
 import { getRuntime, sendScore } from "../langfuse.js";
 import { ensureConfig } from "../config.js";
-import { shapePayload, truncate, extractFinalAssistant, extractAssistantOutput } from "../utils.js";
+import { shapePayload, truncate, extractFinalAssistant, extractAssistantOutput, getCapturePolicy } from "../utils.js";
 import { closeDanglingObservations } from "./tool.js";
+import { applyCapturePolicy } from "../capture-policy.js";
+
+function stringMetadata(metadata: Record<string, unknown> | undefined): Record<string, string> | undefined {
+  if (!metadata) {
+    return undefined;
+  }
+
+  const output: Record<string, string> = {};
+  for (const [key, value] of Object.entries(metadata)) {
+    if (typeof value === "string") {
+      output[key] = value;
+    } else if (typeof value === "number" || typeof value === "boolean" || typeof value === "bigint") {
+      output[key] = String(value);
+    }
+  }
+  return Object.keys(output).length > 0 ? output : undefined;
+}
 
 export function updateTraceIO(input?: unknown, output?: unknown) {
   const root = state.agentState?.root;
@@ -45,15 +62,28 @@ export async function startAgentRun(event: Record<string, unknown>, ctx: any) {
       // Ignore if getSystemPrompt is not available or fails
     }
 
-    const promptInput = shapePayload({
+    const rawPromptInput = shapePayload({
       prompt: event.prompt,
       images: event.images,
       context: event.context ?? event.attachments,
     });
+    const captured = applyCapturePolicy(
+      {
+        input: rawPromptInput,
+        metadata: {
+          cwd,
+          ...(state.currentModel ? { model: state.currentModel } : {}),
+          ...(state.currentProvider ? { provider: state.currentProvider } : {}),
+          sessionId: state.currentSessionId || undefined,
+        },
+        systemPrompt: systemPrompt ? truncate(String(systemPrompt), 20000) : undefined,
+      },
+      getCapturePolicy(),
+    );
 
     state.agentState = {
       cwd,
-      promptInput,
+      promptInput: captured.input,
       generationSeq: 0,
       activeGenerations: new Map(),
       generationOrder: [],
@@ -65,32 +95,25 @@ export async function startAgentRun(event: Record<string, unknown>, ctx: any) {
       {
         sessionId: state.currentSessionId ? truncate(state.currentSessionId, 200) : undefined,
         traceName: "pi-agent",
-        metadata: {
-          cwd: truncate(cwd, 200),
-          ...(state.currentModel ? { model: truncate(state.currentModel, 200) } : {}),
-          ...(state.currentProvider ? { provider: truncate(state.currentProvider, 200) } : {}),
-        },
+        metadata: stringMetadata(captured.metadata),
       },
       () =>
         rt.startObservation(
           "pi-agent",
-          {
-            input: promptInput,
-            metadata: {
-              cwd,
-              model: state.currentModel || undefined,
-              provider: state.currentProvider || undefined,
-              sessionId: state.currentSessionId || undefined,
-              ...(systemPrompt ? { systemPrompt: truncate(String(systemPrompt), 20000) } : {}),
+            {
+              input: captured.input,
+              metadata: {
+                ...(captured.metadata ?? {}),
+                ...(captured.systemPrompt ? { systemPrompt: captured.systemPrompt } : {}),
+              },
             },
-          },
           { asType: "agent" },
         ),
     );
 
     state.agentState.root = root;
     state.agentState.traceId = root.traceId;
-    updateTraceIO(promptInput, undefined);
+    updateTraceIO(captured.input, undefined);
   } catch (e) {
     console.warn("📊 Langfuse: Failed to create agent observation", e);
     state.isTracingDisabled = true;
@@ -104,7 +127,21 @@ export async function finishAgentRun(event: Record<string, unknown> = {}) {
   }
 
   const lastAssistant = extractFinalAssistant(event.messages);
-  const output = lastAssistant ? extractAssistantOutput(lastAssistant) : state.agentState.latestAssistantOutput;
+  const rawOutput = lastAssistant ? extractAssistantOutput(lastAssistant) : state.agentState.latestAssistantOutput;
+  const captured = applyCapturePolicy(
+    {
+      output: rawOutput,
+      metadata: {
+        cwd: state.agentState.cwd,
+        completed: true,
+        model: state.currentModel || undefined,
+        provider: state.currentProvider || undefined,
+        totalTools: state.toolCallCount,
+        ...computeEvaluationScores(),
+      },
+    },
+    getCapturePolicy(),
+  );
   const scores = computeEvaluationScores();
 
   closeDanglingObservations("Agent run ended before observation finalized");
@@ -112,18 +149,11 @@ export async function finishAgentRun(event: Record<string, unknown> = {}) {
   try {
     state.agentState.root
       .update({
-        output,
-        metadata: {
-          cwd: state.agentState.cwd,
-          completed: true,
-          model: state.currentModel || undefined,
-          provider: state.currentProvider || undefined,
-          totalTools: state.toolCallCount,
-          ...scores,
-        },
+        output: captured.output,
+        metadata: captured.metadata,
       })
       .end();
-    updateTraceIO(state.agentState.promptInput, output);
+    updateTraceIO(state.agentState.promptInput, captured.output);
 
     await sendScore("tool_call_count", scores.tool_call_count, { traceId: state.agentState.traceId });
     await sendScore("turn_count", scores.turn_count, { traceId: state.agentState.traceId });

package/src/handlers/generation.ts

@@ -1,5 +1,6 @@
 import { state } from "../state.js";
 import { getRuntime } from "../langfuse.js";
+import { startChildObservation } from "../observation.js";
 import {
   getRequestKey,
   getProviderPayload,
@@ -9,8 +10,10 @@ import {
   extractAssistantOutput,
   extractUsage,
   extractCostDetails,
+  getCapturePolicy,
 } from "../utils.js";
 import type { GenerationState, ObservationUpdate } from "../types.js";
+import { applyCapturePolicy } from "../capture-policy.js";
 
 export function getOpenGeneration(): GenerationState | undefined {
   if (state.isTracingDisabled || !state.agentState) {
@@ -44,33 +47,32 @@ export async function startGeneration(event: Record<string, unknown>) {
       url: event.url,
       method: event.method,
     }) as Record<string, unknown>;
+    const captured = applyCapturePolicy(
+      {
+        input: shapePayload(payload),
+        metadata,
+      },
+      getCapturePolicy(),
+    );
 
     const parent = state.agentState.activeTurn ?? state.agentState.root;
-    const generation = parent.startObservation
-      ? parent.startObservation(
-          "llm-generation",
-          {
-            input: shapePayload(payload),
-            model: model || undefined,
-            metadata,
-          },
-          { asType: "generation" },
-        )
-      : (await getRuntime()).startObservation(
-          "llm-generation",
-          {
-            input: shapePayload(payload),
-            model: model || undefined,
-            metadata,
-          },
-          { asType: "generation" },
-        );
+    const generation = await startChildObservation({
+      parent,
+      runtime: getRuntime,
+      name: "llm-generation",
+      body: {
+        input: captured.input,
+        model: model || undefined,
+        metadata: captured.metadata,
+      },
+      asType: "generation",
+    });
 
     state.agentState.activeGenerations.set(key, {
       observation: generation,
       requestKey: key,
       ended: false,
-      metadata,
+      metadata: captured.metadata ?? {},
     });
     state.agentState.generationOrder.push(key);
   } catch (e) {
@@ -84,7 +86,7 @@ export function updateGenerationMetadata(event: Record<string, unknown>) {
   }
 
   const key = getRequestKey(event, "");
-  const metadata = extractResponseMetadata(event);
+  const metadata = applyCapturePolicy({ metadata: extractResponseMetadata(event) }, getCapturePolicy()).metadata ?? {};
   if (!key) {
     const generation = getOpenGeneration();
     if (generation) {
@@ -160,7 +162,9 @@ export async function finishGenerationFromMessage(event: Record<string, unknown>
   }
 
   const generation = getOpenGeneration();
-  const output = extractAssistantOutput(message);
+  const rawOutput = extractAssistantOutput(message);
+  const captured = applyCapturePolicy({ output: rawOutput }, getCapturePolicy());
+  const output = captured.output;
   state.agentState.latestAssistantOutput = output;
 
   if (!generation) {
@@ -180,6 +184,7 @@ export async function finishGenerationFromMessage(event: Record<string, unknown>
       finishReason: message.finishReason ?? message.stopReason ?? event.finishReason,
     },
   };
+  update.metadata = applyCapturePolicy({ metadata: update.metadata }, getCapturePolicy()).metadata;
 
   try {
     generation.observation.update(update).end();
@@ -198,38 +203,32 @@ export async function createFallbackGenerationFromTurn(event: Record<string, unk
     const usageDetails = extractUsage({ ...event, message });
     const costDetails = extractCostDetails({ ...event, message });
     const model = String(message.model ?? event.model ?? state.currentModel ?? "");
+    const captured = applyCapturePolicy(
+      {
+        input: state.agentState.promptInput,
+        output: extractAssistantOutput(message),
+        metadata: {
+          provider: state.currentProvider || undefined,
+          sourceEvent: "turn_end",
+        },
+      },
+      getCapturePolicy(),
+    );
     const parent = state.agentState.activeTurn ?? state.agentState.root;
-    const generation = parent.startObservation
-      ? parent.startObservation(
-          "llm-generation",
-          {
-            input: state.agentState.promptInput,
-            output: extractAssistantOutput(message),
-            model: model || undefined,
-            usageDetails,
-            costDetails,
-            metadata: {
-              provider: state.currentProvider || undefined,
-              sourceEvent: "turn_end",
-            },
-          },
-          { asType: "generation" },
-        )
-      : (await getRuntime()).startObservation(
-          "llm-generation",
-          {
-            input: state.agentState.promptInput,
-            output: extractAssistantOutput(message),
-            model: model || undefined,
-            usageDetails,
-            costDetails,
-            metadata: {
-              provider: state.currentProvider || undefined,
-              sourceEvent: "turn_end",
-            },
-          },
-          { asType: "generation" },
-        );
+    const generation = await startChildObservation({
+      parent,
+      runtime: getRuntime,
+      name: "llm-generation",
+      body: {
+        input: captured.input,
+        output: captured.output,
+        model: model || undefined,
+        usageDetails,
+        costDetails,
+        metadata: captured.metadata,
+      },
+      asType: "generation",
+    });
 
     generation.end();
     state.agentState.generationOrder.push("turn-end-fallback");

package/src/handlers/tool.ts

@@ -1,5 +1,6 @@
 import { state } from "../state.js";
 import { getRuntime, sendScore } from "../langfuse.js";
+import { startChildObservation } from "../observation.js";
 import {
   getToolCallId,
   getToolName,
@@ -8,8 +9,11 @@ import {
   extractTextContent,
   truncate,
   estimatePayloadBytes,
+  getCapturePolicy,
 } from "../utils.js";
 import { MAX_TOOL_PAYLOAD_LENGTH } from "../constants.js";
+import { applyCapturePolicy } from "../capture-policy.js";
+import { redactString } from "../redaction.js";
 
 export async function startToolObservation(event: Record<string, unknown>) {
   if (state.isTracingDisabled || !state.agentState?.root) {
@@ -25,25 +29,25 @@ export async function startToolObservation(event: Record<string, unknown>) {
     const toolName = getToolName(event);
     const toolInput = getToolInput(event);
     const shapedInput = shapePayload(toolInput, { maxString: MAX_TOOL_PAYLOAD_LENGTH });
-    const inputBytes = estimatePayloadBytes(shapedInput, MAX_TOOL_PAYLOAD_LENGTH);
+    const captured = applyCapturePolicy(
+      {
+        toolInput: shapedInput,
+        metadata: { toolName, toolCallId },
+      },
+      getCapturePolicy(),
+    );
+    const inputBytes = estimatePayloadBytes(captured.toolInput, MAX_TOOL_PAYLOAD_LENGTH);
     const parent = state.agentState.activeTurn ?? state.agentState.root;
-    const tool = parent.startObservation
-      ? parent.startObservation(
-          toolName,
-          {
-            input: shapedInput,
-            metadata: { toolName, toolCallId, inputBytes },
-          },
-          { asType: "tool" },
-        )
-      : (await getRuntime()).startObservation(
-          toolName,
-          {
-            input: shapedInput,
-            metadata: { toolName, toolCallId, inputBytes },
-          },
-          { asType: "tool" },
-        );
+    const tool = await startChildObservation({
+      parent,
+      runtime: getRuntime,
+      name: toolName,
+      body: {
+        input: captured.toolInput,
+        metadata: { ...(captured.metadata ?? {}), inputBytes },
+      },
+      asType: "tool",
+    });
 
     state.toolCallCount++;
     state.agentState.activeTools.set(toolCallId, {
@@ -84,18 +88,27 @@ export async function finishToolObservation(event: Record<string, unknown>) {
 
   try {
     const shapedOutput = shapePayload(output, { maxString: MAX_TOOL_PAYLOAD_LENGTH });
-    const outputBytes = estimatePayloadBytes(shapedOutput, MAX_TOOL_PAYLOAD_LENGTH);
+    const captured = applyCapturePolicy(
+      {
+        toolOutput: shapedOutput,
+        metadata: {
+          toolName: activeTool.toolName,
+          toolCallId,
+          isError,
+        },
+      },
+      getCapturePolicy(),
+    );
+    const outputBytes = estimatePayloadBytes(captured.toolOutput, MAX_TOOL_PAYLOAD_LENGTH);
     const durationMs = Math.max(0, Date.now() - activeTool.startedAt);
 
     activeTool.observation
       .update({
-        output: shapedOutput,
+        output: captured.toolOutput,
         level: isError ? "ERROR" : "DEFAULT",
-        statusMessage: isError ? truncate(String(event.error ?? output), 1_000) : undefined,
+        statusMessage: isError ? redactString(truncate(String(event.error ?? output), 1_000)) : undefined,
         metadata: {
-          toolName: activeTool.toolName,
-          toolCallId,
-          isError,
+          ...(captured.metadata ?? {}),
           durationMs,
           inputBytes: activeTool.inputBytes,
           outputBytes,

package/src/handlers/turn.ts

@@ -1,6 +1,8 @@
 import { state } from "../state.js";
 import { getRuntime } from "../langfuse.js";
-import { shapePayload } from "../utils.js";
+import { startChildObservation } from "../observation.js";
+import { shapePayload, getCapturePolicy } from "../utils.js";
+import { applyCapturePolicy } from "../capture-policy.js";
 
 export async function startTurnObservation(event: Record<string, unknown>) {
   if (state.isTracingDisabled || !state.agentState?.root) {
@@ -15,23 +17,23 @@ export async function startTurnObservation(event: Record<string, unknown>) {
 
   try {
     const turnIndex = event.turnIndex ?? state.turnCount;
-    const observation = state.agentState.root.startObservation
-      ? state.agentState.root.startObservation(
-          "turn",
-          {
-            input: shapePayload(event.context ?? event),
-            metadata: { turnIndex },
-          },
-          { asType: "span" },
-        )
-      : (await getRuntime()).startObservation(
-          "turn",
-          {
-            input: shapePayload(event.context ?? event),
-            metadata: { turnIndex },
-          },
-          { asType: "span" },
-        );
+    const captured = applyCapturePolicy(
+      {
+        input: shapePayload(event.context ?? event),
+        metadata: { turnIndex },
+      },
+      getCapturePolicy(),
+    );
+    const observation = await startChildObservation({
+      parent: state.agentState.root,
+      runtime: getRuntime,
+      name: "turn",
+      body: {
+        input: captured.input,
+        metadata: captured.metadata,
+      },
+      asType: "span",
+    });
 
     state.agentState.activeTurn = observation;
   } catch (e) {
@@ -39,7 +41,7 @@ export async function startTurnObservation(event: Record<string, unknown>) {
   }
 }
 
-export function finishTurnObservation(event?: Record<string, unknown>) {
+export function finishTurnObservation(_event?: Record<string, unknown>) {
   if (state.isTracingDisabled || !state.agentState?.activeTurn) {
     return;
   }

package/src/langfuse.ts

@@ -3,6 +3,7 @@ import { state } from "./state.js";
 import { randomUUID } from "node:crypto";
 
 let runtime: LangfuseRuntime | null = null;
+const activeSessions = new Set<string>();
 
 type FallbackObservationType = "SPAN" | "GENERATION";
 
@@ -294,6 +295,14 @@ export async function getRuntime(): Promise<LangfuseRuntime> {
     throw new Error("Langfuse config is not set");
   }
 
+  // Track the current session as a runtime consumer.
+  // Multiple sessions can share the same runtime; shutdown is deferred
+  // until the last session releases it.
+  const sessionId = state.currentSessionId;
+  if (sessionId) {
+    activeSessions.add(sessionId);
+  }
+
   if (!runtime) {
     const [{ BasicTracerProvider }, { LangfuseSpanProcessor }, tracing, { LangfuseClient }] = await Promise.all([
       import("@opentelemetry/sdk-trace-base"),
@@ -337,23 +346,58 @@ export async function getRuntime(): Promise<LangfuseRuntime> {
   return runtime as LangfuseRuntime;
 }
 
-export async function shutdownRuntime(): Promise<void> {
-  if (!runtime) {
-    return;
-  }
+function doShutdownRuntime(): Promise<void> {
+  return (async () => {
+    if (!runtime) {
+      return;
+    }
 
-  try {
-    await runtime.tracerProvider?.forceFlush?.();
-    await fallbackToRestIngestion(runtime);
-    await runtime.scoreClient.flush?.();
-    await runtime.scoreClient.shutdown?.();
-    await runtime.tracerProvider?.shutdown?.();
-  } catch (e) {
-    console.warn("📊 Langfuse: Failed to flush/shutdown cleanly", e);
-  } finally {
-    runtime.clearTracerProvider?.();
+    const rt = runtime;
     runtime = null;
+
+    try {
+      await rt.tracerProvider?.forceFlush?.();
+      await fallbackToRestIngestion(rt);
+      await rt.scoreClient.flush?.();
+      await rt.scoreClient.shutdown?.();
+      await rt.tracerProvider?.shutdown?.();
+    } catch (e) {
+      console.warn("📊 Langfuse: Failed to flush/shutdown cleanly", e);
+    } finally {
+      if (!runtime) {
+        rt.clearTracerProvider?.();
+      }
+    }
+  })();
+}
+
+/**
+ * Release the current session's reference to the Langfuse runtime.
+ * Only actually shuts down the runtime when the last session releases it.
+ * Accepts an optional sessionId for use outside of withSession (e.g. deferred callbacks).
+ */
+export async function shutdownRuntime(sessionId?: string): Promise<void> {
+  const sid = sessionId ?? state.currentSessionId;
+  if (sid) {
+    activeSessions.delete(sid);
   }
+
+  // Still have active sessions — keep the runtime alive.
+  if (activeSessions.size > 0) {
+    return;
+  }
+
+  await doShutdownRuntime();
+}
+
+/**
+ * Force-shutdown the Langfuse runtime regardless of active session references.
+ * Used when the user manually reconfigures (e.g. /langfuse-setup) and needs
+ * a fresh runtime with new credentials.
+ */
+export async function forceShutdownRuntime(): Promise<void> {
+  activeSessions.clear();
+  await doShutdownRuntime();
 }
 
 export async function sendScore(name: string, value: number, options: { traceId?: string; observationId?: string } = {}) {

package/src/observation.ts

@@ -0,0 +1,21 @@
+import type { LangfuseObservation, LangfuseRuntime, ObservationUpdate } from "./types.js";
+
+export async function startChildObservation({
+  parent,
+  runtime,
+  name,
+  body,
+  asType,
+}: {
+  parent: LangfuseObservation;
+  runtime: () => Promise<LangfuseRuntime>;
+  name: string;
+  body?: ObservationUpdate;
+  asType: "generation" | "tool" | "span";
+}): Promise<LangfuseObservation> {
+  if (parent.startObservation) {
+    return parent.startObservation(name, body, { asType });
+  }
+
+  return (await runtime()).startObservation(name, body, { asType });
+}