pi-langfuse 1.0.0 → 1.2.0

package/package.json

@@ -1,7 +1,15 @@
 {
 	"name": "pi-langfuse",
-	"version": "1.0.0",
+	"version": "1.2.0",
 	"description": "Langfuse extension for Pi coding agent",
+	"repository": {
+		"type": "git",
+		"url": "git+https://github.com/gooyoung/pi-langfuse.git"
+	},
+	"bugs": {
+		"url": "https://github.com/gooyoung/pi-langfuse/issues"
+	},
+	"homepage": "https://github.com/gooyoung/pi-langfuse#readme",
 	"type": "module",
 	"main": "index.ts",
 	"scripts": {
@@ -19,10 +27,14 @@
 	"pi": {
 		"extensions": [
 			"./index.ts"
-		]
+		],
+		"image": "https://github.com/gooyoung/pi-langfuse/raw/main/media/image.png"
 	},
 	"dependencies": {
-		"langfuse": "^3.0.0"
+		"@langfuse/client": "^5.3.0",
+		"@langfuse/otel": "^5.3.0",
+		"@langfuse/tracing": "^5.3.0",
+		"@opentelemetry/sdk-node": "^0.218.0"
 	},
 	"peerDependencies": {
 		"@earendil-works/pi-coding-agent": "*"
@@ -34,5 +46,8 @@
 	"license": "MIT",
 	"engines": {
 		"node": ">=22"
+	},
+	"devDependencies": {
+		"typescript": "^6.0.3"
 	}
 }

package/src/config.ts

@@ -0,0 +1,98 @@
+import { readFileSync, existsSync, writeFileSync } from "node:fs";
+import type { Config } from "./types.js";
+import { CONFIG_PATH, DEFAULT_LANGFUSE_HOST } from "./constants.js";
+import { state } from "./state.js";
+import { shutdownRuntime } from "./langfuse.js";
+
+export function loadConfigFromFile(): Config | null {
+  if (existsSync(CONFIG_PATH)) {
+    try {
+      const content = readFileSync(CONFIG_PATH, "utf-8");
+      const config = JSON.parse(content) as Config;
+      if (config.publicKey && config.secretKey) {
+        return {
+          publicKey: config.publicKey,
+          secretKey: config.secretKey,
+          host: config.host || DEFAULT_LANGFUSE_HOST,
+        };
+      }
+    } catch (e) {
+      console.warn("📊 Langfuse: Failed to load config.json", e);
+    }
+  }
+
+  return null;
+}
+
+export function loadConfigFromEnv(): Config | null {
+  const publicKey = process.env.LANGFUSE_PUBLIC_KEY || "";
+  const secretKey = process.env.LANGFUSE_SECRET_KEY || "";
+  if (!publicKey || !secretKey) {
+    return null;
+  }
+
+  return {
+    publicKey,
+    secretKey,
+    host: process.env.LANGFUSE_BASE_URL || process.env.LANGFUSE_HOST || DEFAULT_LANGFUSE_HOST,
+  };
+}
+
+export function saveConfig(config: Config) {
+  writeFileSync(CONFIG_PATH, `${JSON.stringify(config, null, 2)}\n`, "utf-8");
+}
+
+export async function ensureConfig(ctx: any): Promise<boolean> {
+  if (state.config) {
+    return true;
+  }
+
+  if (state.setupAttemptedThisSession) {
+    return false;
+  }
+  state.setupAttemptedThisSession = true;
+
+  if (!ctx.hasUI) {
+    console.log("📊 Langfuse: Missing config. Run this extension in Pi UI to complete setup, or set LANGFUSE_PUBLIC_KEY / LANGFUSE_SECRET_KEY / LANGFUSE_BASE_URL.");
+    return false;
+  }
+
+  ctx.ui.notify("Langfuse setup required. Enter your API keys to enable tracing.", "info");
+
+  const publicKey = (await ctx.ui.input("Langfuse public key:", "pk-lf-..."))?.trim();
+  if (!publicKey) {
+    ctx.ui.notify("Langfuse setup cancelled.", "warning");
+    return false;
+  }
+
+  const secretKey = (await ctx.ui.input("Langfuse secret key:", "sk-lf-..."))?.trim();
+  if (!secretKey) {
+    ctx.ui.notify("Langfuse setup cancelled.", "warning");
+    return false;
+  }
+
+  const hostInput = (await ctx.ui.input("Langfuse host:", DEFAULT_LANGFUSE_HOST))?.trim();
+  state.config = {
+    publicKey,
+    secretKey,
+    host: hostInput || DEFAULT_LANGFUSE_HOST,
+  };
+
+  try {
+    saveConfig(state.config);
+    ctx.ui.notify(`Langfuse config saved to ${CONFIG_PATH}`, "info");
+    return true;
+  } catch (error) {
+    console.warn("📊 Langfuse: Failed to save config.json", error);
+    ctx.ui.notify("Failed to save Langfuse config.json. Check extension directory permissions.", "error");
+    state.config = null;
+    return false;
+  }
+}
+
+export async function promptForConfig(ctx: any): Promise<boolean> {
+  state.setupAttemptedThisSession = false;
+  state.config = null;
+  await shutdownRuntime();
+  return ensureConfig(ctx);
+}

package/src/constants.ts

@@ -0,0 +1,15 @@
+import { resolve, dirname } from "node:path";
+import { fileURLToPath } from "node:url";
+
+// Since we are in src/constants.ts, EXT_DIR should point to the parent of src/ (the root of the extension)
+// We use `import.meta.url` which points to `src/constants.ts`.
+// `dirname` gives us `src/`, and `resolve(..., '..')` gives us the root.
+export const EXT_DIR = resolve(dirname(fileURLToPath(import.meta.url)), "..");
+export const CONFIG_PATH = resolve(EXT_DIR, "config.json");
+export const DEFAULT_LANGFUSE_HOST = "https://cloud.langfuse.com";
+
+export const MAX_STRING_LENGTH = 12_000;
+export const MAX_TOOL_PAYLOAD_LENGTH = 24_000;
+export const MAX_DEPTH = 6;
+export const MAX_ARRAY_ITEMS = 50;
+export const MAX_OBJECT_KEYS = 80;

package/src/handlers/agent.ts

@@ -0,0 +1,136 @@
+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 { closeDanglingObservations } from "./tool.js";
+
+export function updateTraceIO(input?: unknown, output?: unknown) {
+  const root = state.agentState?.root;
+  if (!root?.setTraceIO) {
+    return;
+  }
+
+  try {
+    root.setTraceIO({ input, output });
+  } catch {
+    // Older SDKs may omit setTraceIO; root IO still mirrors trace IO in current Langfuse.
+  }
+}
+
+export async function startAgentRun(event: Record<string, unknown>, ctx: any) {
+  if (!(await ensureConfig(ctx))) {
+    return;
+  }
+
+  try {
+    const rt = await getRuntime();
+    const cwd = String(
+      (event.systemPromptOptions && typeof event.systemPromptOptions === "object"
+        ? (event.systemPromptOptions as Record<string, unknown>).cwd
+        : undefined) ?? process.cwd(),
+    );
+
+    if (!state.currentModel && ctx.model) {
+      state.currentModel = ctx.model.id || "";
+      state.currentProvider = ctx.model.provider || "";
+    }
+
+    let systemPrompt = undefined;
+    try {
+      if (ctx.getSystemPrompt) {
+        systemPrompt = await ctx.getSystemPrompt();
+      }
+    } catch {
+      // Ignore if getSystemPrompt is not available or fails
+    }
+
+    const promptInput = shapePayload({
+      prompt: event.prompt,
+      images: event.images,
+      context: event.context ?? event.attachments,
+    });
+
+    state.agentState = {
+      cwd,
+      promptInput,
+      generationSeq: 0,
+      activeGenerations: new Map(),
+      generationOrder: [],
+      activeTools: new Map(),
+      providerMetadataByRequest: new Map(),
+    };
+
+    const root = rt.propagateAttributes(
+      {
+        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) } : {}),
+        },
+      },
+      () =>
+        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) } : {}),
+            },
+          },
+          { asType: "agent" },
+        ),
+    );
+
+    state.agentState.root = root;
+    state.agentState.traceId = root.traceId;
+    updateTraceIO(promptInput, undefined);
+  } catch (e) {
+    console.warn("📊 Langfuse: Failed to create agent observation", e);
+  }
+}
+
+export async function finishAgentRun(event: Record<string, unknown> = {}) {
+  if (!state.agentState?.root) {
+    resetRunState();
+    return;
+  }
+
+  const lastAssistant = extractFinalAssistant(event.messages);
+  const output = lastAssistant ? extractAssistantOutput(lastAssistant) : state.agentState.latestAssistantOutput;
+  const scores = computeEvaluationScores();
+
+  closeDanglingObservations("Agent run ended before observation finalized");
+
+  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,
+        },
+      })
+      .end();
+    updateTraceIO(state.agentState.promptInput, 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 });
+    await sendScore("total_tool_errors", scores.total_tool_errors, { traceId: state.agentState.traceId });
+    await sendScore("tool_success_rate", scores.tool_success_rate, { traceId: state.agentState.traceId });
+    await sendScore("session_had_errors", scores.session_had_errors, { traceId: state.agentState.traceId });
+  } catch (e) {
+    console.warn("📊 Langfuse: Failed to finish agent observation", e);
+  } finally {
+    resetRunState();
+  }
+}

