pi-ext

import {
	type Api,
	clampThinkingLevel,
	type Model,
	type SimpleStreamOptions,
	type TranscriptContext,
} from "@earendil-works/pi-ai";
import {
	streamOpenAICodexResponses,
	streamSimpleOpenAICodexResponses,
} from "@earendil-works/pi-ai/compat";
import type {
	ExtensionAPI,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { closeDebug, dbg, span } from "./src/debug.ts";
import {
	parseBooleanSetting,
	resolveSetting,
	type SettingDeclaration,
} from "./src/pi-ext-settings.ts";
import { isStaleContextError } from "./src/pi-ext-stale-context.ts";

const PROVIDER = "openai-codex";
const API = "openai-codex-responses";
const STATUS_ID = "fast";
const FAST_ICON = "󱐋";

// Precedence: trusted project, user, false.
const ENABLE_SETTING: SettingDeclaration<boolean> = {
	key: "fast.enable",
	parse: parseBooleanSetting,
	default: false,
};

function fastActive(
	ctx: Pick<ExtensionContext, "model">,
	enabled: boolean,
): boolean {
	return enabled && ctx.model?.provider === PROVIDER;
}

export default function fast(pi: ExtensionAPI): void {
	// Until a session reports trust, project settings stay unread.
	let settingsCtx = { cwd: process.cwd(), isProjectTrusted: () => false };
	let enabled: boolean | undefined;
	let settingsError: string | undefined;
	let generation = 0;

	function ensureEnabled(): boolean {
		if (enabled === undefined) {
			const setting = resolveSetting(pi, settingsCtx, ENABLE_SETTING);
			enabled = setting.ok && setting.value;
			settingsError = setting.ok ? undefined : setting.error;
		}
		return enabled;
	}

	function clear(ctx?: ExtensionContext): void {
		pi.events.emit("model-route:fast", { active: false });
		if (ctx?.hasUI) ctx.ui.setStatus(STATUS_ID, undefined);
	}

	function publish(ctx: ExtensionContext): void {
		if (!ctx.hasUI) return;
		if (!fastActive(ctx, enabled ?? false)) {
			clear(ctx);
			return;
		}

		ctx.ui.setStatus(STATUS_ID, FAST_ICON);
		pi.events.emit("model-route:fast", { active: true });
	}

	pi.registerProvider(PROVIDER, {
		api: API,
		streamSimple(
			model: Model<Api>,
			context: TranscriptContext,
			options?: SimpleStreamOptions,
		) {
			if (!ensureEnabled())
				return streamSimpleOpenAICodexResponses(
					model as Model<typeof API>,
					context,
					options,
				);

			const clampedReasoning = options?.reasoning
				? clampThinkingLevel(model, options.reasoning)
				: undefined;
			const reasoningEffort =
				clampedReasoning === "off" ? undefined : clampedReasoning;

			return streamOpenAICodexResponses(model as Model<typeof API>, context, {
				...options,
				reasoningEffort,
				serviceTier: "priority",
			});
		},
	});

	function scheduleRefresh(ctx: ExtensionContext): void {
		const ticket = ++generation;
		const end = span?.("settings.refresh", { generation: ticket });
		const trusted = ctx.isProjectTrusted();
		settingsCtx = { cwd: ctx.cwd, isProjectTrusted: () => trusted };
		enabled = undefined;
		const timer = setTimeout(() => {
			if (ticket !== generation) {
				end?.("finish", { status: "stale" });
				return;
			}
			try {
				const isEnabled = ensureEnabled();
				if (settingsError && ctx.hasUI)
					ctx.ui.notify(`${settingsError}; fast mode off`, "warning");
				publish(ctx);
				end?.("finish", {
					status: settingsError
						? "invalid"
						: !isEnabled
							? "disabled"
							: ctx.model?.provider === PROVIDER
								? "active"
								: "provider-mismatch",
				});
			} catch (error) {
				end?.("error", { type: "publish" });
				if (!isStaleContextError(error)) throw error;
			}
		}, 0);
		timer.unref?.();
	}

	pi.on("session_start", (_event, ctx) => {
		dbg?.("session.start", { mode: ctx.mode });
		scheduleRefresh(ctx);
	});
	pi.on("model_select", (_event, ctx) => scheduleRefresh(ctx));
	pi.on("session_shutdown", (_event, ctx) => {
		dbg?.("session.shutdown", { mode: ctx.mode });
		generation++;
		enabled = undefined;
		clear(ctx);
		closeDebug();
	});
}

export const _test = { fastActive };