import { chmodSync, mkdirSync, readFileSync, renameSync, unlinkSync, writeFileSync } from "node:fs"; import { randomUUID } from "node:crypto"; import { dirname, join } from "node:path"; import { getAgentDir, type ExtensionAPI, type ExtensionCommandContext, type ExtensionContext, } from "@earendil-works/pi-coding-agent"; export const FAST_MODE_STATUS_KEY = "codex-fast-mode"; export const FAST_MODE_STATE_ENTRY_TYPE = "codex-fast-mode"; export const FAST_MODE_SERVICE_TIER = "priority"; export const FAST_MODE_CONFIG_DIRECTORY = "pi-extension-codex-fast-mode"; const OWNER_ONLY_DIRECTORY_MODE = 0o700; const OWNER_ONLY_FILE_MODE = 0o600; export type FastModeState = { enabled: boolean; }; export type FastModeModel = { provider?: unknown; api?: unknown; }; export type FastModeCommand = "toggle" | "on" | "off" | "status" | "invalid"; export type FastModeConfigLoadResult = { state: FastModeState; warning?: string; }; export type FastModeConfigSaveResult = | { success: true } | { success: false; error: string }; export type FastModeExtensionDependencies = { configPath?: string; loadGlobalState?: (path: string) => FastModeConfigLoadResult; saveGlobalState?: (path: string, state: FastModeState) => FastModeConfigSaveResult; }; /** Returns true only for object records that can safely receive a shallow request rewrite. */ export function isPlainObject(value: unknown): value is Record { if (value === null || typeof value !== "object") return false; const prototype = Object.getPrototypeOf(value); return prototype === Object.prototype || prototype === null; } /** Fast mode is limited to Pi's subscription-backed Codex Responses provider. */ export function isFastModeEligibleModel(model: FastModeModel | undefined): boolean { return model?.provider === "openai-codex" && model.api === "openai-codex-responses"; } /** * Applies Fast-mode request intent without mutating the provider's serialized payload. * Undefined is intentional: Pi then retains the original payload unchanged. */ export function transformFastModeRequest( enabled: boolean, model: FastModeModel | undefined, payload: unknown, ): Record | undefined { if (!enabled || !isFastModeEligibleModel(model) || !isPlainObject(payload)) return undefined; return { ...payload, service_tier: FAST_MODE_SERVICE_TIER }; } /** Resolves the owner-only global preference file beneath Pi's agent directory. */ export function getFastModeConfigPath(agentDir: string): string { return join(agentDir, "extensions", FAST_MODE_CONFIG_DIRECTORY, "config.json"); } /** Reads the cross-session default. Missing or invalid configuration fails safely to off. */ export function loadGlobalFastModeState(path: string): FastModeConfigLoadResult { try { const parsed = JSON.parse(readFileSync(path, "utf8")) as unknown; if (!isPlainObject(parsed) || typeof parsed.enabled !== "boolean") { return { state: { enabled: false }, warning: `Ignoring invalid Fast mode config at ${path}; expected { "enabled": boolean }.`, }; } return { state: { enabled: parsed.enabled } }; } catch (error) { if (typeof error === "object" && error !== null && "code" in error && error.code === "ENOENT") { return { state: { enabled: false } }; } return { state: { enabled: false }, warning: `Could not read Fast mode config at ${path}: ${error instanceof Error ? error.message : String(error)}`, }; } } /** Atomically writes the cross-session default with owner-only permissions. */ export function saveGlobalFastModeState(path: string, state: FastModeState): FastModeConfigSaveResult { const directory = dirname(path); const temporaryPath = `${path}.${process.pid}.${randomUUID()}.tmp`; try { mkdirSync(directory, { recursive: true, mode: OWNER_ONLY_DIRECTORY_MODE }); chmodSync(directory, OWNER_ONLY_DIRECTORY_MODE); writeFileSync(temporaryPath, `${JSON.stringify(state, null, 2)}\n`, { encoding: "utf8", flag: "wx", mode: OWNER_ONLY_FILE_MODE, }); chmodSync(temporaryPath, OWNER_ONLY_FILE_MODE); renameSync(temporaryPath, path); chmodSync(path, OWNER_ONLY_FILE_MODE); return { success: true }; } catch (error) { try { unlinkSync(temporaryPath); } catch { // The temporary file may not have been created. } return { success: false, error: error instanceof Error ? error.message : String(error), }; } } /** Reconstructs the latest branch snapshot, falling back to the cross-session default. */ export function reconstructFastModeState( entries: readonly unknown[], defaultEnabled = false, ): FastModeState { let enabled = defaultEnabled; for (const entry of entries) { if (!isPlainObject(entry)) continue; if (entry.type !== "custom" || entry.customType !== FAST_MODE_STATE_ENTRY_TYPE) continue; if (!isPlainObject(entry.data) || typeof entry.data.enabled !== "boolean") continue; enabled = entry.data.enabled; } return { enabled }; } export function parseFastModeCommand(args: string): FastModeCommand { const normalized = args.trim().toLowerCase(); if (!normalized) return "toggle"; if (normalized === "on" || normalized === "off" || normalized === "status") return normalized; return "invalid"; } export function fastModeArgumentCompletions(prefix: string) { const normalized = prefix.trim().toLowerCase(); return ["on", "off", "status"] .filter((value) => value.startsWith(normalized)) .map((value) => ({ value, label: value })); } function isBusy(ctx: Pick): boolean { return !ctx.isIdle() || ctx.hasPendingMessages(); } function publishStatus(ctx: Pick, enabled: boolean): void { ctx.ui.setStatus(FAST_MODE_STATUS_KEY, enabled ? "on" : "off"); } function formatStatus(enabled: boolean, globalEnabled: boolean): string { const state = enabled ? "on" : "off"; const globalState = globalEnabled ? "on" : "off"; return `Fast mode: ${state}. New sessions default to ${globalState}. It only requests priority service for openai-codex/openai-codex-responses.`; } export default function codexFastModeExtension( pi: ExtensionAPI, dependencies: FastModeExtensionDependencies = {}, ): void { const configPath = dependencies.configPath ?? getFastModeConfigPath(getAgentDir()); const loadGlobalState = dependencies.loadGlobalState ?? loadGlobalFastModeState; const saveGlobalState = dependencies.saveGlobalState ?? saveGlobalFastModeState; let enabled = false; let globalEnabled = false; let lastConfigWarning: string | undefined; const restoreState = (ctx: ExtensionContext): void => { const loaded = loadGlobalState(configPath); globalEnabled = loaded.state.enabled; enabled = reconstructFastModeState(ctx.sessionManager.getBranch(), globalEnabled).enabled; publishStatus(ctx, enabled); if (loaded.warning && loaded.warning !== lastConfigWarning) { ctx.ui.notify(loaded.warning, "warning"); } lastConfigWarning = loaded.warning; }; const persistSessionState = (): void => { pi.appendEntry(FAST_MODE_STATE_ENTRY_TYPE, { enabled }); }; const setEnabled = (ctx: ExtensionCommandContext, nextEnabled: boolean): void => { const stateChanged = enabled !== nextEnabled; const saveResult = saveGlobalState(configPath, { enabled: nextEnabled }); if (saveResult.success) { globalEnabled = nextEnabled; } if (stateChanged) { enabled = nextEnabled; persistSessionState(); } publishStatus(ctx, enabled); if (stateChanged) { ctx.ui.notify(enabled ? "Fast mode enabled. Supported Codex requests will use priority service, and new sessions will inherit this setting." : "Fast mode disabled. Supported Codex requests will keep their existing service tier, and new sessions will inherit this setting.", "info"); } else { ctx.ui.notify(`Fast mode is already ${enabled ? "on" : "off"}. The global default was refreshed.`, "info"); } if (!saveResult.success) { ctx.ui.notify(`The current session was updated, but the Fast mode global default could not be saved: ${saveResult.error}`, "warning"); } }; pi.on("session_start", (_event, ctx) => { restoreState(ctx); }); pi.on("session_tree", (_event, ctx) => { restoreState(ctx); }); pi.on("before_provider_request", (event, ctx) => { return transformFastModeRequest(enabled, ctx.model, event.payload); }); pi.registerCommand("fast-mode", { description: "Toggle persistent Codex subscription Fast mode. Usage: /fast-mode [on|off|status]", getArgumentCompletions: fastModeArgumentCompletions, handler: async (args, ctx) => { const command = parseFastModeCommand(args); if (command === "status") { publishStatus(ctx, enabled); ctx.ui.notify(formatStatus(enabled, globalEnabled), "info"); return; } if (command === "invalid") { ctx.ui.notify("Usage: /fast-mode [on|off|status]", "warning"); return; } if (isBusy(ctx)) { ctx.ui.notify("Fast mode cannot be changed while the session is busy. Run /fast-mode status to inspect it.", "warning"); return; } setEnabled(ctx, command === "toggle" ? !enabled : command === "on"); }, }); }