package/src/handlers/generation.ts

@@ -0,0 +1,239 @@
+import { state } from "../state.js";
+import { getRuntime } from "../langfuse.js";
+import {
+  getRequestKey,
+  getProviderPayload,
+  shapePayload,
+  extractResponseMetadata,
+  getMessageFromEvent,
+  extractAssistantOutput,
+  extractUsage,
+  extractCostDetails,
+} from "../utils.js";
+import type { GenerationState, ObservationUpdate } from "../types.js";
+
+export function getOpenGeneration(): GenerationState | undefined {
+  if (!state.agentState) {
+    return undefined;
+  }
+
+  for (let i = state.agentState.generationOrder.length - 1; i >= 0; i--) {
+    const key = state.agentState.generationOrder[i];
+    const genState = state.agentState.activeGenerations.get(key);
+    if (genState && !genState.ended) {
+      return genState;
+    }
+  }
+
+  return undefined;
+}
+
+export async function startGeneration(event: Record<string, unknown>) {
+  if (!state.agentState?.root) {
+    return;
+  }
+
+  try {
+    const key = getRequestKey(event, `generation-${++state.agentState.generationSeq}`);
+    const payload = getProviderPayload(event);
+    const model = String(event.model ?? event.modelId ?? state.currentModel ?? "");
+    const provider = String(event.provider ?? state.currentProvider ?? "");
+    const metadata = shapePayload({
+      provider,
+      requestId: key,
+      url: event.url,
+      method: event.method,
+    }) as Record<string, unknown>;
+
+    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" },
+        );
+
+    state.agentState.activeGenerations.set(key, {
+      observation: generation,
+      requestKey: key,
+      ended: false,
+      metadata,
+    });
+    state.agentState.generationOrder.push(key);
+  } catch (e) {
+    console.warn("📊 Langfuse: Failed to start generation", e);
+  }
+}
+
+export function updateGenerationMetadata(event: Record<string, unknown>) {
+  if (!state.agentState) {
+    return;
+  }
+
+  const key = getRequestKey(event, "");
+  const metadata = extractResponseMetadata(event);
+  if (!key) {
+    const generation = getOpenGeneration();
+    if (generation) {
+      generation.metadata = { ...generation.metadata, ...metadata };
+      
+      const isError = 
+        (typeof metadata.status === "number" && metadata.status >= 400) || 
+        event.error || 
+        event.isError;
+        
+      if (isError) {
+        generation.observation.update({ 
+          metadata: generation.metadata,
+          level: "ERROR",
+          statusMessage: String(event.error ?? metadata.statusMessage ?? "Provider request failed")
+        }).end();
+        generation.ended = true;
+      } else {
+        generation.observation.update({ metadata: generation.metadata });
+      }
+    }
+    return;
+  }
+
+  const generation = state.agentState.activeGenerations.get(key) ?? getOpenGeneration();
+  if (generation) {
+    generation.metadata = { ...generation.metadata, ...metadata };
+    
+    const isError = 
+      (typeof metadata.status === "number" && metadata.status >= 400) || 
+      event.error || 
+      event.isError;
+      
+    if (isError) {
+      generation.observation.update({ 
+        metadata: generation.metadata,
+        level: "ERROR",
+        statusMessage: String(event.error ?? metadata.statusMessage ?? "Provider request failed")
+      }).end();
+      generation.ended = true;
+    } else {
+      generation.observation.update({ metadata: generation.metadata });
+    }
+  }
+}
+
+export function recordTTFT(event: Record<string, unknown>) {
+  if (!state.agentState) {
+    return;
+  }
+
+  const key = getRequestKey(event, "");
+  const generation = key ? state.agentState.activeGenerations.get(key) : getOpenGeneration();
+  
+  if (generation && !generation.ttftRecorded && !generation.ended) {
+    generation.ttftRecorded = true;
+    try {
+      generation.observation.update({ completionStartTime: new Date() });
+    } catch (e) {
+      // Ignore
+    }
+  }
+}
+
+export async function finishGenerationFromMessage(event: Record<string, unknown>) {
+  if (!state.agentState) {
+    return;
+  }
+
+  const message = getMessageFromEvent(event);
+  if (!message || message.role !== "assistant") {
+    return;
+  }
+
+  const generation = getOpenGeneration();
+  const output = extractAssistantOutput(message);
+  state.agentState.latestAssistantOutput = output;
+
+  if (!generation) {
+    return;
+  }
+
+  const usageDetails = extractUsage({ ...event, message });
+  const costDetails = extractCostDetails({ ...event, message });
+  const model = String(message.model ?? event.model ?? state.currentModel ?? "");
+  const update: ObservationUpdate = {
+    output,
+    model: model || undefined,
+    usageDetails,
+    costDetails,
+    metadata: {
+      ...generation.metadata,
+      finishReason: message.finishReason ?? message.stopReason ?? event.finishReason,
+    },
+  };
+
+  try {
+    generation.observation.update(update).end();
+    generation.ended = true;
+  } catch (e) {
+    console.warn("📊 Langfuse: Failed to finish generation", e);
+  }
+}
+
+export async function createFallbackGenerationFromTurn(event: Record<string, unknown>, message: Record<string, unknown>) {
+  if (!state.agentState?.root || state.agentState.generationOrder.length > 0) {
+    return;
+  }
+
+  try {
+    const usageDetails = extractUsage({ ...event, message });
+    const costDetails = extractCostDetails({ ...event, message });
+    const model = String(message.model ?? event.model ?? state.currentModel ?? "");
+    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" },
+        );
+
+    generation.end();
+    state.agentState.generationOrder.push("turn-end-fallback");
+  } catch (e) {
+    console.warn("📊 Langfuse: Failed to create fallback generation", e);
+  }
+}

