pi-langfuse 1.4.3 → 1.4.5

package/src/commands.ts

@@ -0,0 +1,261 @@
+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;
+  };
+}
+
+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;
+}
+
+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();
+    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;
+  }
+
+  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 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;
+      },
+    );
+    await rt.tracerProvider?.forceFlush?.();
+    await rt.scoreClient.flush?.();
+    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,19 +1,27 @@
 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 { 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 +32,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,17 +42,18 @@ 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) {
-  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> {
@@ -72,6 +81,7 @@ async function collectConfigFromUI(ctx: any, reason: string): Promise<Config | n
     publicKey,
     secretKey,
     host: hostInput || DEFAULT_LANGFUSE_HOST,
+    capturePolicy: createCapturePolicy(),
   };
 }
 
@@ -115,7 +125,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");