my-pi/pi-extension-codex-fast-mode/index.ts

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<string, unknown> {
  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<string, unknown> | 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<ExtensionCommandContext, "isIdle" | "hasPendingMessages">): boolean {
  return !ctx.isIdle() || ctx.hasPendingMessages();
}

function publishStatus(ctx: Pick<ExtensionContext, "ui">, 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<FastModeState>(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");
    },
  });
}