@narumitw/pi-langfuse 0.14.1

package/src/langfuse.ts

@@ -0,0 +1,295 @@
+import type {
+	ExtensionAPI,
+	ExtensionCommandContext,
+} from "@earendil-works/pi-coding-agent";
+import {
+	DEFAULT_BASE_URL,
+	type LangfuseConfig,
+	type LangfuseConfigResult,
+	loadLangfuseConfig,
+	normalizeLangfuseConfig,
+	writeLangfuseConfig,
+} from "./config.js";
+import { 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>;
+}
+
+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" },
+];
+
+export function createLangfuseExtension(
+	dependencies: Partial<ExtensionDependencies> = {},
+): (pi: ExtensionAPI) => void {
+	const loadConfig = dependencies.loadConfig ?? loadLangfuseConfig;
+	const writeConfig = dependencies.writeConfig ?? writeLangfuseConfig;
+	const createBackend =
+		dependencies.createBackend ??
+		(async (config) => {
+			const { createProductionBackend } = await import("./runtime.js");
+			return createProductionBackend(config);
+		});
+
+	return function langfuse(pi: ExtensionAPI) {
+		let recorder: TraceRecorder | undefined;
+		let activeConfig: LangfuseConfig | undefined;
+		let configPath: string | undefined;
+		let initializationError: 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;
+					}
+					ctx.ui.notify(
+						initializationError ??
+							"Langfuse tracing is disabled because its credentials are not configured.",
+						"warning",
+					);
+					return;
+				}
+				if (action === "flush") {
+					if (!recorder) {
+						ctx.ui.notify("Langfuse tracing is not enabled.", "warning");
+						return;
+					}
+					try {
+						await recorder.flush();
+						ctx.ui.notify("Langfuse traces flushed.", "info");
+					} catch (error) {
+						ctx.ui.notify(`Langfuse flush failed: ${formatError(error)}`, "error");
+					}
+					return;
+				}
+				if (action === "init") {
+					if (!ctx.hasUI) {
+						ctx.ui.notify(
+							"/langfuse init requires interactive UI. Edit pi-langfuse.json manually.",
+							"warning",
+						);
+						return;
+					}
+					const loaded = await loadConfig(configPath);
+					configPath = loaded.path;
+					const next = await promptForConfig(ctx, loaded.ok ? loaded.config : undefined);
+					if (!next) 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");
+					} catch (error) {
+						ctx.ui.notify(`Failed to save Langfuse config: ${formatError(error)}`, "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;
+				}
+				ctx.ui.notify("Usage: /langfuse [status|flush|help|init]", "warning");
+			},
+		});
+
+		pi.on("session_start", async (_event, ctx) => {
+			recorder = undefined;
+			activeConfig = undefined;
+			configPath = undefined;
+			initializationError = undefined;
+
+			const result = await loadConfig();
+			configPath = result.path;
+			for (const warning of result.warnings) ctx.ui.notify(warning, "warning");
+			if (!result.ok) {
+				initializationError = formatConfigError(result);
+				ctx.ui.notify(initializationError, "warning");
+				return;
+			}
+
+			try {
+				const backend = await createBackend(result.config);
+				activeConfig = result.config;
+				recorder = new TraceRecorder(backend, {
+					sessionId: ctx.sessionManager.getSessionId(),
+					cwd: ctx.cwd,
+					mode: ctx.mode,
+					captureContent: result.config.captureContent,
+				});
+			} catch (error) {
+				initializationError = `Langfuse tracing could not start: ${formatError(error)}`;
+				ctx.ui.notify(initializationError, "warning");
+			}
+		});
+
+		pi.on("before_agent_start", (event, ctx) => {
+			recorder?.beginAgent({
+				prompt: event.prompt,
+				images: event.images,
+				model: ctx.model ? { provider: ctx.model.provider, id: ctx.model.id } : 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,
+				});
+			}
+			recorder.beginTurn(event.turnIndex);
+		});
+
+		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();
+		});
+
+		pi.on("after_provider_response", (event) => {
+			recorder?.recordProviderResponse(event.status);
+		});
+
+		pi.on("message_end", (event) => {
+			if (event.message.role === "assistant") recorder?.markGenerationEnd();
+		});
+
+		pi.on("turn_end", (event) => {
+			if (event.message.role === "assistant") recorder?.finishAssistant(event.message);
+			recorder?.finishTurn(event.turnIndex, {
+				message: event.message,
+				toolResultCount: event.toolResults.length,
+			});
+		});
+
+		pi.on("tool_execution_start", (event) => {
+			recorder?.beginTool(event.toolCallId, event.toolName);
+		});
+
+		pi.on("tool_result", (event) => {
+			recorder?.updateToolInput(event.toolCallId, event.input);
+		});
+
+		pi.on("tool_execution_end", (event) => {
+			recorder?.finishTool(event.toolCallId, {
+				content: event.result.content,
+				details: event.result.details,
+				isError: event.isError,
+			});
+		});
+
+		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();
+		});
+
+		pi.on("session_shutdown", async (event, ctx) => {
+			const activeRecorder = recorder;
+			recorder = undefined;
+			activeConfig = undefined;
+			if (!activeRecorder) return;
+			try {
+				if (event.reason === "quit") await activeRecorder.shutdown();
+				else {
+					activeRecorder.settle();
+					await activeRecorder.flush();
+				}
+			} catch (error) {
+				ctx.ui.notify(`Langfuse shutdown export failed: ${formatError(error)}`, "error");
+			}
+		});
+	};
+}
+
+async function promptForConfig(
+	ctx: ExtensionCommandContext,
+	current: LangfuseConfig | undefined,
+): Promise<LangfuseConfig | undefined> {
+	const secretKey = await ctx.ui.input("Langfuse secret key (leave blank to keep existing):");
+	if (secretKey === undefined) {
+		ctx.ui.notify("Cancelled", "info");
+		return undefined;
+	}
+	const publicKey = await ctx.ui.input("Langfuse public key (leave blank to keep existing):");
+	if (publicKey === undefined) {
+		ctx.ui.notify("Cancelled", "info");
+		return undefined;
+	}
+	const baseUrl = await ctx.ui.input(
+		`Langfuse base URL (leave blank for default ${DEFAULT_BASE_URL}):`,
+		DEFAULT_BASE_URL,
+	);
+	if (baseUrl === undefined) {
+		ctx.ui.notify("Cancelled", "info");
+		return undefined;
+	}
+
+	const normalized = normalizeLangfuseConfig({
+		...current,
+		publicKey: publicKey.trim() || current?.publicKey || "",
+		secretKey: secretKey.trim() || current?.secretKey || "",
+		baseUrl: baseUrl.trim() || DEFAULT_BASE_URL,
+		captureContent: current?.captureContent ?? true,
+	});
+	if (!normalized.ok) {
+		ctx.ui.notify(`Invalid Langfuse config: ${normalized.reason}`, "error");
+		return undefined;
+	}
+	return normalized.config;
+}
+
+function formatConfigError(result: Extract<LangfuseConfigResult, { ok: false }>): string {
+	const setupHint = result.reason.startsWith("Configuration file not found:")
+		? " Run /langfuse init to create it."
+		: "";
+	return `Langfuse tracing is disabled: ${result.reason}${setupHint}`;
+}
+
+function formatError(error: unknown): string {
+	return error instanceof Error ? error.message : String(error);
+}
+
+export default createLangfuseExtension();

