pi-langfuse 1.4.3 → 1.4.5

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

package/src/redaction.ts

@@ -0,0 +1,115 @@
+import { createHash } from "node:crypto";
+
+export const REDACTED = "[REDACTED_SECRET]";
+
+export interface RedactOptions {
+  maxDepth: number;
+  maxArrayItems: number;
+  maxObjectKeys: number;
+  maxStringLength: number;
+}
+
+const DEFAULT_OPTIONS: RedactOptions = {
+  maxDepth: 6,
+  maxArrayItems: 50,
+  maxObjectKeys: 80,
+  maxStringLength: 12_000,
+};
+
+const SECRET_ASSIGNMENT_RE =
+  /\b([A-Z0-9_]*(?:SECRET|TOKEN|PASSWORD|PASS|API[_-]?KEY|PRIVATE[_-]?KEY|AUTH|COOKIE)[A-Z0-9_]*)\s*=\s*([^\s"'`]+)/gi;
+const PRIVATE_KEY_RE = /-----BEGIN [A-Z ]*PRIVATE KEY-----[\s\S]*?-----END [A-Z ]*PRIVATE KEY-----/g;
+const BEARER_RE = /\bBearer\s+[A-Za-z0-9._~+/=-]{12,}/gi;
+const KNOWN_TOKEN_RE =
+  /\b(?:sk-(?:lf|ant|proj|live|test)[A-Za-z0-9_-]*|pk-lf-[A-Za-z0-9_-]+|gh[pousr]_[A-Za-z0-9_]{20,}|npm_[A-Za-z0-9_-]{20,}|AKIA[0-9A-Z]{16})\b/g;
+const ABSOLUTE_PATH_RE =
+  /(?:\/Users\/[^/\s]+|\/home\/[^/\s]+|\/private\/tmp|\/tmp|[A-Za-z]:\\Users\\[^\\\s]+)(?:[^\s"'`]*)/g;
+const SENSITIVE_FIELD_RE =
+  /^(authorization|cookie|setcookie|xapikey|apikey|token|accesstoken|refreshtoken|secret|secretkey|password|passwd|privatekey)$/;
+
+export function hashPath(path: string): string {
+  return `[PATH_HASH:${createHash("sha256").update(path).digest("hex").slice(0, 12)}]`;
+}
+
+function truncate(value: string, maxStringLength: number): string {
+  return value.length > maxStringLength ? `${value.slice(0, maxStringLength)}... [truncated]` : value;
+}
+
+export function redactString(value: string, options: Partial<RedactOptions> = {}): string {
+  const merged = { ...DEFAULT_OPTIONS, ...options };
+  const truncated = truncate(value, merged.maxStringLength);
+  return truncated
+    .replace(PRIVATE_KEY_RE, REDACTED)
+    .replace(BEARER_RE, REDACTED)
+    .replace(KNOWN_TOKEN_RE, REDACTED)
+    .replace(SECRET_ASSIGNMENT_RE, (_match, key: string) => `${key}=${REDACTED}`)
+    .replace(ABSOLUTE_PATH_RE, (path: string) => {
+      const envSuffix = path.match(/([/\\]\.env(?:\.[A-Za-z0-9_-]+)?)$/)?.[1];
+      return `${hashPath(envSuffix ? path.slice(0, -envSuffix.length) : path)}${envSuffix ?? ""}`;
+    });
+}
+
+function visit(value: unknown, options: RedactOptions, depth: number, seen: WeakSet<object>): unknown {
+  if (value === null || value === undefined || typeof value === "number" || typeof value === "boolean") {
+    return value;
+  }
+
+  if (typeof value === "bigint") {
+    return value.toString();
+  }
+
+  if (typeof value === "string") {
+    return redactString(value, options);
+  }
+
+  if (typeof value === "function" || typeof value === "symbol") {
+    return `[${typeof value}]`;
+  }
+
+  if (depth <= 0) {
+    return `[max depth ${options.maxDepth} reached]`;
+  }
+
+  if (value instanceof Error) {
+    return {
+      name: redactString(value.name, options),
+      message: redactString(value.message, options),
+      stack: value.stack ? redactString(value.stack, options) : undefined,
+    };
+  }
+
+  if (typeof value !== "object") {
+    return redactString(String(value), options);
+  }
+
+  if (seen.has(value)) {
+    return "[circular]";
+  }
+  seen.add(value);
+
+  if (Array.isArray(value)) {
+    const output = value
+      .slice(0, options.maxArrayItems)
+      .map((item) => visit(item, options, depth - 1, seen));
+    if (value.length > options.maxArrayItems) {
+      output.push(`[${value.length - options.maxArrayItems} truncated items]`);
+    }
+    return output;
+  }
+
+  const entries = Object.entries(value as Record<string, unknown>);
+  const output: Record<string, unknown> = {};
+  for (const [key, item] of entries.slice(0, options.maxObjectKeys)) {
+    const normalizedKey = key.replace(/[^a-zA-Z0-9]/g, "").toLowerCase();
+    output[key] = SENSITIVE_FIELD_RE.test(normalizedKey) ? REDACTED : visit(item, options, depth - 1, seen);
+  }
+  if (entries.length > options.maxObjectKeys) {
+    output.__truncatedKeys = entries.length - options.maxObjectKeys;
+  }
+  return output;
+}
+
+export function redactValue(value: unknown, options: Partial<RedactOptions> = {}): unknown {
+  const merged: RedactOptions = { ...DEFAULT_OPTIONS, ...options };
+  return visit(value, merged, merged.maxDepth, new WeakSet<object>());
+}

package/src/state.ts

@@ -9,6 +9,7 @@ export interface SessionRunState {
   errorCount: number;
   turnCount: number;
   tracingDisabled: boolean;
+  setupAttemptedThisSession: boolean;
 }
 
 const DEFAULT_SESSION_ID = "__pi_langfuse_default_session__";
@@ -25,6 +26,7 @@ function createSessionRunState(): SessionRunState {
     errorCount: 0,
     turnCount: 0,
     tracingDisabled: false,
+    setupAttemptedThisSession: false,
   };
 }
 
@@ -59,7 +61,6 @@ export function runWithSession<T>(sessionId: string | undefined, fn: () => T): T
 
 export const state = {
   config: null as Config | null,
-  setupAttemptedThisSession: false,
   sessionStates: new Map<string, SessionRunState>(),
 
   get currentSessionId() {
@@ -118,10 +119,23 @@ export const state = {
   set isTracingDisabled(disabled: boolean) {
     getSessionRunState().tracingDisabled = disabled;
   },
+
+  get setupAttemptedThisSession() {
+    return getSessionRunState().setupAttemptedThisSession;
+  },
+  set setupAttemptedThisSession(attempted: boolean) {
+    getSessionRunState().setupAttemptedThisSession = attempted;
+  },
 };
 
 export function resetRunState(sessionId = getActiveSessionId()) {
-  state.sessionStates.set(normalizeSessionId(sessionId), createSessionRunState());
+  const normalizedSessionId = normalizeSessionId(sessionId);
+  const setupAttemptedThisSession =
+    state.sessionStates.get(normalizedSessionId)?.setupAttemptedThisSession ?? false;
+  state.sessionStates.set(normalizedSessionId, {
+    ...createSessionRunState(),
+    setupAttemptedThisSession,
+  });
 }
 
 export function clearAllSessionStates() {

package/src/types.ts

@@ -1,7 +1,10 @@
+import type { CapturePolicy } from "./capture-policy.js";
+
 export interface Config {
   publicKey: string;
   secretKey: string;
   host: string;
+  capturePolicy?: CapturePolicy;
 }
 
 export interface LangfuseObservation {

package/src/utils.ts

@@ -6,6 +6,13 @@ import {
   MAX_STRING_LENGTH,
   MAX_TOOL_PAYLOAD_LENGTH,
 } from "./constants.js";
+import { createCapturePolicy, type CapturePolicy } from "./capture-policy.js";
+import { redactValue } from "./redaction.js";
+import { state } from "./state.js";
+
+export function getCapturePolicy(): CapturePolicy {
+  return state.config?.capturePolicy ?? createCapturePolicy();
+}
 
 export function truncate(value: string, maxLength = MAX_STRING_LENGTH): string {
   return value.length > maxLength ? `${value.slice(0, maxLength)}... [truncated]` : value;
@@ -26,7 +33,10 @@ export function tryParseJson(value: string): unknown {
 
 const PAYLOAD_TOO_LARGE = "[payload too large]";
 
-export function shapePayload(value: unknown, options: { maxString?: number; depth?: number; maxNodes?: number } = {}): unknown {
+export function shapePayload(
+  value: unknown,
+  options: { maxString?: number; depth?: number; maxNodes?: number; redact?: boolean } = {},
+): unknown {
   const maxString = options.maxString ?? MAX_STRING_LENGTH;
   const depth = options.depth ?? MAX_DEPTH;
   const maxNodes = options.maxNodes ?? MAX_PAYLOAD_NODES;
@@ -117,7 +127,15 @@ export function shapePayload(value: unknown, options: { maxString?: number; dept
     return String(item);
   }
 
-  return visit(value, depth, new WeakSet<object>());
+  const shaped = visit(value, depth, new WeakSet<object>());
+  return options.redact === false
+    ? shaped
+    : redactValue(shaped, {
+        maxDepth: depth,
+        maxStringLength: maxString,
+        maxArrayItems: MAX_ARRAY_ITEMS,
+        maxObjectKeys: MAX_OBJECT_KEYS,
+      });
 }
 
 export function safeSerialize(value: unknown, maxLength = MAX_TOOL_PAYLOAD_LENGTH): string {