feat(pi): add Codex fast mode extension
This commit is contained in:
parent
9fdb2d7b7d
commit
b9ab5a392e
5 changed files with 447 additions and 0 deletions
167
pi/extensions/openai-codex-fast/tests/runtime.test.cjs
Normal file
167
pi/extensions/openai-codex-fast/tests/runtime.test.cjs
Normal file
|
|
@ -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 });
|
||||
}
|
||||
});
|
||||
75
pi/extensions/openai-codex-fast/tests/state.test.ts
Normal file
75
pi/extensions/openai-codex-fast/tests/state.test.ts
Normal file
|
|
@ -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);
|
||||
});
|
||||
Loading…
Add table
Add a link
Reference in a new issue