pi-langfuse 1.4.2 → 1.4.4

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";
 
@@ -42,7 +43,8 @@ interface RestFallbackStore {
   attempted: boolean;
 }
 
-const OTEL_VISIBILITY_DELAY_MS = 1_500;
+const OTEL_VISIBILITY_TIMEOUT_MS = 1_500;
+const OTEL_VISIBILITY_POLL_INTERVAL_MS = 200;
 
 function nowIso() {
   return new Date().toISOString();
@@ -190,6 +192,23 @@ async function traceExists(rt: LangfuseRuntime, traceId: string): Promise<boolea
   }
 }
 
+async function waitForTraceVisibility(rt: LangfuseRuntime, traceId: string): Promise<boolean> {
+  const deadline = Date.now() + OTEL_VISIBILITY_TIMEOUT_MS;
+
+  while (true) {
+    if (await traceExists(rt, traceId)) {
+      return true;
+    }
+
+    const remainingMs = deadline - Date.now();
+    if (remainingMs <= 0) {
+      return false;
+    }
+
+    await delay(Math.min(OTEL_VISIBILITY_POLL_INTERVAL_MS, remainingMs));
+  }
+}
+
 function eventTimestamp(record: { endTime?: string; startTime?: string; timestamp?: string }) {
   return record.endTime ?? record.startTime ?? record.timestamp ?? nowIso();
 }
@@ -201,8 +220,7 @@ async function fallbackToRestIngestion(rt: LangfuseRuntime) {
   }
   store.attempted = true;
 
-  await delay(OTEL_VISIBILITY_DELAY_MS);
-  if (await traceExists(rt, store.trace.id)) {
+  if (await waitForTraceVisibility(rt, store.trace.id)) {
     return;
   }
 
@@ -277,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"),
@@ -320,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

@@ -8,6 +8,8 @@ export interface SessionRunState {
   toolCallCount: number;
   errorCount: number;
   turnCount: number;
+  tracingDisabled: boolean;
+  setupAttemptedThisSession: boolean;
 }
 
 const DEFAULT_SESSION_ID = "__pi_langfuse_default_session__";
@@ -23,6 +25,8 @@ function createSessionRunState(): SessionRunState {
     toolCallCount: 0,
     errorCount: 0,
     turnCount: 0,
+    tracingDisabled: false,
+    setupAttemptedThisSession: false,
   };
 }
 
@@ -57,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() {
@@ -109,10 +112,30 @@ export const state = {
   set turnCount(turnCount: number) {
     getSessionRunState().turnCount = turnCount;
   },
+
+  get isTracingDisabled() {
+    return getSessionRunState().tracingDisabled;
+  },
+  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

@@ -2,9 +2,17 @@ import {
   MAX_ARRAY_ITEMS,
   MAX_DEPTH,
   MAX_OBJECT_KEYS,
+  MAX_PAYLOAD_NODES,
   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;
@@ -23,11 +31,28 @@ export function tryParseJson(value: string): unknown {
   }
 }
 
-export function shapePayload(value: unknown, options: { maxString?: number; depth?: number } = {}): unknown {
+const PAYLOAD_TOO_LARGE = "[payload too large]";
+
+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;
+  const budget = { exhausted: false, nodeCount: 0 };
 
   function visit(item: unknown, remainingDepth: number, seen: WeakSet<object>): unknown {
+    if (budget.exhausted) {
+      return PAYLOAD_TOO_LARGE;
+    }
+
+    budget.nodeCount++;
+    if (budget.nodeCount > maxNodes) {
+      budget.exhausted = true;
+      return PAYLOAD_TOO_LARGE;
+    }
+
     if (typeof item === "string") {
       const truncated = truncate(item, maxString);
       const parsed = tryParseJson(truncated);
@@ -59,7 +84,15 @@ export function shapePayload(value: unknown, options: { maxString?: number; dept
     }
 
     if (Array.isArray(item)) {
-      return item.slice(0, MAX_ARRAY_ITEMS).map((entry) => visit(entry, remainingDepth - 1, seen));
+      const output: unknown[] = [];
+      const limit = Math.min(item.length, MAX_ARRAY_ITEMS);
+      for (let index = 0; index < limit; index++) {
+        output.push(visit(item[index], remainingDepth - 1, seen));
+        if (budget.exhausted) {
+          break;
+        }
+      }
+      return output;
     }
 
     if (item instanceof Error) {
@@ -77,8 +110,16 @@ export function shapePayload(value: unknown, options: { maxString?: number; dept
       seen.add(item);
 
       const output: Record<string, unknown> = {};
-      for (const [key, entry] of Object.entries(item as Record<string, unknown>).slice(0, MAX_OBJECT_KEYS)) {
-        output[key] = visit(entry, remainingDepth - 1, seen);
+      let keyCount = 0;
+      for (const key in item as Record<string, unknown>) {
+        if (!Object.hasOwn(item, key)) {
+          continue;
+        }
+        output[key] = visit((item as Record<string, unknown>)[key], remainingDepth - 1, seen);
+        keyCount++;
+        if (budget.exhausted || keyCount >= MAX_OBJECT_KEYS) {
+          break;
+        }
       }
       return output;
     }
@@ -86,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 {

package/types/langfuse-runtime-shims.d.ts

@@ -0,0 +1,49 @@
+declare module "@opentelemetry/sdk-trace-base" {
+  export class BasicTracerProvider {
+    constructor(options?: { spanProcessors?: unknown[] });
+    forceFlush?(): Promise<void>;
+    shutdown?(): Promise<void>;
+  }
+}
+
+declare module "@langfuse/otel" {
+  export class LangfuseSpanProcessor {
+    constructor(options: {
+      publicKey: string;
+      secretKey: string;
+      baseUrl: string;
+    });
+    forceFlush?(): Promise<void>;
+    shutdown?(): Promise<void>;
+  }
+}
+
+declare module "@langfuse/tracing" {
+  export function setLangfuseTracerProvider(provider: unknown): void;
+
+  export function startObservation(
+    name: string,
+    body?: Record<string, unknown>,
+    options?: { asType?: string },
+  ): unknown;
+
+  export function propagateAttributes<T>(
+    params: {
+      sessionId?: string;
+      traceName?: string;
+      metadata?: Record<string, string>;
+      tags?: string[];
+    },
+    fn: () => T,
+  ): T;
+}
+
+declare module "@langfuse/client" {
+  export class LangfuseClient {
+    constructor(options: {
+      publicKey: string;
+      secretKey: string;
+      baseUrl: string;
+    });
+  }
+}