mirror of
https://bitbucket.org/siakitem/my-pi.git
synced 2026-08-28 08:35:57 +00:00
feat: vendor and persist Codex fast mode
This commit is contained in:
@@ -0,0 +1,261 @@
|
||||
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");
|
||||
},
|
||||
});
|
||||
}
|
||||
Reference in New Issue
Block a user