From b9ab5a392ec6723a7545a677f3b92fc3a695cc8a Mon Sep 17 00:00:00 2001 From: 0xrsydn Date: Tue, 13 Oct 2026 16:02:22 +0700 Subject: [PATCH] feat(pi): add Codex fast mode extension --- pi/extensions/openai-codex-fast/README.md | 81 +++++++++ pi/extensions/openai-codex-fast/index.ts | 71 ++++++++ pi/extensions/openai-codex-fast/state.ts | 53 ++++++ .../openai-codex-fast/tests/runtime.test.cjs | 167 ++++++++++++++++++ .../openai-codex-fast/tests/state.test.ts | 75 ++++++++ 5 files changed, 447 insertions(+) create mode 100644 pi/extensions/openai-codex-fast/README.md create mode 100644 pi/extensions/openai-codex-fast/index.ts create mode 100644 pi/extensions/openai-codex-fast/state.ts create mode 100644 pi/extensions/openai-codex-fast/tests/runtime.test.cjs create mode 100644 pi/extensions/openai-codex-fast/tests/state.test.ts diff --git a/pi/extensions/openai-codex-fast/README.md b/pi/extensions/openai-codex-fast/README.md new file mode 100644 index 0000000..bcc81bd --- /dev/null +++ b/pi/extensions/openai-codex-fast/README.md @@ -0,0 +1,81 @@ +# OpenAI Codex fast mode + +Adds `/fast` for Pi's `openai-codex` provider and `openai-codex-responses` API. +When enabled, the extension adds `service_tier: "priority"` to matching requests. +It keeps the selected model, reasoning level, tools, and conversation input. + +## Usage + +- `/fast` toggles fast mode. +- `/fast on` requests priority service. Usage may increase. +- `/fast off` stops this extension's override and restores provider defaults. +- `/fast status` reports the setting without changing it. + +New sessions start with fast mode off. Changes are stored as custom session +entries, outside model context. Resume and reload restore the active branch's +setting. Forks inherit only settings before the fork point. Switching to another +provider makes the setting inactive; switching back to Codex reactivates it. +Changes require an idle session. Status reports remain available during a turn. + +The footer shows `Codex fast: priority` while this extension requests priority +service. This shows the requested tier, not confirmation from OpenAI. + +## Availability and usage + +OpenAI controls priority access for each model and account. Priority service can +increase subscription usage or cost. The extension does not change reasoning +settings, query Codex's model catalog, verify account eligibility, or guarantee +faster responses. Unlike the Codex CLI, it does not hide the command for models +without an advertised fast tier. API failures use Pi's normal error handling; +the extension does not silently retry with another tier. + +Other providers are untouched, including `openai`, OpenRouter, and providers +that reuse the Codex API under another name. While enabled, this extension +replaces any existing `service_tier` value with `priority`. Disabling it leaves +native request settings and other extensions' overrides unchanged. + +Your `openai-server-compaction` extension can remain enabled. Both request +patches preserve unrelated fields. Fast mode applies to normal session requests, +not the compaction extension's separate HTTP compaction or summary calls. + +## Installation + +Home Manager deploys this directory through `rsydn.aiTools.piExtensions` along +with the other extensions. Apply your usual dotfiles rebuild, then run `/reload`. +To test before deploying: + +```sh +pi -e ./pi/extensions/openai-codex-fast/index.ts +``` + +## Codex implementation reference + +Reviewed OpenAI Codex commit +[`1a69ad3237fc`](https://github.com/openai/codex/tree/1a69ad3237fc485273ccfa9731239fdf789b43e4). + +- [Service-tier UI](https://github.com/openai/codex/blob/1a69ad3237fc485273ccfa9731239fdf789b43e4/codex-rs/tui/src/chatwidget/service_tiers.rs) + toggles the model's advertised Fast tier and persists the selection. +- [Model catalog](https://github.com/openai/codex/blob/1a69ad3237fc485273ccfa9731239fdf789b43e4/codex-rs/models-manager/models.json) + identifies Fast with the request value `priority`. +- [Request builder](https://github.com/openai/codex/blob/1a69ad3237fc485273ccfa9731239fdf789b43e4/codex-rs/core/src/client.rs) + sends the selected tier in `service_tier`. + +Pi's `before_provider_request` hook changes the native payload. No provider +replacement, custom transport, or additional runtime dependency is required. + +## Tests + +Requires Node.js 22.19 or newer and an installed `pi` executable. Tested with +Pi 1.1.0. + +```sh +node --experimental-strip-types --test \ + pi/extensions/openai-codex-fast/tests/*.test.ts \ + pi/extensions/openai-codex-fast/tests/*.test.cjs +``` + +Unit tests cover scope, payload preservation, branch state, and both compaction +patch orders. The runtime test uses Pi's real loader and native providers with a +local HTTP fixture. It checks outgoing payloads, commands, the idle guard, +provider switching, resume, and forks. It sends no requests to OpenAI and uses no +real credentials. diff --git a/pi/extensions/openai-codex-fast/index.ts b/pi/extensions/openai-codex-fast/index.ts new file mode 100644 index 0000000..23ad548 --- /dev/null +++ b/pi/extensions/openai-codex-fast/index.ts @@ -0,0 +1,71 @@ +import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent"; +import { + FAST_STATE_ENTRY, + isCodexModel, + parseFastCommand, + patchFastPayload, + restoreFastMode, +} from "./state.ts"; + +export default function openaiCodexFast(pi: ExtensionAPI) { + let enabled = false; + + function updateStatus(ctx: ExtensionContext, model = ctx.model): void { + if (!ctx.hasUI) return; + ctx.ui.setStatus( + FAST_STATE_ENTRY, + enabled && isCodexModel(model) ? "Codex fast: priority" : undefined, + ); + } + + function restore(ctx: ExtensionContext): void { + enabled = restoreFastMode(ctx.sessionManager.getBranch()); + updateStatus(ctx); + } + + pi.on("session_start", (_event, ctx) => restore(ctx)); + pi.on("session_tree", (_event, ctx) => restore(ctx)); + pi.on("model_select", (event, ctx) => updateStatus(ctx, event.model)); + + pi.on("before_provider_request", (event, ctx) => + patchFastPayload({ enabled, model: ctx.model, payload: event.payload }), + ); + + pi.registerCommand("fast", { + description: "Toggle Codex priority service. Usage: /fast [on|off|status]. May increase usage.", + handler: async (args, ctx) => { + const command = parseFastCommand(args); + if (!command) { + ctx.ui.notify("Usage: /fast [on|off|status]", "warning"); + return; + } + if (command === "status") { + const state = enabled + ? isCodexModel(ctx.model) ? "on, requesting priority service" : "on, inactive for this provider" + : "off, using provider defaults"; + ctx.ui.notify(`Codex fast mode is ${state}.`, "info"); + return; + } + if (!ctx.isIdle()) { + ctx.ui.notify("Wait for the current turn to finish before changing fast mode.", "warning"); + return; + } + const next = command === "toggle" ? !enabled : command === "on"; + if (next && !isCodexModel(ctx.model)) { + ctx.ui.notify("Fast mode requires the openai-codex provider and openai-codex-responses API.", "warning"); + return; + } + if (next !== enabled) { + pi.appendEntry(FAST_STATE_ENTRY, { enabled: next }); + enabled = next; + } + updateStatus(ctx); + ctx.ui.notify( + enabled + ? "Codex fast mode on. Requesting priority service; usage may increase. OpenAI controls availability." + : "Codex fast mode off. Using provider defaults.", + "info", + ); + }, + }); +} diff --git a/pi/extensions/openai-codex-fast/state.ts b/pi/extensions/openai-codex-fast/state.ts new file mode 100644 index 0000000..7ded097 --- /dev/null +++ b/pi/extensions/openai-codex-fast/state.ts @@ -0,0 +1,53 @@ +export const FAST_STATE_ENTRY = "openai-codex-fast"; + +export type FastCommand = "toggle" | "on" | "off" | "status"; + +type ModelIdentity = { provider: string; api: string; id: string }; + +export function parseFastCommand(args: string): FastCommand | undefined { + const command = args.trim().toLowerCase(); + if (command === "") return "toggle"; + if (command === "on" || command === "off" || command === "status") return command; + return undefined; +} + +export function isCodexModel(model: ModelIdentity | undefined): boolean { + return model?.provider === "openai-codex" && model.api === "openai-codex-responses"; +} + +export function restoreFastMode(branch: readonly unknown[]): boolean { + let enabled = false; + for (const entry of branch) { + if ( + typeof entry !== "object" || entry === null || + !("type" in entry) || entry.type !== "custom" || + !("customType" in entry) || entry.customType !== FAST_STATE_ENTRY || + !("data" in entry) + ) continue; + const data = entry.data; + if ( + typeof data === "object" && data !== null && + "enabled" in data && typeof data.enabled === "boolean" + ) enabled = data.enabled; + } + return enabled; +} + +export function patchFastPayload({ + enabled, + model, + payload, +}: { + enabled: boolean; + model: ModelIdentity | undefined; + payload: unknown; +}): unknown { + if ( + !enabled || !isCodexModel(model) || + typeof payload !== "object" || payload === null || Array.isArray(payload) || + !("model" in payload) || payload.model !== model?.id || + !("input" in payload) || !Array.isArray(payload.input) + ) return undefined; + + return { ...payload, service_tier: "priority" }; +} diff --git a/pi/extensions/openai-codex-fast/tests/runtime.test.cjs b/pi/extensions/openai-codex-fast/tests/runtime.test.cjs new file mode 100644 index 0000000..b8b809b --- /dev/null +++ b/pi/extensions/openai-codex-fast/tests/runtime.test.cjs @@ -0,0 +1,167 @@ +const assert = require("node:assert/strict"); +const { spawn } = require("node:child_process"); +const { once } = require("node:events"); +const fs = require("node:fs"); +const http = require("node:http"); +const os = require("node:os"); +const path = require("node:path"); +const test = require("node:test"); +const zlib = require("node:zlib"); + +test("installed Pi sends priority only for Codex and restores session state", { timeout: 60000 }, async () => { + const dir = fs.mkdtempSync(path.join(os.tmpdir(), "pi-codex-fast-test-")); + const requests = []; + let holdNextResponse; + const server = http.createServer(async (req, res) => { + try { + const chunks = []; + for await (const chunk of req) chunks.push(chunk); + let bytes = Buffer.concat(chunks); + if (req.headers["content-encoding"] === "zstd") bytes = zlib.zstdDecompressSync(bytes); + requests.push(JSON.parse(bytes.toString())); + const held = holdNextResponse; + holdNextResponse = undefined; + server.emit("fixture_request"); + if (held) await held; + res.writeHead(200, { "Content-Type": "text/event-stream" }); + const response = { + id: "resp_fixture", status: "completed", output: [], + usage: { input_tokens: 10, output_tokens: 0, total_tokens: 10, input_tokens_details: { cached_tokens: 0 } }, + }; + res.end(`data: ${JSON.stringify({ type: "response.completed", response })}\n\n`); + } catch (error) { + res.writeHead(500); + res.end(String(error)); + } + }); + server.listen(0, "127.0.0.1"); + await once(server, "listening"); + const baseUrl = `http://127.0.0.1:${server.address().port}`; + const token = `e30.${Buffer.from(JSON.stringify({ "https://api.openai.com/auth": { chatgpt_account_id: "fixture" } })).toString("base64")}.fixture`; + const model = { id: "gpt-fast-fixture", reasoning: true, input: ["text"], contextWindow: 128000, maxTokens: 4096 }; + fs.writeFileSync(path.join(dir, "settings.json"), JSON.stringify({ + transport: "sse", compaction: { enabled: false }, retry: { enabled: false }, enableInstallTelemetry: false, + })); + fs.writeFileSync(path.join(dir, "models.json"), JSON.stringify({ providers: { + "openai-codex": { baseUrl, api: "openai-codex-responses", apiKey: token, models: [model] }, + openai: { baseUrl, api: "openai-responses", apiKey: "fixture", models: [model] }, + } })); + const child = spawn("pi", [ + "--offline", "--no-extensions", "--no-mcp", "--no-skills", "--no-prompt-templates", "--no-context-files", + "-e", path.resolve(__dirname, "../index.ts"), + "-e", path.resolve(__dirname, "../../openai-server-compaction/index.ts"), + "--model", "openai-codex/gpt-fast-fixture", "--thinking", "high", "--no-tools", "--mode", "rpc", + ], { cwd: dir, env: { ...process.env, PI_CODING_AGENT_DIR: dir } }); + let buffer = "", stderr = "", seq = 0; + const events = [], waiters = new Set(); + child.stderr.on("data", chunk => { stderr += chunk; }); + child.stdout.on("data", chunk => { + buffer += chunk; + let newline; + while ((newline = buffer.indexOf("\n")) !== -1) { + const line = buffer.slice(0, newline); + buffer = buffer.slice(newline + 1); + let record; + try { record = JSON.parse(line); } catch { continue; } + events.push(record); + for (const waiter of [...waiters]) if (waiter.predicate(record)) { + waiters.delete(waiter); + clearTimeout(waiter.timer); + waiter.resolve(record); + } + } + }); + function wait(predicate) { + return new Promise((resolve, reject) => { + const waiter = { predicate, resolve, timer: setTimeout(() => { + waiters.delete(waiter); + reject(new Error(`RPC timeout. stderr: ${stderr}`)); + }, 10000) }; + waiters.add(waiter); + }); + } + async function command(type, fields = {}) { + const id = String(++seq); + const pending = wait(record => record.type === "response" && record.id === id); + child.stdin.write(JSON.stringify({ id, type, ...fields }) + "\n"); + const record = await pending; + assert.equal(record.success, true, JSON.stringify(record)); + return record.data; + } + async function slash(message) { + const start = events.length; + await command("prompt", { message }); + return events.slice(start).filter(record => record.type === "extension_ui_request" && record.method === "notify").map(record => record.message).join("\n"); + } + async function turn(message) { + const start = events.length; + const settled = wait(record => record.type === "agent_settled"); + await command("prompt", { message }); + await settled; + const answer = events.slice(start).find(record => record.type === "message_end" && record.message?.role === "assistant"); + assert.equal(answer?.message.stopReason, "stop", JSON.stringify(answer)); + return requests.at(-1); + } + try { + const commands = await command("get_commands"); + assert.ok(commands.commands.some(item => item.name === "fast")); + assert.match(await slash("/fast status"), /off/); + const first = await turn("BEFORE_FAST"); + assert.equal(first.service_tier, undefined); + const sessionFile = (await command("get_state")).sessionFile; + assert.match(await slash("/fast on"), /usage may increase/); + const fast = await turn("FAST_REQUEST"); + assert.equal(fast.service_tier, "priority"); + assert.equal(fast.model, first.model); + assert.deepEqual(fast.reasoning, first.reasoning); + assert.match(await slash("/fast on extra"), /Usage:/); + assert.match(await slash("/fast status"), /on, requesting priority/); + await command("set_model", { provider: "openai", modelId: model.id }); + assert.equal((await turn("OTHER_PROVIDER")).service_tier, undefined); + assert.match(await slash("/fast status"), /inactive/); + await slash("/fast off"); + assert.match(await slash("/fast on"), /requires the openai-codex provider/); + await command("set_model", { provider: "openai-codex", modelId: model.id }); + assert.equal((await turn("DISABLED")).service_tier, undefined); + await slash("/fast"); + assert.equal((await turn("TOGGLED_ON")).service_tier, "priority"); + const snapshots = (await command("get_entries")).entries.filter(entry => entry.type === "custom" && entry.customType === "openai-codex-fast"); + assert.deepEqual(snapshots.map(entry => entry.data.enabled), [true, false, true]); + let releaseResponse; + holdNextResponse = new Promise(resolve => { releaseResponse = resolve; }); + const received = once(server, "fixture_request"); + const settled = wait(record => record.type === "agent_settled"); + await command("prompt", { message: "BUSY_TURN" }); + await received; + try { + assert.match(await slash("/fast off"), /Wait for the current turn/); + assert.match(await slash("/fast status"), /on, requesting priority/); + } finally { + releaseResponse(); + } + await settled; + assert.match(await slash("/fast"), /off/); + assert.match(await slash("/fast"), /on/); + await command("new_session"); + assert.match(await slash("/fast status"), /off/); + await command("switch_session", { sessionPath: sessionFile }); + assert.match(await slash("/fast status"), /on, requesting priority/); + assert.equal((await turn("RESUMED")).service_tier, "priority"); + const forks = await command("get_fork_messages"); + await command("fork", { entryId: forks.messages.find(message => message.text === "BEFORE_FAST").entryId }); + assert.match(await slash("/fast status"), /off/); + const entries = await command("get_entries"); + assert.ok(!entries.entries.some(entry => entry.type === "custom" && entry.customType === "openai-codex-fast")); + assert.doesNotMatch(stderr, /Failed to load extension|Extension error/); + } finally { + for (const waiter of waiters) clearTimeout(waiter.timer); + const closed = once(child, "close"); + child.stdin.end(); + const timer = setTimeout(() => child.kill("SIGKILL"), 3000); + await closed; + clearTimeout(timer); + server.closeAllConnections(); + await new Promise(resolve => server.close(resolve)); + fs.rmSync(dir, { recursive: true, force: true }); + } +}); diff --git a/pi/extensions/openai-codex-fast/tests/state.test.ts b/pi/extensions/openai-codex-fast/tests/state.test.ts new file mode 100644 index 0000000..a20fedc --- /dev/null +++ b/pi/extensions/openai-codex-fast/tests/state.test.ts @@ -0,0 +1,75 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { isRecord } from "../../openai-server-compaction/config.ts"; +import { applyRemoteHistoryPayloadPatch } from "../../openai-server-compaction/openai.ts"; +import { FAST_STATE_ENTRY, isCodexModel, parseFastCommand, patchFastPayload, restoreFastMode } from "../state.ts"; + +const model = { provider: "openai-codex", api: "openai-codex-responses", id: "gpt-5.5" }; +const payload = { + model: model.id, + input: [{ role: "user", content: "Hello" }], + reasoning: { effort: "high", summary: "auto" }, + store: false, + stream: true, +}; +const entry = (enabled: unknown) => ({ type: "custom", customType: FAST_STATE_ENTRY, data: { enabled } }); + +test("parses toggle, explicit settings, and status without accepting extra arguments", () => { + assert.equal(parseFastCommand(""), "toggle"); + assert.equal(parseFastCommand(" ON "), "on"); + assert.equal(parseFastCommand("off"), "off"); + assert.equal(parseFastCommand("status"), "status"); + for (const args of ["on extra", "toggle", "true"]) assert.equal(parseFastCommand(args), undefined); +}); + +test("targets only the exact Codex provider and API", () => { + assert.equal(isCodexModel(model), true); + assert.equal(isCodexModel(undefined), false); + assert.equal(isCodexModel({ ...model, provider: "openai" }), false); + assert.equal(isCodexModel({ ...model, api: "openai-responses" }), false); + assert.equal(isCodexModel({ ...model, provider: "proxy" }), false); +}); + +test("adds priority without changing model, reasoning, input, or the original payload", () => { + const frozen = Object.freeze({ ...payload, service_tier: "flex" }); + assert.deepEqual(patchFastPayload({ enabled: true, model, payload: frozen }), { + ...frozen, service_tier: "priority", + }); + assert.equal(frozen.service_tier, "flex"); +}); + +test("leaves disabled, unrelated, malformed, and mismatched requests unchanged", () => { + assert.equal(patchFastPayload({ enabled: false, model, payload }), undefined); + assert.equal(patchFastPayload({ enabled: true, model: undefined, payload }), undefined); + assert.equal(patchFastPayload({ enabled: true, model: { ...model, provider: "openai" }, payload }), undefined); + assert.equal(patchFastPayload({ enabled: true, model: { ...model, api: "openai-responses" }, payload }), undefined); + for (const value of [null, 1, "payload", [], {}, { ...payload, model: "other" }, { model: model.id }, { ...payload, input: {} }]) { + assert.equal(patchFastPayload({ enabled: true, model, payload: value }), undefined); + } +}); + +test("restores the last valid state on the active branch and defaults new sessions to off", () => { + assert.equal(restoreFastMode([]), false); + assert.equal(restoreFastMode([entry(true)]), true); + assert.equal(restoreFastMode([entry(true), entry(false)]), false); + assert.equal(restoreFastMode([entry(false), entry(true)]), true); + assert.equal(restoreFastMode([entry(true), { type: "compaction" }]), true); + assert.equal(restoreFastMode([entry(true)].slice(0, 0)), false); +}); + +test("ignores invalid snapshots and entries owned by other extensions", () => { + const invalid = [null, 1, {}, entry("true"), entry(null), { ...entry(false), customType: "other" }, { ...entry(false), type: "custom_message" }]; + assert.equal(restoreFastMode(invalid), false); + assert.equal(restoreFastMode([entry(true), ...invalid]), true); +}); + +test("composes with the existing compaction payload patch in either order", () => { + const explicitHistory = [{ type: "compaction", encrypted_content: "fixture" }]; + const request = { ...payload, previous_response_id: "old" }; + const compacted = applyRemoteHistoryPayloadPatch({ payload: request, explicitHistory }); + const expected = { ...payload, input: explicitHistory, service_tier: "priority" }; + assert.deepEqual(patchFastPayload({ enabled: true, model, payload: compacted }), expected); + const fast = patchFastPayload({ enabled: true, model, payload: request }); + assert.ok(isRecord(fast)); + assert.deepEqual(applyRemoteHistoryPayloadPatch({ payload: fast, explicitHistory }), expected); +});