pi-ext

import {
	lazyStream,
	type Provider,
	type SimpleStreamOptions,
} from "@earendil-works/pi-ai";
import type {
	ExtensionAPI,
	ExtensionContext,
	ModelRegistry,
} from "@earendil-works/pi-coding-agent";
import type { QueryCoordinator } from "./src/coordinator.js";
import { closeDebug, dbg, debugSpan } from "./src/debug.js";
import { KLAUS_BASE_URL, KLAUS_PROVIDER, klausModels } from "./src/models.js";

dbg?.("index.module");

export default function klaus(pi: ExtensionAPI): void | Promise<void> {
	dbg?.("index.extension.start");
	let context: ExtensionContext | undefined;
	/** Last registry Pi handed to Klaus; `getModels()` projects from it on demand. */
	let registry: ModelRegistry | undefined;
	let warnedFable = false;
	let warnedPromptDrift = false;
	let coordinatorPromise: Promise<QueryCoordinator> | undefined;
	const getCoordinator = () =>
		(coordinatorPromise ??= import("./src/coordinator.js").then(
			({ QueryCoordinator }) =>
				new QueryCoordinator(
					async (modelId) => {
						const finish = debugSpan?.("index.oauth");
						try {
							if (!context)
								throw new Error("Klaus has no active Pi session context.");
							const source = context.modelRegistry.find("anthropic", modelId);
							dbg?.("index.oauth.model", { found: Boolean(source) });
							if (!source)
								throw new Error(
									`Anthropic model ${modelId} is unavailable in Pi.`,
								);
							const usingOAuth = context.modelRegistry.isUsingOAuth(source);
							dbg?.("index.oauth.kind", { usingOAuth });
							if (!usingOAuth) {
								throw new Error(
									"Klaus requires Anthropic OAuth. Use Pi /login for Anthropic first.",
								);
							}
							const auth =
								await context.modelRegistry.getProviderAuth("anthropic");
							const token = auth?.auth.apiKey;
							dbg?.("index.oauth.resolved", { hasAuth: Boolean(auth) });
							if (!token)
								throw new Error(
									"Pi could not resolve Anthropic OAuth for Klaus.",
								);
							if (modelId === "claude-fable-5-1" && !warnedFable) {
								dbg?.("index.oauth.fableWarning");
								warnedFable = true;
								context.ui.notify(
									"Klaus Fable may consume Anthropic usage credits when subscription discovery omits it.",
									"warning",
								);
							}
							finish?.();
							return token;
						} catch (error) {
							finish?.("error", { kind: "oauth" });
							throw error;
						}
					},
					() => {
						const cwd = context?.cwd ?? process.cwd();
						dbg?.("index.currentCwd");
						return cwd;
					},
					() => {
						const scope = context
							? {
									id: context.sessionManager.getSessionId(),
									persisted:
										context.sessionManager.getSessionFile() !== undefined,
									dir: context.sessionManager.getSessionDir(),
								}
							: undefined;
						dbg?.("index.currentSession", {
							hasScope: Boolean(scope),
							...(scope ? { persisted: scope.persisted } : {}),
						});
						return scope;
					},
					(message) => {
						dbg?.("index.promptDriftWarning", {
							alreadyWarned: warnedPromptDrift,
						});
						if (warnedPromptDrift || !context) return;
						warnedPromptDrift = true;
						context.ui.notify(message, "warning");
					},
				),
		));

	pi.registerCommand("klaus-setup-token", {
		description: "Mint a long-lived Claude OAuth token with your subscription",
		async handler(_args, commandCtx) {
			const { runSetupTokenCommand } = await import("./src/setup-token.js");
			await runSetupTokenCommand(commandCtx);
		},
	});

	let preparedSession: string | undefined;
	let preparation: Promise<void> | undefined;
	const readyCoordinator = async () => {
		const coordinator = await getCoordinator();
		const active = context;
		const session = active
			? `${active.sessionManager.getSessionId()}\0${active.sessionManager.getSessionDir()}`
			: "none";
		if (preparedSession === session) return coordinator;
		preparation ??= coordinator
			.preparePrimaryCache(
				active
					? {
							id: active.sessionManager.getSessionId(),
							persisted: active.sessionManager.getSessionFile() !== undefined,
							dir: active.sessionManager.getSessionDir(),
						}
					: { id: "", persisted: false, dir: process.cwd() },
			)
			.then(() => {
				preparedSession = session;
			})
			.finally(() => {
				preparation = undefined;
			});
		await preparation;
		return coordinator;
	};

	const streamSimple: Provider["streamSimple"] = (
		model,
		requestContext,
		options,
	) => {
		dbg?.("index.streamSimple", {
			messageCount: requestContext.messages.length,
		});
		return lazyStream(model, async () =>
			(await readyCoordinator()).stream(model, requestContext, options),
		);
	};
	// Lazy Provider form: Pi calls `getModels()` whenever it recomposes its
	// catalog (registration, models.json reload, refresh), so projection follows
	// the Anthropic catalog without Klaus polling for changes.
	const provider: Provider = {
		id: KLAUS_PROVIDER,
		name: "Klaus",
		baseUrl: KLAUS_BASE_URL,
		auth: {
			apiKey: {
				name: "Klaus",
				// Anthropic OAuth is resolved per request; this placeholder only
				// marks the provider as configured and is never sent anywhere.
				resolve: async () => ({
					auth: { apiKey: KLAUS_PROVIDER },
					source: "Pi Anthropic OAuth",
				}),
			},
		},
		getModels: () => klausModels(registry),
		// Pi's agent loop uses streamSimple; raw stream options map onto the same route.
		stream: (model, requestContext, options) =>
			streamSimple(
				model,
				requestContext,
				options as SimpleStreamOptions | undefined,
			),
		streamSimple,
	};

	dbg?.("index.register");
	pi.registerProvider(provider);

	pi.on("session_start", (_event, nextContext) => {
		dbg?.("index.session_start", {
			persisted: nextContext.sessionManager.getSessionFile() !== undefined,
		});
		context = nextContext;
		preparedSession = undefined;
		warnedFable = false;
		warnedPromptDrift = false;
		if (registry !== nextContext.modelRegistry) {
			// First sight of Pi's registry: re-register once so Pi re-evaluates
			// `getModels()` against the session catalog instead of the builtin one.
			registry = nextContext.modelRegistry;
			dbg?.("index.register");
			pi.registerProvider(provider);
		}
	});
	pi.on("turn_end", async (_event, activeContext) => {
		const leafId = activeContext.sessionManager.getLeafId();
		dbg?.("index.turn_end");
		const coordinator = await coordinatorPromise?.catch(() => undefined);
		if (coordinator && leafId)
			await coordinator.persistCheckpoint(
				activeContext.sessionManager.getSessionId(),
				leafId,
			);
	});
	const flushAndClose = async (
		activeContext: ExtensionContext,
		reason: string,
	): Promise<void> => {
		const coordinator = await coordinatorPromise?.catch(() => undefined);
		if (!coordinator) return;
		const leafId = activeContext.sessionManager.getLeafId();
		const finish = debugSpan?.("index.flush", { hasLeaf: Boolean(leafId) });
		try {
			if (leafId)
				await coordinator.persistCheckpoint(
					activeContext.sessionManager.getSessionId(),
					leafId,
				);
			await coordinator.closeAll(reason);
			finish?.();
		} catch (error) {
			finish?.("error", { kind: "flush" });
			throw error;
		}
	};
	pi.on("session_before_switch", async (_event, activeContext) => {
		dbg?.("index.session_before_switch");
		await flushAndClose(activeContext, "Pi switched sessions.");
	});
	pi.on("session_before_fork", async (_event, activeContext) => {
		dbg?.("index.session_before_fork");
		await flushAndClose(activeContext, "Pi forked the session.");
	});
	pi.on("session_before_compact", async (_event, activeContext) => {
		dbg?.("index.session_before_compact");
		await flushAndClose(activeContext, "Pi compacted the session.");
	});
	pi.on("session_shutdown", async (_event, activeContext) => {
		dbg?.("index.session_shutdown");
		try {
			await flushAndClose(activeContext, "Pi shut down the session.");
		} finally {
			context = undefined;
			preparedSession = undefined;
			closeDebug();
		}
	});
}