package/src/handlers/tool.ts

@@ -0,0 +1,126 @@
+import { state } from "../state.js";
+import { getRuntime, sendScore } from "../langfuse.js";
+import {
+  getToolCallId,
+  getToolName,
+  getToolInput,
+  shapePayload,
+  extractTextContent,
+  truncate,
+} from "../utils.js";
+import { MAX_TOOL_PAYLOAD_LENGTH } from "../constants.js";
+
+export async function startToolObservation(event: Record<string, unknown>) {
+  if (!state.agentState?.root) {
+    return;
+  }
+
+  const toolCallId = getToolCallId(event);
+  if (!toolCallId || state.agentState.activeTools.has(toolCallId)) {
+    return;
+  }
+
+  try {
+    const toolName = getToolName(event);
+    const parent = state.agentState.activeTurn ?? state.agentState.root;
+    const tool = parent.startObservation
+      ? parent.startObservation(
+          toolName,
+          {
+            input: shapePayload(getToolInput(event), { maxString: MAX_TOOL_PAYLOAD_LENGTH }),
+            metadata: { toolName, toolCallId },
+          },
+          { asType: "tool" },
+        )
+      : (await getRuntime()).startObservation(
+          toolName,
+          {
+            input: shapePayload(getToolInput(event), { maxString: MAX_TOOL_PAYLOAD_LENGTH }),
+            metadata: { toolName, toolCallId },
+          },
+          { asType: "tool" },
+        );
+
+    state.toolCallCount++;
+    state.agentState.activeTools.set(toolCallId, { observation: tool, toolName, ended: false });
+  } catch (e) {
+    console.warn("📊 Langfuse: Failed to start tool observation", e);
+  }
+}
+
+export async function finishToolObservation(event: Record<string, unknown>) {
+  if (!state.agentState) {
+    return;
+  }
+
+  const toolCallId = getToolCallId(event);
+  if (!toolCallId) {
+    return;
+  }
+
+  const activeTool = state.agentState.activeTools.get(toolCallId);
+  if (!activeTool || activeTool.ended) {
+    return;
+  }
+
+  const isError = Boolean(event.isError ?? event.error ?? event.status === "error");
+  const output =
+    extractTextContent(event.content, MAX_TOOL_PAYLOAD_LENGTH) ??
+    event.output ??
+    event.result ??
+    event.error ??
+    event.content ??
+    event;
+
+  try {
+    activeTool.observation
+      .update({
+        output: shapePayload(output, { maxString: MAX_TOOL_PAYLOAD_LENGTH }),
+        level: isError ? "ERROR" : "DEFAULT",
+        statusMessage: isError ? truncate(String(event.error ?? output), 1_000) : undefined,
+        metadata: {
+          toolName: activeTool.toolName,
+          toolCallId,
+          isError,
+        },
+      })
+      .end();
+    activeTool.ended = true;
+
+    if (isError) {
+      state.errorCount++;
+      await sendScore("tool_is_error", 1, {
+        traceId: state.agentState.traceId,
+        observationId: activeTool.observation.id,
+      });
+    }
+  } catch (e) {
+    console.warn("📊 Langfuse: Failed to finish tool observation", e);
+  } finally {
+    state.agentState.activeTools.delete(toolCallId);
+  }
+}
+
+export function closeDanglingObservations(statusMessage: string) {
+  if (!state.agentState) {
+    return;
+  }
+
+  for (const activeTool of state.agentState.activeTools.values()) {
+    if (!activeTool.ended) {
+      activeTool.observation
+        .update({ level: "WARNING", statusMessage, metadata: { toolName: activeTool.toolName, cancelled: true } })
+        .end();
+      activeTool.ended = true;
+    }
+  }
+
+  for (const generation of state.agentState.activeGenerations.values()) {
+    if (!generation.ended) {
+      generation.observation.update({ level: "WARNING", statusMessage, metadata: { ...generation.metadata, cancelled: true } }).end();
+      generation.ended = true;
+    }
+  }
+
+  state.agentState.activeTools.clear();
+}