package/src/runtime.ts

@@ -0,0 +1,247 @@
+import { createHash } from "node:crypto";
+import { LangfuseSpanProcessor } from "@langfuse/otel";
+import {
+	type LangfuseObservation,
+	type LangfuseObservationAttributes,
+	type LangfuseTraceAttributes,
+	setLangfuseTracerProvider,
+	startObservation,
+} from "@langfuse/tracing";
+import type { SpanProcessor } from "@opentelemetry/sdk-trace-base";
+import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node";
+import type { LangfuseConfig } from "./config.js";
+import type {
+	Observation,
+	ObservationAttributes,
+	ObservationType,
+	TraceBackend,
+} from "./tracing.js";
+
+const RUNTIME_KEY = Symbol.for("@narumitw/pi-langfuse/runtime/v1");
+
+interface SharedRuntime {
+	fingerprint: string;
+	backend: ProductionTraceBackend;
+	shutdown: boolean;
+}
+
+type GlobalWithRuntime = typeof globalThis & {
+	[RUNTIME_KEY]?: Promise<SharedRuntime>;
+};
+
+export interface RuntimeFactories {
+	createProcessor(config: LangfuseConfig): SpanProcessor;
+	createProvider(processor: SpanProcessor): NodeTracerProvider;
+	selectProvider(provider: NodeTracerProvider): void;
+}
+
+class ProductionObservation implements Observation {
+	constructor(readonly native: LangfuseObservation) {}
+
+	update(attributes: ObservationAttributes): Observation {
+		this.native.updateOtelSpanAttributes(attributes as LangfuseObservationAttributes);
+		return this;
+	}
+
+	updateTrace(attributes: ObservationAttributes): Observation {
+		this.native.updateTrace(attributes as LangfuseTraceAttributes);
+		return this;
+	}
+
+	end(endTime?: number): Observation {
+		this.native.end(endTime);
+		return this;
+	}
+}
+
+class ProductionTraceBackend implements TraceBackend {
+	constructor(
+		private readonly provider: NodeTracerProvider,
+		private readonly processor: SpanProcessor,
+	) {}
+
+	start(
+		name: string,
+		attributes: ObservationAttributes,
+		options: { asType: ObservationType; parent?: Observation },
+	): Observation {
+		const parent = options.parent;
+		if (parent instanceof ProductionObservation) {
+			return new ProductionObservation(
+				startChild(parent.native, name, attributes, options.asType),
+			);
+		}
+		return new ProductionObservation(startRoot(name, attributes, options.asType));
+	}
+
+	async forceFlush(): Promise<void> {
+		await this.processor.forceFlush();
+	}
+
+	async shutdown(): Promise<void> {
+		await this.provider.shutdown();
+	}
+}
+
+export async function createProductionBackend(
+	config: LangfuseConfig,
+	factoryOverrides: Partial<RuntimeFactories> = {},
+): Promise<TraceBackend> {
+	const globalRuntime = globalThis as GlobalWithRuntime;
+	const fingerprint = configFingerprint(config);
+	const existing = globalRuntime[RUNTIME_KEY];
+	if (existing) {
+		const runtime = await existing;
+		if (runtime.shutdown) {
+			throw new Error("Langfuse tracing was already shut down; restart Pi to enable it again.");
+		}
+		if (runtime.fingerprint !== fingerprint) {
+			throw new Error("Langfuse configuration changed; restart Pi to apply the new credentials.");
+		}
+		return runtime.backend;
+	}
+
+	const factories = { ...defaultFactories, ...factoryOverrides };
+	const initializing = initializeRuntime(config, fingerprint, factories);
+	globalRuntime[RUNTIME_KEY] = initializing;
+	try {
+		return (await initializing).backend;
+	} catch (error) {
+		delete globalRuntime[RUNTIME_KEY];
+		throw error;
+	}
+}
+
+async function initializeRuntime(
+	config: LangfuseConfig,
+	fingerprint: string,
+	factories: RuntimeFactories,
+): Promise<SharedRuntime> {
+	const processor = factories.createProcessor(config);
+	let provider: NodeTracerProvider | undefined;
+	try {
+		provider = factories.createProvider(processor);
+		factories.selectProvider(provider);
+	} catch (error) {
+		await (provider?.shutdown() ?? processor.shutdown()).catch(() => undefined);
+		throw error;
+	}
+	const backend = new ProductionTraceBackend(provider, processor);
+	const runtime: SharedRuntime = { fingerprint, backend, shutdown: false };
+	const originalShutdown = backend.shutdown.bind(backend);
+	backend.shutdown = async () => {
+		if (runtime.shutdown) return;
+		runtime.shutdown = true;
+		await originalShutdown();
+	};
+	return runtime;
+}
+
+const defaultFactories: RuntimeFactories = {
+	createProcessor: (config) => {
+		const secrets = [config.secretKey, config.publicKey];
+		return new LangfuseSpanProcessor({
+			publicKey: config.publicKey,
+			secretKey: config.secretKey,
+			baseUrl: config.baseUrl,
+			environment: config.environment ?? "",
+			release: config.release ?? "",
+			flushAt: 512,
+			flushInterval: 5,
+			timeout: 5,
+			mask: ({ data }) => maskSecrets(data, secrets),
+			shouldExportSpan: ({ otelSpan }) =>
+				typeof otelSpan.attributes["langfuse.observation.type"] === "string",
+		});
+	},
+	createProvider: (processor) => new NodeTracerProvider({ spanProcessors: [processor] }),
+	selectProvider: setLangfuseTracerProvider,
+};
+
+function startRoot(
+	name: string,
+	attributes: ObservationAttributes,
+	type: ObservationType,
+): LangfuseObservation {
+	if (type === "agent") {
+		return startObservation(name, attributes as LangfuseObservationAttributes, { asType: "agent" });
+	}
+	if (type === "generation") {
+		return startObservation(name, attributes as LangfuseObservationAttributes, {
+			asType: "generation",
+		});
+	}
+	if (type === "tool") {
+		return startObservation(name, attributes as LangfuseObservationAttributes, { asType: "tool" });
+	}
+	return startObservation(name, attributes as LangfuseObservationAttributes, { asType: "span" });
+}
+
+function startChild(
+	parent: LangfuseObservation,
+	name: string,
+	attributes: ObservationAttributes,
+	type: ObservationType,
+): LangfuseObservation {
+	if (type === "agent") {
+		return parent.startObservation(name, attributes as LangfuseObservationAttributes, {
+			asType: "agent",
+		});
+	}
+	if (type === "generation") {
+		return parent.startObservation(name, attributes as LangfuseObservationAttributes, {
+			asType: "generation",
+		});
+	}
+	if (type === "tool") {
+		return parent.startObservation(name, attributes as LangfuseObservationAttributes, {
+			asType: "tool",
+		});
+	}
+	return parent.startObservation(name, attributes as LangfuseObservationAttributes, {
+		asType: "span",
+	});
+}
+
+function configFingerprint(config: LangfuseConfig): string {
+	return createHash("sha256")
+		.update(
+			`${config.publicKey}\0${config.secretKey}\0${config.baseUrl}\0${config.environment ?? ""}\0${config.release ?? ""}`,
+		)
+		.digest("hex");
+}
+
+export function maskSecrets(data: unknown, secrets: readonly string[]): unknown {
+	return maskSecretValue(data, secrets, new WeakSet<object>());
+}
+
+function maskSecretValue(
+	data: unknown,
+	secrets: readonly string[],
+	active: WeakSet<object>,
+): unknown {
+	if (typeof data === "string") {
+		let masked = data.replace(/\b(?:sk|pk)-lf-[A-Za-z0-9_-]+\b/g, "[LANGFUSE_KEY_REDACTED]");
+		for (const secret of secrets) {
+			if (secret) masked = masked.replaceAll(secret, "[LANGFUSE_KEY_REDACTED]");
+		}
+		return masked;
+	}
+	if (!data || typeof data !== "object") return data;
+	if (active.has(data)) return "[circular]";
+
+	active.add(data);
+	try {
+		if (Array.isArray(data)) {
+			return data.map((item) => maskSecretValue(item, secrets, active));
+		}
+		return Object.fromEntries(
+			Object.entries(data).map(([key, value]) => [
+				key,
+				maskSecretValue(value, secrets, active),
+			]),
+		);
+	} finally {
+		active.delete(data);
+	}
+}