@narumitw/pi-langfuse 0.18.0 → 0.26.0

package/src/langfuse.ts

@@ -1,6 +1,7 @@
 import type {
 	ExtensionAPI,
 	ExtensionCommandContext,
+	ExtensionContext,
 } from "@earendil-works/pi-coding-agent";
 import {
 	DEFAULT_BASE_URL,
@@ -10,20 +11,24 @@ import {
 	normalizeLangfuseConfig,
 	writeLangfuseConfig,
 } from "./config.js";
-import { type TraceBackend, TraceRecorder } from "./tracing.js";
+import {
+	type ContextSnapshot,
+	type GitMetadata,
+	type TraceBackend,
+	TraceRecorder,
+} from "./tracing.js";
 
 interface ExtensionDependencies {
 	loadConfig(path?: string): Promise<LangfuseConfigResult>;
 	writeConfig(config: LangfuseConfig, path?: string): Promise<LangfuseConfig>;
 	createBackend(config: LangfuseConfig): Promise<TraceBackend>;
+	resolveGitMetadata(cwd: string): Promise<GitMetadata | undefined>;
 }
 
-const COMMAND_COMPLETIONS = [
-	{ value: "status", label: "status", description: "Show Langfuse tracing status" },
-	{ value: "flush", label: "flush", description: "Export all completed traces now" },
-	{ value: "help", label: "help", description: "Show Langfuse command help" },
-	{ value: "init", label: "init", description: "Interactively create or update config" },
-];
+const FLUSH_ACTION = "Flush completed traces for this session";
+const SET_UP_ACTION = "Set up Langfuse for this Pi agent directory (restart required)";
+const UPDATE_ACTION = "Update Langfuse for this Pi agent directory (restart required)";
+const HELP_ACTION = "Show setup and privacy help";
 
 export function createLangfuseExtension(
 	dependencies: Partial<ExtensionDependencies> = {},
@@ -38,97 +43,97 @@ export function createLangfuseExtension(
 		});
 
 	return function langfuse(pi: ExtensionAPI) {
+		const resolveGit =
+			dependencies.resolveGitMetadata ??
+			((cwd: string) =>
+				resolveGitMetadata((command, args, options) => pi.exec(command, args, options), cwd));
 		let recorder: TraceRecorder | undefined;
 		let activeConfig: LangfuseConfig | undefined;
 		let configPath: string | undefined;
 		let initializationError: string | undefined;
+		let hasStoredConfig = false;
+		let configurationNotice: string | undefined;
+		let sessionGeneration = 0;
+		let nextAttemptReason: string | undefined;
 
 		pi.registerCommand("langfuse", {
-			description: "Manage Langfuse tracing and configuration",
-			getArgumentCompletions: (prefix) => {
-				const normalized = prefix.trimStart().toLowerCase();
-				if (/\s/.test(normalized)) return null;
-				const matches = COMMAND_COMPLETIONS.filter((item) => item.value.startsWith(normalized));
-				return matches.length > 0 ? matches : null;
-			},
-			handler: async (args, ctx) => {
-				const action = args.trim().toLowerCase() || "status";
-				if (action === "status") {
-					if (activeConfig) {
-						ctx.ui.notify(
-							[
-								"Langfuse tracing is enabled.",
-								`Endpoint: ${activeConfig.baseUrl}`,
-								`Content capture: ${activeConfig.captureContent ? "enabled" : "disabled"}`,
-								`Configuration: ${configPath ?? "unknown"}`,
-							].join("\n"),
-							"info",
-						);
-						return;
-					}
+			description: "Open interactive Langfuse tracing controls",
+			handler: async (_args, ctx) => {
+				if (!ctx.hasUI) {
 					ctx.ui.notify(
-						initializationError ??
-							"Langfuse tracing is disabled because its credentials are not configured.",
+						formatNonInteractiveStatus(activeConfig, configPath, initializationError),
 						"warning",
 					);
 					return;
 				}
-				if (action === "flush") {
-					if (!recorder) {
-						ctx.ui.notify("Langfuse tracing is not enabled.", "warning");
+
+				const menuGeneration = sessionGeneration;
+				const menuRecorder = recorder;
+				const menuConfig = activeConfig;
+				const choice = await ctx.ui.select(
+					formatMenuTitle(activeConfig, configPath, initializationError, configurationNotice),
+					[
+						...(menuRecorder ? [FLUSH_ACTION] : []),
+						hasStoredConfig ? UPDATE_ACTION : SET_UP_ACTION,
+						HELP_ACTION,
+					],
+				);
+				if (!choice || menuGeneration !== sessionGeneration) return;
+
+				if (choice === FLUSH_ACTION) {
+					if (!menuRecorder) {
+						ctx.ui.notify("Langfuse tracing is not enabled for this session.", "warning");
 						return;
 					}
 					try {
-						await recorder.flush();
-						ctx.ui.notify("Langfuse traces flushed.", "info");
+						await menuRecorder.flush();
+						if (menuGeneration !== sessionGeneration) return;
+						ctx.ui.notify("Langfuse traces flushed for this session.", "info");
 					} catch (error) {
-						ctx.ui.notify(`Langfuse flush failed: ${formatError(error)}`, "error");
+						if (menuGeneration !== sessionGeneration) return;
+						ctx.ui.notify(`Langfuse flush failed: ${formatError(error, menuConfig)}`, "error");
 					}
 					return;
 				}
-				if (action === "init") {
-					if (!ctx.hasUI) {
-						ctx.ui.notify(
-							"/langfuse init requires interactive UI. Edit pi-langfuse.json manually.",
-							"warning",
-						);
-						return;
-					}
+
+				if (choice === SET_UP_ACTION || choice === UPDATE_ACTION) {
 					const loaded = await loadConfig(configPath);
+					if (menuGeneration !== sessionGeneration) return;
 					configPath = loaded.path;
 					const next = await promptForConfig(ctx, loaded.ok ? loaded.config : undefined);
-					if (!next) return;
+					if (!next || menuGeneration !== sessionGeneration) return;
 					try {
 						await writeConfig(next, loaded.path);
-						initializationError =
-							"Langfuse config was saved. Restart Pi to apply the new configuration.";
-						ctx.ui.notify(`Saved Langfuse config to ${loaded.path}. Restart Pi to apply it.`, "info");
+						if (menuGeneration !== sessionGeneration) return;
+						hasStoredConfig = true;
+						configurationNotice =
+							"Saved; restart each Pi process to use it in subsequent sessions.";
+						ctx.ui.notify(
+							`Saved Langfuse config to ${loaded.path} for this Pi agent directory. Restart each Pi process to apply it to subsequent sessions.`,
+							"info",
+						);
 					} catch (error) {
-						ctx.ui.notify(`Failed to save Langfuse config: ${formatError(error)}`, "error");
+						if (menuGeneration !== sessionGeneration) return;
+						ctx.ui.notify(`Failed to save Langfuse config: ${formatError(error, next)}`, "error");
 					}
 					return;
 				}
-				if (action === "help") {
-					ctx.ui.notify(
-						[
-							"Usage: /langfuse [status|flush|help|init]",
-							"status: show tracing state without credentials",
-							"flush: wait for completed traces to export",
-							"init: interactively create or update the private config",
-						].join("\n"),
-						"info",
-					);
-					return;
+
+				if (choice === HELP_ACTION) {
+					ctx.ui.notify(formatHelp(configPath), "info");
 				}
-				ctx.ui.notify("Usage: /langfuse [status|flush|help|init]", "warning");
 			},
 		});
 
 		pi.on("session_start", async (_event, ctx) => {
+			sessionGeneration += 1;
 			recorder = undefined;
 			activeConfig = undefined;
 			configPath = undefined;
 			initializationError = undefined;
+			hasStoredConfig = false;
+			configurationNotice = undefined;
+			nextAttemptReason = undefined;
 
 			const result = await loadConfig();
 			configPath = result.path;
@@ -139,6 +144,7 @@ export function createLangfuseExtension(
 				return;
 			}
 
+			hasStoredConfig = true;
 			try {
 				const backend = await createBackend(result.config);
 				activeConfig = result.config;
@@ -149,43 +155,59 @@ export function createLangfuseExtension(
 					captureContent: result.config.captureContent,
 				});
 			} catch (error) {
-				initializationError = `Langfuse tracing could not start: ${formatError(error)}`;
+				initializationError = `Langfuse tracing could not start: ${formatError(error, result.config)}`;
 				ctx.ui.notify(initializationError, "warning");
 			}
 		});
 
-		pi.on("before_agent_start", (event, ctx) => {
-			recorder?.beginAgent({
+		pi.on("before_agent_start", async (event, ctx) => {
+			nextAttemptReason = undefined;
+			const activeRecorder = recorder;
+			if (!activeRecorder) return;
+			const git = await resolveGit(ctx.cwd).catch(() => undefined);
+			if (recorder !== activeRecorder) return;
+			activeRecorder.beginAgent({
 				prompt: event.prompt,
 				images: event.images,
-				model: ctx.model ? { provider: ctx.model.provider, id: ctx.model.id } : undefined,
+				model: ctx.model
+					? { provider: ctx.model.provider, id: ctx.model.id, api: ctx.model.api }
+					: undefined,
+				git,
+				snapshot: contextSnapshot(ctx),
 			});
 		});
 
+		pi.on("agent_start", () => {
+			if (!recorder) return;
+			recorder.beginAttempt(nextAttemptReason ? { reason: nextAttemptReason } : undefined);
+			nextAttemptReason = undefined;
+		});
+
 		pi.on("turn_start", (event, ctx) => {
 			if (!recorder) return;
-			if (!recorder.hasActiveTrace()) {
-				recorder.beginAgent({
-					prompt: "[automatic continuation]",
-					model: ctx.model ? { provider: ctx.model.provider, id: ctx.model.id } : undefined,
-				});
-			}
+			ensureActiveRun(recorder, ctx);
 			recorder.beginTurn(event.turnIndex);
 		});
 
-		pi.on("before_provider_request", (_event, ctx) => {
+		pi.on("before_provider_request", (event, ctx) => {
 			if (!recorder) return;
-			if (!recorder.hasActiveTrace()) {
-				recorder.beginAgent({
-					prompt: "[automatic continuation]",
-					model: ctx.model ? { provider: ctx.model.provider, id: ctx.model.id } : undefined,
-				});
-			}
-			recorder.beginGeneration();
+			ensureActiveRun(recorder, ctx);
+			recorder.beginGeneration({
+				payload: event.payload,
+				payloadStage: "before_provider_request",
+				model: ctx.model
+					? { provider: ctx.model.provider, id: ctx.model.id, api: ctx.model.api }
+					: undefined,
+				thinkingLevel: pi.getThinkingLevel(),
+			});
 		});
 
 		pi.on("after_provider_response", (event) => {
-			recorder?.recordProviderResponse(event.status);
+			recorder?.recordProviderResponse(event.status, event.headers);
+		});
+
+		pi.on("message_update", (event) => {
+			if (isRealOutputDelta(event.assistantMessageEvent)) recorder?.markGenerationFirstOutput();
 		});
 
 		pi.on("message_end", (event) => {
@@ -201,11 +223,15 @@ export function createLangfuseExtension(
 		});
 
 		pi.on("tool_execution_start", (event) => {
-			recorder?.beginTool(event.toolCallId, event.toolName);
+			recorder?.beginTool(event.toolCallId, event.toolName, event.args);
+		});
+
+		pi.on("tool_execution_update", (event) => {
+			recorder?.recordToolProgress(event.toolCallId);
 		});
 
 		pi.on("tool_result", (event) => {
-			recorder?.updateToolInput(event.toolCallId, event.input);
+			recorder?.recordToolInput(event.toolCallId, event.input);
 		});
 
 		pi.on("tool_execution_end", (event) => {
@@ -217,33 +243,177 @@ export function createLangfuseExtension(
 		});
 
 		pi.on("agent_end", (event) => {
-			for (let index = event.messages.length - 1; index >= 0; index -= 1) {
-				const message = event.messages[index];
-				if (message?.role !== "assistant") continue;
-				recorder?.finishAssistant(message);
-				break;
-			}
-			recorder?.settle();
+			const message = findLastAssistant(event.messages);
+			recorder?.finishAttempt(message);
+		});
+
+		pi.on("session_before_compact", (event) => {
+			recorder?.beginCompaction({
+				reason: event.reason,
+				willRetry: event.willRetry,
+				tokensBefore: event.preparation.tokensBefore,
+				messagesToSummarize: event.preparation.messagesToSummarize.length,
+				turnPrefixMessages: event.preparation.turnPrefixMessages.length,
+				branchEntries: event.branchEntries.length,
+				isSplitTurn: event.preparation.isSplitTurn,
+			});
+		});
+
+		pi.on("session_compact", (event) => {
+			const entry = event.compactionEntry as typeof event.compactionEntry & {
+				usage?: Parameters<TraceRecorder["finishCompaction"]>[0]["usage"];
+			};
+			recorder?.finishCompaction({
+				reason: event.reason,
+				willRetry: event.willRetry,
+				fromExtension: event.fromExtension,
+				tokensBefore: entry.tokensBefore,
+				details: entry.details,
+				usage: entry.usage,
+			});
+			if (event.willRetry && recorder?.hasActiveTrace()) nextAttemptReason = "post_compaction";
+		});
+
+		pi.on("agent_settled", (_event, ctx) => {
+			recorder?.settle(contextSnapshot(ctx));
+			nextAttemptReason = undefined;
 		});
 
 		pi.on("session_shutdown", async (event, ctx) => {
+			sessionGeneration += 1;
 			const activeRecorder = recorder;
+			const shutdownConfig = activeConfig;
 			recorder = undefined;
 			activeConfig = undefined;
 			if (!activeRecorder) return;
 			try {
-				if (event.reason === "quit") await activeRecorder.shutdown();
+				const snapshot = contextSnapshot(ctx);
+				if (event.reason === "quit") await activeRecorder.shutdown(snapshot);
 				else {
-					activeRecorder.settle();
+					activeRecorder.interrupt(
+						`Pi session ended before settlement (${event.reason}).`,
+						snapshot,
+					);
 					await activeRecorder.flush();
 				}
 			} catch (error) {
-				ctx.ui.notify(`Langfuse shutdown export failed: ${formatError(error)}`, "error");
+				ctx.ui.notify(
+					`Langfuse shutdown export failed: ${formatError(error, shutdownConfig)}`,
+					"error",
+				);
 			}
 		});
 	};
 }
 
+function ensureActiveRun(recorder: TraceRecorder, ctx: ExtensionContext): void {
+	if (!recorder.hasActiveTrace()) {
+		recorder.beginAgent({
+			prompt: "[automatic continuation]",
+			model: ctx.model
+				? { provider: ctx.model.provider, id: ctx.model.id, api: ctx.model.api }
+				: undefined,
+			snapshot: contextSnapshot(ctx),
+		});
+	}
+	if (!recorder.hasActiveAttempt()) recorder.beginAttempt();
+}
+
+function contextSnapshot(ctx: ExtensionContext): ContextSnapshot {
+	return {
+		leafId:
+			typeof ctx.sessionManager.getLeafId === "function"
+				? ctx.sessionManager.getLeafId()
+				: undefined,
+		contextUsage: ctx.getContextUsage(),
+	};
+}
+
+function findLastAssistant<T extends { role?: string }>(messages: readonly T[]): T | undefined {
+	for (let index = messages.length - 1; index >= 0; index -= 1) {
+		const message = messages[index];
+		if (message?.role === "assistant") return message;
+	}
+	return undefined;
+}
+
+function isRealOutputDelta(event: { type: string; delta?: unknown }): boolean {
+	return (
+		(event.type === "text_delta" ||
+			event.type === "thinking_delta" ||
+			event.type === "toolcall_delta") &&
+		typeof event.delta === "string" &&
+		event.delta.length > 0
+	);
+}
+
+const GIT_LOOKUP_TIMEOUT_MS = 1_000;
+const MAX_GIT_BRANCH_LENGTH = 256;
+
+type GitExecutor = ExtensionAPI["exec"];
+
+export async function resolveGitMetadata(
+	exec: GitExecutor,
+	cwd: string,
+): Promise<GitMetadata | undefined> {
+	const [branchResult, commit] = await Promise.all([
+		resolveGitBranch(exec, cwd),
+		resolveGitCommit(exec, cwd),
+	]);
+	if (!branchResult.resolved) return undefined;
+	if (branchResult.branch) {
+		return {
+			branch: branchResult.branch,
+			...(commit ? { commit } : {}),
+			detached: false,
+		};
+	}
+	return commit ? { commit, detached: true } : undefined;
+}
+
+async function resolveGitBranch(
+	exec: GitExecutor,
+	cwd: string,
+): Promise<{ resolved: boolean; branch?: string }> {
+	try {
+		const result = await exec("git", ["symbolic-ref", "--quiet", "--short", "HEAD"], {
+			cwd,
+			timeout: GIT_LOOKUP_TIMEOUT_MS,
+		});
+		if (result.killed) return { resolved: false };
+		if (result.code === 1) return { resolved: true };
+		if (result.code !== 0) return { resolved: false };
+		const branch = normalizeGitBranch(result.stdout);
+		return branch ? { resolved: true, branch } : { resolved: false };
+	} catch {
+		return { resolved: false };
+	}
+}
+
+async function resolveGitCommit(exec: GitExecutor, cwd: string): Promise<string | undefined> {
+	try {
+		const result = await exec("git", ["rev-parse", "--verify", "--short=12", "HEAD"], {
+			cwd,
+			timeout: GIT_LOOKUP_TIMEOUT_MS,
+		});
+		if (result.code !== 0 || result.killed) return undefined;
+		const commit = result.stdout.trim();
+		return /^[0-9a-f]{4,64}$/iu.test(commit) ? commit.toLowerCase() : undefined;
+	} catch {
+		return undefined;
+	}
+}
+
+function normalizeGitBranch(value: string): string | undefined {
+	const branch = value.trim();
+	if (!branch || branch.length > MAX_GIT_BRANCH_LENGTH) return undefined;
+	for (const character of branch) {
+		const code = character.codePointAt(0) ?? 0;
+		if (code <= 31 || code === 127) return undefined;
+	}
+	return branch;
+}
+
 async function promptForConfig(
 	ctx: ExtensionCommandContext,
 	current: LangfuseConfig | undefined,
@@ -281,15 +451,88 @@ async function promptForConfig(
 	return normalized.config;
 }
 
+function formatMenuTitle(
+	activeConfig: LangfuseConfig | undefined,
+	configPath: string | undefined,
+	initializationError: string | undefined,
+	configurationNotice: string | undefined,
+): string {
+	const lines = [
+		"Langfuse",
+		"",
+		"Current session:",
+		`  Tracing: ${activeConfig ? "enabled" : "disabled"}`,
+	];
+	if (activeConfig) {
+		lines.push(
+			`  Endpoint: ${activeConfig.baseUrl}`,
+			`  Content capture: ${activeConfig.captureContent ? "enabled" : "disabled"}`,
+		);
+	} else if (configurationNotice) {
+		lines.push("  State: tracing remains disabled until Pi restarts.");
+	} else if (initializationError) {
+		lines.push(`  Reason: ${initializationError}`);
+	}
+	lines.push(
+		"",
+		"Agent-directory configuration:",
+		`  Configuration: ${configPath ?? "unknown"}`,
+		"  Scope: this Pi agent directory; restart each Pi process to apply changes.",
+	);
+	if (configurationNotice) lines.push(`  Pending: ${configurationNotice}`);
+	lines.push("", "What do you want to do?");
+	return lines.join("\n");
+}
+
+function formatNonInteractiveStatus(
+	activeConfig: LangfuseConfig | undefined,
+	configPath: string | undefined,
+	initializationError: string | undefined,
+): string {
+	return [
+		"/langfuse requires interactive UI.",
+		`Current session tracing: ${activeConfig ? "enabled" : "disabled"}`,
+		...(activeConfig
+			? [
+					`Endpoint: ${activeConfig.baseUrl}`,
+					`Content capture: ${activeConfig.captureContent ? "enabled" : "disabled"}`,
+				]
+			: initializationError
+				? [`Reason: ${initializationError}`]
+				: []),
+		`Configuration: ${configPath ?? "unknown"}`,
+		"Edit the private config manually, then restart each Pi process to apply it to subsequent sessions.",
+	].join("\n");
+}
+
+function formatHelp(configPath: string | undefined): string {
+	return [
+		"Langfuse setup and privacy:",
+		"Run /langfuse in interactive Pi to manage tracing for the current session.",
+		`Configuration: ${configPath ?? "pi-langfuse.json"}`,
+		"The file belongs to this Pi agent directory; restart each Pi process after changing it.",
+		"Trace content may contain prompts, responses, tool arguments, tool results, and source code.",
+		"Git branch, commit, cwd, model, and usage remain in metadata when content capture is disabled.",
+		'Use "captureContent": false in the private config to export metadata only.',
+	].join("\n");
+}
+
 function formatConfigError(result: Extract<LangfuseConfigResult, { ok: false }>): string {
 	const setupHint = result.reason.startsWith("Configuration file not found:")
-		? " Run /langfuse init to create it."
+		? " Run /langfuse and choose Set up Langfuse."
 		: "";
 	return `Langfuse tracing is disabled: ${result.reason}${setupHint}`;
 }
 
-function formatError(error: unknown): string {
-	return error instanceof Error ? error.message : String(error);
+function formatError(error: unknown, config?: LangfuseConfig): string {
+	let message = error instanceof Error ? error.message : String(error);
+	const secrets = [config?.publicKey, config?.secretKey]
+		.filter((secret): secret is string => Boolean(secret))
+		.sort((left, right) => right.length - left.length);
+	for (const secret of secrets) {
+		message = message.split(secret).join("[LANGFUSE_KEY_REDACTED]");
+	}
+	return message;
 }
 
 export default createLangfuseExtension();