diff --git a/.changeset/codex-controls.md b/.changeset/codex-controls.md new file mode 100644 index 0000000..aa867d7 --- /dev/null +++ b/.changeset/codex-controls.md @@ -0,0 +1,5 @@ +--- +"@eleith-pi/codex": minor +--- + +Rename context-codex to codex while preserving context controls and branch state. Add Codex account usage reporting and explicit branch-local, per-model fast controls using Pi's existing login, conservative catalog eligibility, and higher-consumption confirmation for priority requests on supported direct legacy Codex selections. diff --git a/extensions/codex/src/capabilities.test.ts b/extensions/codex/src/capabilities.test.ts new file mode 100644 index 0000000..4c3cecb --- /dev/null +++ b/extensions/codex/src/capabilities.test.ts @@ -0,0 +1,74 @@ +import { expect, it } from "vitest"; +import { codexFastEligibility } from "./capabilities.ts"; + +const modelId = "future-codex-model"; + +it.each(["priority", "fast"])( + "recognizes the advertised %s service tier without guessing model names", + (id) => { + const catalog = { models: [{ slug: modelId, service_tiers: [{ id }] }] }; + expect(codexFastEligibility(catalog, modelId)).toBe("supported"); + }, +); + +it("supports the legacy fast advertisement even alongside an empty newer field", () => { + const catalog = { + models: [{ slug: modelId, service_tiers: [], additional_speed_tiers: ["fast"] }], + }; + expect(codexFastEligibility(catalog, modelId)).toBe("supported"); +}); + +it("accepts legacy-only catalogs", () => { + const catalog = { models: [{ slug: modelId, additional_speed_tiers: ["fast"] }] }; + expect(codexFastEligibility(catalog, modelId)).toBe("supported"); +}); + +it.each([ + { service_tiers: [] }, + { service_tiers: [{ id: "default" }, { id: "ultrafast" }] }, + { additional_speed_tiers: [] }, +])("distinguishes a valid list without fast mode from unknown metadata: %j", (metadata) => { + expect(codexFastEligibility({ models: [{ slug: modelId, ...metadata }] }, modelId)).toBe( + "unsupported", + ); +}); + +it.each([ + {}, + { default_service_tier: "priority" }, + { service_tiers: ["priority"] }, + { service_tiers: [{ id: "priority" }, null] }, + { service_tiers: [], additional_speed_tiers: ["fast", null] }, + { service_tiers: [{ id: "priority" }], additional_speed_tiers: ["fast", null] }, + { service_tiers: [null], additional_speed_tiers: ["fast"] }, +])("does not infer eligibility from missing or malformed fields: %j", (metadata) => { + expect(codexFastEligibility({ models: [{ slug: modelId, ...metadata }] }, modelId)).toBe( + "unknown", + ); +}); + +it("does not apply another model's capability to the selected model", () => { + const catalog = { + models: [ + { slug: "other-model", service_tiers: [{ id: "priority" }] }, + { slug: modelId, service_tiers: [] }, + ], + }; + expect(codexFastEligibility(catalog, modelId)).toBe("unsupported"); + expect(codexFastEligibility(catalog, "missing-model")).toBe("unknown"); +}); + +it("treats duplicate model entries as ambiguous", () => { + const entry = { slug: modelId, service_tiers: [{ id: "priority" }] }; + expect(codexFastEligibility({ models: [entry, entry] }, modelId)).toBe("unknown"); +}); + +it.each([ + null, + [], + {}, + { models: {} }, + { models: [{ id: modelId, service_tiers: [{ id: "priority" }] }] }, +])("keeps unrecognized catalog schemas unknown: %j", (catalog) => { + expect(codexFastEligibility(catalog, modelId)).toBe("unknown"); +}); diff --git a/extensions/codex/src/capabilities.ts b/extensions/codex/src/capabilities.ts new file mode 100644 index 0000000..34f4cbc --- /dev/null +++ b/extensions/codex/src/capabilities.ts @@ -0,0 +1,55 @@ +type CatalogModel = Record; +export type FastEligibility = "supported" | "unsupported" | "unknown"; + +export function codexFastEligibility(catalog: unknown, modelId: string): FastEligibility { + const model = findCatalogModel(catalog, modelId); + if (!model) return "unknown"; + + const serviceTierSupport = readServiceTierSupport(model.service_tiers); + const legacyTierSupport = readLegacyTierSupport(model.additional_speed_tiers); + const invalidServiceTiers = "service_tiers" in model && serviceTierSupport === undefined; + const invalidLegacyTiers = "additional_speed_tiers" in model && legacyTierSupport === undefined; + if (invalidServiceTiers || invalidLegacyTiers) return "unknown"; + + if (serviceTierSupport === true || legacyTierSupport === true) return "supported"; + if (serviceTierSupport === false || legacyTierSupport === false) return "unsupported"; + return "unknown"; +} + +function findCatalogModel(catalog: unknown, modelId: string): CatalogModel | undefined { + if (!modelId) return undefined; + const models = readObject(catalog)?.models; + if (!Array.isArray(models)) return undefined; + + let matched: CatalogModel | undefined; + for (const value of models) { + const model = readObject(value); + if (model?.slug !== modelId) continue; + if (matched) return undefined; + matched = model; + } + return matched; +} + +function readServiceTierSupport(value: unknown): boolean | undefined { + if (!Array.isArray(value)) return undefined; + let supported = false; + for (const item of value) { + const tier = readObject(item); + if (typeof tier?.id !== "string" || !tier.id.trim()) return undefined; + if (tier.id === "priority" || tier.id === "fast") supported = true; + } + return supported; +} + +function readLegacyTierSupport(value: unknown): boolean | undefined { + if (!Array.isArray(value)) return undefined; + if (!value.every((tier) => typeof tier === "string" && tier.trim().length > 0)) return undefined; + return value.includes("fast"); +} + +function readObject(value: unknown): CatalogModel | undefined { + return value !== null && typeof value === "object" && !Array.isArray(value) + ? (value as CatalogModel) + : undefined; +} diff --git a/extensions/codex/src/client.test.ts b/extensions/codex/src/client.test.ts index 1572cef..c36f422 100644 --- a/extensions/codex/src/client.test.ts +++ b/extensions/codex/src/client.test.ts @@ -1,6 +1,12 @@ +import { createHash } from "node:crypto"; import type { Api, AuthResult, Model } from "@earendil-works/pi-ai"; -import { afterEach, expect, it, vi } from "vitest"; -import { fetchCodexUsage, type CodexAccountRegistry } from "./client.ts"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { + fetchCodexModelCatalog, + fetchCodexUsage, + readCodexAuthScope, + type CodexAccountRegistry, +} from "./client.ts"; const BASE_URL = "https://chatgpt.com/backend-api"; const model: Model = { @@ -66,6 +72,142 @@ it("uses Pi's existing Codex auth for the fixed account endpoint", async () => { }); }); +it("uses the same Pi credentials for the fixed versioned model catalog", async () => { + const f = fixture(); + const catalog = { models: [{ slug: "future-model", service_tiers: [{ id: "priority" }] }] }; + f.fetcher.mockResolvedValue(json(catalog)); + const result = await fetchCodexModelCatalog(f.registry, undefined, f.fetcher); + expect(result.data).toEqual(catalog); + expect(f.registry.getProviderAuth).toHaveBeenCalledExactlyOnceWith("openai-codex"); + expect(f.fetcher).toHaveBeenCalledExactlyOnceWith( + `${BASE_URL}/codex/models?client_version=0.162.0`, + expect.objectContaining({ method: "GET", redirect: "error" }), + ); + const headers = new Headers(f.fetcher.mock.calls[0][1]?.headers); + expect(headers.get("authorization")).toBe(`Bearer ${f.auth.auth.apiKey}`); + expect(headers.get("chatgpt-account-id")).toBe("test-account"); + expect(result.authScope).toBe( + createHash("sha256") + .update(JSON.stringify([headers.get("authorization"), headers.get("chatgpt-account-id")])) + .digest("hex"), + ); + expect(Object.keys(result).sort()).toEqual(["authScope", "data"]); + expect(JSON.stringify(result)).not.toContain(f.auth.auth.apiKey); + expect(JSON.stringify(result)).not.toContain("test-account"); +}); + +it("reads a stable auth scope without fetching and ignores unused provider headers", async () => { + const f = fixture(); + const fetchSpy = vi.spyOn(globalThis, "fetch").mockRejectedValue(new Error("must not fetch")); + try { + const scope = await readCodexAuthScope(f.registry); + f.auth.auth.headers = { + "ChatGPT-Account-Id": "unused-account", + Authorization: "unused-token", + "X-Unrelated-Secret": "unused-secret", + }; + expect(await readCodexAuthScope(f.registry)).toBe(scope); + expect(scope).toMatch(/^[a-f0-9]{64}$/); + expect(scope).not.toContain(f.auth.auth.apiKey); + expect(scope).not.toContain("test-account"); + expect(f.registry.getProviderAuth).toHaveBeenCalledTimes(2); + expect(fetchSpy).not.toHaveBeenCalled(); + const catalog = await fetchCodexModelCatalog(f.registry, undefined, f.fetcher); + expect(catalog.authScope).toBe(scope); + } finally { + fetchSpy.mockRestore(); + } +}); + +it.each(["token refresh", "account change"] as const)( + "changes the scope on %s and binds the catalog to its own resolved credentials", + async (change) => { + const f = fixture(); + const previous = await readCodexAuthScope(f.registry); + f.auth.auth.apiKey = + change === "token refresh" ? `${token()}-refreshed` : token("another-account"); + const catalog = await fetchCodexModelCatalog(f.registry, undefined, f.fetcher); + expect(catalog.authScope).not.toBe(previous); + expect(await readCodexAuthScope(f.registry)).toBe(catalog.authScope); + }, +); + +describe("auth-only validation without HTTP", () => { + const unexpectedFetch = vi.fn(async () => { + throw new Error("must not fetch"); + }); + beforeEach(() => { + unexpectedFetch.mockClear(); + vi.stubGlobal("fetch", unexpectedFetch); + }); + afterEach(() => { + vi.unstubAllGlobals(); + expect(unexpectedFetch).not.toHaveBeenCalled(); + }); + const unavailableAuth = { + code: "auth", + message: "Codex account access is unavailable. Check your Pi Codex login.", + }; + const unsupportedProvider = { + code: "unsupported", + message: "Account requests require the official Codex provider.", + }; + + it("rejects an overridden provider endpoint", async () => { + const account = fixture(); + account.provider.baseUrl = "https://gateway.invalid"; + await expect(readCodexAuthScope(account.registry)).rejects.toMatchObject(unsupportedProvider); + }); + + it("rejects a catalog model at a different endpoint", async () => { + const account = fixture(); + account.provider.getModels = () => [{ ...model, baseUrl: "https://gateway.invalid" }]; + await expect(readCodexAuthScope(account.registry)).rejects.toMatchObject(unsupportedProvider); + }); + + it("rejects an endpoint override in resolved credentials", async () => { + const account = fixture(); + account.auth.auth.baseUrl = "https://gateway.invalid"; + await expect(readCodexAuthScope(account.registry)).rejects.toMatchObject({ + code: "unsupported", + message: "Account requests do not support endpoint overrides.", + }); + }); + + it("requires an OAuth availability snapshot", async () => { + const account = fixture(); + account.registry.isUsingOAuth.mockReturnValue(false); + await expect(readCodexAuthScope(account.registry)).rejects.toMatchObject(unavailableAuth); + }); + + it("rejects a non-OAuth resolution even when the snapshot says OAuth", async () => { + const account = fixture(); + account.auth.source = "models.json"; + await expect(readCodexAuthScope(account.registry)).rejects.toMatchObject(unavailableAuth); + }); + + it("requires the token's account-routing claim", async () => { + const account = fixture(); + account.auth.auth.apiKey = "e30.e30.private-token"; + await expect(readCodexAuthScope(account.registry)).rejects.toMatchObject(unavailableAuth); + }); + + it("does not repeat private Pi refresh errors", async () => { + const account = fixture(); + account.registry.getProviderAuth.mockRejectedValue(new Error("private refresh details")); + await expect(readCodexAuthScope(account.registry)).rejects.toMatchObject(unavailableAuth); + }); +}); + +it("does not forward catalog credentials through an endpoint override", async () => { + const f = fixture(); + f.auth.auth.baseUrl = "https://gateway.invalid"; + await expect(fetchCodexModelCatalog(f.registry, undefined, f.fetcher)).rejects.toMatchObject({ + code: "unsupported", + }); + expect(f.fetcher).not.toHaveBeenCalled(); +}); + it("derives the account from the token like Pi's legacy Codex transport", async () => { const f = fixture(); f.auth.auth.headers = { "ChatGPT-Account-Id": "selected-account" }; @@ -207,23 +349,29 @@ it("does not follow or display a redirected endpoint", async () => { expect(f.fetcher).toHaveBeenCalledOnce(); }); -it("cancels a chunked response before retaining more than the byte limit", async () => { - const f = fixture(); - const cancel = vi.fn(); - const body = new ReadableStream({ - start(controller) { - controller.enqueue(new Uint8Array(1024 * 1024)); - controller.enqueue(new Uint8Array(1)); - }, - cancel, - }); - f.fetcher.mockResolvedValue(new Response(body)); - await expect(fetchCodexUsage(f.registry, undefined, f.fetcher)).rejects.toMatchObject({ - code: "response", - message: "Codex account response was too large.", - }); - expect(cancel).toHaveBeenCalledOnce(); -}); +it.each([ + ["usage", fetchCodexUsage], + ["catalog", fetchCodexModelCatalog], +] as const)( + "bounds chunked %s responses before retaining more than the byte limit", + async (_resource, fetchData) => { + const f = fixture(); + const cancel = vi.fn(); + const body = new ReadableStream({ + start(controller) { + controller.enqueue(new Uint8Array(1024 * 1024)); + controller.enqueue(new Uint8Array(1)); + }, + cancel, + }); + f.fetcher.mockResolvedValue(new Response(body)); + await expect(fetchData(f.registry, undefined, f.fetcher)).rejects.toMatchObject({ + code: "response", + message: "Codex account response was too large.", + }); + expect(cancel).toHaveBeenCalledOnce(); + }, +); it("reports invalid JSON without repeating returned content", async () => { const f = fixture(); @@ -234,6 +382,65 @@ it("reports invalid JSON without repeating returned content", async () => { }); }); +it.each(["usage", "catalog", "auth"] as const)( + "cancels during the public %s wrapper handoff", + async (resource) => { + const f = fixture(); + const controller = new AbortController(); + const abort = AbortController.prototype.abort; + // The private flow aborts its deadline in finally, immediately before handing off. + const cleanup = vi.spyOn(AbortController.prototype, "abort").mockImplementation(function ( + this: AbortController, + reason?: unknown, + ) { + abort.call(this, reason); + if (this !== controller) abort.call(controller); + }); + try { + let request: Promise; + if (resource === "auth") request = readCodexAuthScope(f.registry, controller.signal); + else if (resource === "usage") + request = fetchCodexUsage(f.registry, controller.signal, f.fetcher); + else request = fetchCodexModelCatalog(f.registry, controller.signal, f.fetcher); + await expect(request).rejects.toMatchObject({ code: "cancelled" }); + } finally { + cleanup.mockRestore(); + } + }, +); + +it("does not resolve auth for an already-cancelled scope read", async () => { + const f = fixture(); + const controller = new AbortController(); + controller.abort(); + await expect(readCodexAuthScope(f.registry, controller.signal)).rejects.toMatchObject({ + code: "cancelled", + }); + expect(f.registry.getProviderAuth).not.toHaveBeenCalled(); +}); + +it("ignores auth resolving after a scope read is cancelled", async () => { + const f = fixture(); + const auth = deferred(); + f.registry.getProviderAuth.mockReturnValue(auth.promise); + const controller = new AbortController(); + const request = readCodexAuthScope(f.registry, controller.signal); + controller.abort(); + await expect(request).rejects.toMatchObject({ code: "cancelled" }); + auth.resolve(f.auth); + await auth.promise; +}); + +it("bounds auth-only resolution by the same ten-second deadline", async () => { + vi.useFakeTimers(); + const f = fixture(); + f.registry.getProviderAuth.mockReturnValue(new Promise(() => {})); + const result = expect(readCodexAuthScope(f.registry)).rejects.toMatchObject({ code: "timeout" }); + await vi.advanceTimersByTimeAsync(10_000); + await result; + expect(vi.getTimerCount()).toBe(0); +}); + it("cancels during handoff of an already-buffered response", async () => { const f = fixture(); const started = deferred(); @@ -321,16 +528,22 @@ it("bounds the wait for response headers", async () => { expect(vi.getTimerCount()).toBe(0); }); -it("keeps the deadline active while consuming the response body", async () => { - vi.useFakeTimers(); - const f = fixture(); - const cancel = vi.fn(); - f.fetcher.mockResolvedValue(new Response(new ReadableStream({ cancel }))); - const result = expect(fetchCodexUsage(f.registry, undefined, f.fetcher)).rejects.toMatchObject({ - code: "timeout", - }); - await vi.advanceTimersByTimeAsync(10_000); - await result; - expect(cancel).toHaveBeenCalledOnce(); - expect(vi.getTimerCount()).toBe(0); -}); +it.each([ + ["usage", fetchCodexUsage], + ["catalog", fetchCodexModelCatalog], +] as const)( + "keeps the deadline active while consuming the %s body", + async (_resource, fetchData) => { + vi.useFakeTimers(); + const f = fixture(); + const cancel = vi.fn(); + f.fetcher.mockResolvedValue(new Response(new ReadableStream({ cancel }))); + const result = expect(fetchData(f.registry, undefined, f.fetcher)).rejects.toMatchObject({ + code: "timeout", + }); + await vi.advanceTimersByTimeAsync(10_000); + await result; + expect(cancel).toHaveBeenCalledOnce(); + expect(vi.getTimerCount()).toBe(0); + }, +); diff --git a/extensions/codex/src/client.ts b/extensions/codex/src/client.ts index 172471f..043ffda 100644 --- a/extensions/codex/src/client.ts +++ b/extensions/codex/src/client.ts @@ -1,9 +1,15 @@ +import { createHash } from "node:crypto"; import type { AuthResult, Provider } from "@earendil-works/pi-ai"; import type { ExtensionContext } from "@earendil-works/pi-coding-agent"; const PROVIDER = "openai-codex"; const BASE_URL = "https://chatgpt.com/backend-api"; -const USAGE_URL = `${BASE_URL}/wham/usage`; +const ACCOUNT_URLS = { + usage: `${BASE_URL}/wham/usage`, + catalog: `${BASE_URL}/codex/models?client_version=0.162.0`, +} as const; +type NetworkResource = keyof typeof ACCOUNT_URLS; +type AccountResource = NetworkResource | "auth"; const TIMEOUT_MS = 10_000; const MAX_RESPONSE_BYTES = 1024 * 1024; @@ -14,6 +20,13 @@ export type CodexAccountRegistry = Pick< getProvider(id: string): Pick | undefined; }; +interface AccountData { + data: unknown; + authScope: string; +} + +export type CodexModelCatalog = AccountData; + type Failure = "auth" | "unsupported" | "cancelled" | "timeout" | "network" | "response" | "http"; export class CodexAccountError extends Error { @@ -33,6 +46,36 @@ export async function fetchCodexUsage( signal?: AbortSignal, fetcher: typeof fetch = fetch, ): Promise { + const result = await readAccountData("usage", registry, signal, fetcher); + if (signal?.aborted) throw new CodexAccountError("cancelled", "Codex account request cancelled."); + return result.data; +} + +export async function fetchCodexModelCatalog( + registry: CodexAccountRegistry, + signal?: AbortSignal, + fetcher: typeof fetch = fetch, +): Promise { + const result = await readAccountData("catalog", registry, signal, fetcher); + if (signal?.aborted) throw new CodexAccountError("cancelled", "Codex account request cancelled."); + return result; +} + +export async function readCodexAuthScope( + registry: CodexAccountRegistry, + signal?: AbortSignal, +): Promise { + const result = await readAccountData("auth", registry, signal); + if (signal?.aborted) throw new CodexAccountError("cancelled", "Codex account request cancelled."); + return result.authScope; +} + +async function readAccountData( + resource: AccountResource, + registry: CodexAccountRegistry, + signal: AbortSignal | undefined, + fetcher: typeof fetch = fetch, +): Promise { const deadline = new AbortController(); const timer = setTimeout(() => deadline.abort(), TIMEOUT_MS); const requestSignal = signal ? AbortSignal.any([signal, deadline.signal]) : deadline.signal; @@ -40,9 +83,14 @@ export async function fetchCodexUsage( requestSignal.throwIfAborted(); const headers = await resolveAccountHeaders(registry, requestSignal); requestSignal.throwIfAborted(); - const data = await requestUsage(headers, requestSignal, fetcher); + // Hash only the normalized credentials we actually send, never persist raw headers. + const authScope = createHash("sha256") + .update(JSON.stringify([headers.get("authorization"), headers.get("chatgpt-account-id")])) + .digest("hex"); + if (resource === "auth") return { data: undefined, authScope }; + const data = await requestAccountJson(resource, headers, requestSignal, fetcher); requestSignal.throwIfAborted(); - return data; + return { data, authScope }; } catch (error) { if (signal?.aborted) throw new CodexAccountError("cancelled", "Codex account request cancelled."); @@ -73,7 +121,7 @@ async function resolveAccountHeaders( if (!officialBase(provider?.baseUrl) || !model) { throw new CodexAccountError( "unsupported", - "Account usage requires the official Codex provider.", + "Account requests require the official Codex provider.", ); } if (!registry.isUsingOAuth(model)) throw authError(); @@ -94,7 +142,7 @@ async function resolveAccountHeaders( ) { throw new CodexAccountError( "unsupported", - "Account usage does not support endpoint overrides.", + "Account requests do not support endpoint overrides.", ); } const accessToken = resolvedAuth.auth.apiKey; @@ -145,14 +193,15 @@ function authError(): CodexAccountError { ); } -async function requestUsage( +async function requestAccountJson( + resource: NetworkResource, headers: Headers, requestSignal: AbortSignal, fetcher: typeof fetch, ): Promise { let receivedResponse: Response | undefined; try { - const request = fetcher(USAGE_URL, { + const request = fetcher(ACCOUNT_URLS[resource], { method: "GET", headers, redirect: "error", diff --git a/extensions/codex/src/fast-sdk.test.ts b/extensions/codex/src/fast-sdk.test.ts new file mode 100644 index 0000000..d5b97c2 --- /dev/null +++ b/extensions/codex/src/fast-sdk.test.ts @@ -0,0 +1,423 @@ +import { rmSync } from "node:fs"; +import { zstdDecompressSync } from "node:zlib"; +import { afterAll, afterEach, beforeEach, expect, it, vi } from "vitest"; +import { + InMemoryModelsStore, + type Api, + type CredentialStore, + type Model, + type OAuthCredential, +} from "@earendil-works/pi-ai"; +import { + createAgentSession, + DefaultResourceLoader, + ModelRuntime, + SessionManager, + SettingsManager, + type ExtensionUIContext, +} from "@earendil-works/pi-coding-agent"; +import codex from "./index.ts"; + +const directory = await vi.hoisted(async () => { + const { mkdtempSync } = await import("node:fs"); + const { tmpdir } = await import("node:os"); + const { join } = await import("node:path"); + const directory = mkdtempSync(join(tmpdir(), "codex-fast-sdk-")); + vi.stubEnv("PI_CODING_AGENT_DIR", directory); + vi.stubEnv("PI_OFFLINE", "1"); + return directory; +}); + +const CATALOG_URL = "https://chatgpt.com/backend-api/codex/models?client_version=0.162.0"; +const RESPONSES_URL = "https://chatgpt.com/backend-api/codex/responses"; +const SOCKET_URL = "wss://chatgpt.com/backend-api/codex/responses"; +const FAST_ENTRY = "eleith-codex-fast"; +type Transport = "sse" | "websocket"; + +beforeEach(blockUnexpectedNetwork); +afterEach(() => { + vi.restoreAllMocks(); + vi.unstubAllGlobals(); +}); +afterAll(() => { + vi.unstubAllEnvs(); + rmSync(directory, { recursive: true, force: true }); +}); + +it.each(["sse", "websocket"] as const)( + "SDK fast on reaches the native Codex %s wire, independently of the observed tier", + async (transport) => { + const network = installResponseTransports(transport); + const { runtime, model, accessToken, credentialWrites, resolveAuth, modelTurns } = + await createOfflineModelRuntime(transport, network.fetch); + const { session, manager, confirm, notify, errors, closeSession } = + await openRealExtensionSession(runtime, model); + try { + const branch = structuredClone(manager.getBranch()); + await session.prompt("/eleith:codex fast"); + await session.prompt("/eleith:codex fast status"); + expect(network.fetch).not.toHaveBeenCalled(); + expect(manager.getBranch()).toEqual(branch); + expect(modelTurns).not.toHaveBeenCalled(); + + await session.prompt("/eleith:codex fast on"); + expect(confirm).toHaveBeenCalledOnce(); + expect(network.catalog).toHaveBeenCalledOnce(); + expect(network.catalog.mock.calls[0][0]).toBe(CATALOG_URL); + const headers = new Headers(network.catalog.mock.calls[0][1]?.headers); + expect(headers.get("authorization")).toBe(`Bearer ${accessToken}`); + expect(headers.get("chatgpt-account-id")).toBe("fast-sdk-account"); + expect(fastChoices(manager)).toEqual([ + { provider: "openai-codex", modelId: model.id, mode: "on" }, + ]); + expect(modelTurns).not.toHaveBeenCalled(); + + resolveAuth.mockClear(); + await session.prompt("Reply with ok."); + // Ordinary input checks the cached catalog's auth scope through Pi resolution only. + expect(resolveAuth).toHaveBeenCalledWith("openai-codex"); + expect(network.catalog).toHaveBeenCalledOnce(); + expect(modelTurns).toHaveBeenCalledOnce(); + expect(network.payloads).toEqual([ + expect.objectContaining({ model: model.id, service_tier: "priority" }), + ]); + expect(session.getLastAssistantText()).toBe("ok"); + expect(network.sse).toHaveBeenCalledTimes(transport === "sse" ? 1 : 0); + expect(network.sockets).toHaveLength(transport === "websocket" ? 1 : 0); + if (transport === "websocket") { + expect(network.payloads[0]).toHaveProperty("type", "response.create"); + expect(network.sockets[0].readyState).toBe(3); + } + + const requests = network.fetch.mock.calls.length; + resolveAuth.mockClear(); + const afterTurn = structuredClone(manager.getBranch()); + for (const command of ["/eleith:codex fast status", "/eleith:codex fast"]) { + await session.prompt(command); + expect(notify).toHaveBeenLastCalledWith(expect.stringContaining("Preference: on"), "info"); + expect(notify).toHaveBeenLastCalledWith( + expect.stringContaining("Last requested tier: priority"), + "info", + ); + expect(notify).toHaveBeenLastCalledWith( + expect.stringContaining("Last observed tier: default"), + "info", + ); + } + expect(network.fetch).toHaveBeenCalledTimes(requests); + expect(resolveAuth).not.toHaveBeenCalled(); + expect(manager.getBranch()).toEqual(afterTurn); + expect(modelTurns).toHaveBeenCalledOnce(); + expect(credentialWrites).not.toHaveBeenCalled(); + expect(errors).toEqual([]); + } finally { + await closeSession(); + } + }, +); + +it("SDK replays the active fast branch alongside context controls, then off and reset change the wire", async () => { + const network = installResponseTransports("sse"); + const { runtime, model, credentialWrites, modelTurns } = await createOfflineModelRuntime( + "sse", + network.fetch, + ); + const manager = SessionManager.inMemory(directory); + const saved = { provider: "openai-codex", modelId: model.id, mode: "on" }; + const onLeaf = manager.appendCustomEntry(FAST_ENTRY, saved); + // Leave a conflicting choice on an inactive sibling: replay must use getBranch(), not all entries. + manager.appendCustomEntry(FAST_ENTRY, { ...saved, mode: "off" }); + manager.branch(onLeaf); + const { session, confirm, notify, errors, closeSession } = await openRealExtensionSession( + runtime, + model, + manager, + ); + try { + expect(network.fetch).not.toHaveBeenCalled(); + const branch = structuredClone(manager.getBranch()); + await session.prompt("/eleith:codex fast"); + await session.prompt("/eleith:codex fast status"); + expect(notify).toHaveBeenLastCalledWith(expect.stringContaining("Preference: on"), "info"); + expect(network.fetch).not.toHaveBeenCalled(); + expect(manager.getBranch()).toEqual(branch); + + const outputLimit = session.model?.maxTokens; + const thinking = session.thinkingLevel; + await session.prompt("/eleith:codex extend"); + const extendedWindow = session.model?.contextWindow; + expect(extendedWindow).toBeGreaterThanOrEqual(model.contextWindow); + expect(session.model?.maxTokens).toBe(outputLimit); + expect(session.thinkingLevel).toBe(thinking); + expect(network.fetch).not.toHaveBeenCalled(); + + // No explicit fast command fetched metadata in this resumed session. + await session.prompt("Reply with ok on the resumed branch."); + expect(network.catalog).toHaveBeenCalledOnce(); + expect(network.payloads[0]).toMatchObject({ model: model.id, service_tier: "priority" }); + expect(fastChoices(manager)).toEqual([saved]); + expect(confirm).not.toHaveBeenCalled(); + + await session.prompt("/eleith:codex fast off"); + expect(fastChoices(manager)).toEqual([saved, { ...saved, mode: "off" }]); + await session.prompt("Reply with ok without priority."); + expect(network.payloads[1]).toMatchObject({ model: model.id, service_tier: "default" }); + + await session.prompt("/eleith:codex fast reset"); + expect(fastChoices(manager)).toEqual([ + saved, + { ...saved, mode: "off" }, + { ...saved, mode: "inherit" }, + ]); + await session.prompt("Reply with ok using the provider default."); + expect(network.payloads[2]).toHaveProperty("model", model.id); + expect(network.payloads[2]).not.toHaveProperty("service_tier"); + expect(network.payloads).toHaveLength(3); + expect(session.getLastAssistantText()).toBe("ok"); + expect(modelTurns).toHaveBeenCalledTimes(3); + expect(network.catalog).toHaveBeenCalledOnce(); + expect(session.model?.contextWindow).toBe(extendedWindow); + expect(session.model?.maxTokens).toBe(outputLimit); + expect(session.thinkingLevel).toBe(thinking); + expect(credentialWrites).not.toHaveBeenCalled(); + expect(errors).toEqual([]); + } finally { + await closeSession(); + } +}); + +// Network policy: globals start closed; the supplied fetch accepts only these two fixed URLs. +function blockUnexpectedNetwork() { + vi.stubGlobal( + "fetch", + vi.fn(async () => { + throw new Error("Unexpected network request"); + }), + ); + vi.stubGlobal( + "WebSocket", + class { + constructor() { + throw new Error("Unexpected WebSocket request"); + } + }, + ); +} + +function installResponseTransports(transport: Transport) { + const payloads: Record[] = []; + const http = createHttpResponses(transport, payloads); + vi.stubGlobal("fetch", http.fetch); + const sockets = transport === "websocket" ? installSocketResponses(payloads) : []; + return { ...http, payloads, sockets }; +} + +function createHttpResponses(transport: Transport, payloads: Record[]) { + const catalog = vi.fn(async () => + // The client wraps this internally as { data, authScope }; the network JSON does not. + Response.json({ models: [{ slug: "gpt-6-sol", service_tiers: [{ id: "priority" }] }] }), + ); + const sse = vi.fn(async (_url, init) => { + if (transport !== "sse") throw new Error("Unexpected SSE fallback"); + payloads.push(readHttpPayload(init)); + const body = responseProtocolData() + .map((event) => `data: ${JSON.stringify(event)}\n\n`) + .join(""); + return new Response(body, { headers: { "Content-Type": "text/event-stream" } }); + }); + const fetcher = vi.fn(async (url, init) => { + if (String(url) === CATALOG_URL) return catalog(url, init); + if (String(url) === RESPONSES_URL) return sse(url, init); + throw new Error("Unexpected network URL"); + }); + return { fetch: fetcher, catalog, sse }; +} + +function installSocketResponses(payloads: Record[]) { + const sockets: FakeWebSocket[] = []; + class FakeWebSocket extends EventTarget { + readyState = 0; + + constructor(url: string | URL) { + super(); + if (String(url) !== SOCKET_URL) throw new Error("Unexpected WebSocket URL"); + sockets.push(this); + queueMicrotask(() => { + if (this.readyState !== 0) return; + this.readyState = 1; + this.dispatchEvent(new Event("open")); + }); + } + + send(data: string) { + if (this.readyState !== 1) throw new Error("WebSocket is not open"); + payloads.push(JSON.parse(data)); + queueMicrotask(() => { + for (const frame of responseProtocolData()) { + this.dispatchEvent(new MessageEvent("message", { data: JSON.stringify(frame) })); + } + }); + } + + close() { + if (this.readyState === 3) return; + this.readyState = 3; + this.dispatchEvent(new Event("close")); + } + } + vi.stubGlobal("WebSocket", FakeWebSocket); + return sockets; +} + +// Offline ModelRuntime: real built-in provider and OAuth resolution, never persistence or refresh I/O. +async function createOfflineModelRuntime(transport: Transport, fetcher: typeof fetch) { + const claim = { "https://api.openai.com/auth": { chatgpt_account_id: "fast-sdk-account" } }; + const accessToken = `e30.${Buffer.from(JSON.stringify(claim)).toString("base64url")}.fake-signature`; + const { credentials, credentialWrites } = createReadOnlyCredentials(accessToken); + const runtime = await ModelRuntime.create({ + credentials, + modelsPath: null, + modelsStore: new InMemoryModelsStore(), + allowModelNetwork: false, + refreshOnCreate: false, + }); + const refresh = await runtime.refresh({ providers: ["openai-codex"], allowNetwork: false }); + expect(refresh.errors.size).toBe(0); + expect(runtime.isUsingOAuth("openai-codex")).toBe(true); + const model = runtime.getModel("openai-codex", "gpt-6-sol"); + if (!model) throw new Error("Built-in Codex model missing"); + const resolveAuth = vi.spyOn(runtime, "getAuth"); + const nativeStream = runtime.streamSimple.bind(runtime); + // Only select the transport at this public seam; retain native serialization and extension hooks. + const modelTurns = vi + .spyOn(runtime, "streamSimple") + .mockImplementation((selected, context, options) => + nativeStream(selected, context, { + ...options, + transport, + cacheRetention: "none", // No socket cache or idle-expiry timer survives a turn. + fetch: fetcher, + }), + ); + return { runtime, model, accessToken, credentialWrites, resolveAuth, modelTurns }; +} + +function createReadOnlyCredentials(accessToken: string) { + const freshOAuth: OAuthCredential = { + type: "oauth", + access: accessToken, + refresh: "must-not-refresh", + expires: Date.now() + 3_600_000, + }; + const credentialWrites = vi.fn(async () => { + throw new Error("Credential writes forbidden"); + }); + const credentials: CredentialStore = { + read: async (id) => (id === "openai-codex" ? freshOAuth : undefined), + list: async () => [{ providerId: "openai-codex", type: "oauth" }], + modify: credentialWrites, + delete: async () => { + throw new Error("Credential deletion forbidden"); + }, + }; + return { credentials, credentialWrites }; +} + +// Real extension session: production commands, branch persistence, and the public agent bridge. +async function openRealExtensionSession( + runtime: ModelRuntime, + model: Model, + manager = SessionManager.inMemory(directory), +) { + const settingsManager = SettingsManager.inMemory({ + packages: [], + enableInstallTelemetry: false, + enableAnalytics: false, + compaction: { enabled: false }, + retry: { enabled: false }, + cacheWarming: "off", + }); + const resourceLoader = new DefaultResourceLoader({ + cwd: directory, + agentDir: directory, + settingsManager, + noExtensions: true, + noSkills: true, + noPromptTemplates: true, + noThemes: true, + noContextFiles: true, + extensionFactories: [codex], + }); + await resourceLoader.reload(); + expect(resourceLoader.getExtensions().errors).toEqual([]); + const { session } = await createAgentSession({ + cwd: directory, + agentDir: directory, + resourceLoader, + settingsManager, + modelRuntime: runtime, + model, + thinkingLevel: "off", + sessionManager: manager, + noTools: "all", + }); + const confirm = vi.fn(async () => true); + const notify = vi.fn(); + const errors: unknown[] = []; + await session.bindExtensions({ + mode: "tui", + uiContext: { notify, confirm } as unknown as ExtensionUIContext, + onError: (error) => errors.push(error), + }); + return { + session, + manager, + confirm, + notify, + errors, + async closeSession() { + await session.extensionRunner.emit({ type: "session_shutdown", reason: "quit" }); + session.dispose(); + }, + }; +} + +function fastChoices(manager: SessionManager) { + return manager + .getBranch() + .flatMap((entry) => + entry.type === "custom" && entry.customType === FAST_ENTRY ? [entry.data] : [], + ); +} + +// Response protocol data: deliberately report default, not the requested priority. +function responseProtocolData() { + const message = { + type: "message", + id: "answer", + role: "assistant", + status: "completed", + content: [{ type: "output_text", text: "ok", annotations: [] }], + }; + return [ + { type: "response.created", response: { id: "response" } }, + { type: "response.output_item.added", output_index: 0, item: { ...message, content: [] } }, + { type: "response.output_text.delta", output_index: 0, delta: "ok" }, + { type: "response.output_item.done", output_index: 0, item: message }, + { + type: "response.completed", + response: { id: "response", status: "completed", service_tier: "default", output: [message] }, + }, + ]; +} + +function readHttpPayload(init: RequestInit | undefined): Record { + if (typeof init?.body === "string") return JSON.parse(init.body); + if ( + init?.body instanceof Uint8Array && + new Headers(init.headers).get("content-encoding") === "zstd" + ) { + return JSON.parse(zstdDecompressSync(init.body).toString("utf8")); + } + throw new Error("Unexpected Codex request body"); +} diff --git a/extensions/codex/src/fast-state.test.ts b/extensions/codex/src/fast-state.test.ts new file mode 100644 index 0000000..70a3f2a --- /dev/null +++ b/extensions/codex/src/fast-state.test.ts @@ -0,0 +1,128 @@ +import type { SessionEntry } from "@earendil-works/pi-coding-agent"; +import { describe, expect, it } from "vitest"; +import { FAST_ENTRY, readFastPreferences, type FastChoice } from "./fast-state.ts"; +import { CONTEXT_ENTRY, readPreferences } from "./state.ts"; + +function preference( + data: unknown, + id = "preference", + parentId: string | null = null, + customType = FAST_ENTRY, +): Extract { + return { type: "custom", id, parentId, timestamp: "2026-01-01T00:00:00Z", customType, data }; +} + +const enabled: FastChoice = { provider: "openai-codex", modelId: "gpt-6-sol", mode: "on" }; + +describe("active branch fast preferences", () => { + it("uses a dedicated transcript identifier and returns an empty map for an empty branch", () => { + expect(FAST_ENTRY).toBe("eleith-codex-fast"); + expect(readFastPreferences([])).toEqual(new Map()); + }); + + it("keeps choices independent by provider and model", () => { + const branch: readonly SessionEntry[] = [ + preference(enabled), + preference({ ...enabled, modelId: "gpt-6-astra", mode: "off" }), + preference({ ...enabled, provider: "other", mode: "inherit" }), + ]; + expect(readFastPreferences(branch)).toEqual( + new Map([ + ["openai-codex/gpt-6-sol", "on"], + ["openai-codex/gpt-6-astra", "off"], + ["other/gpt-6-sol", "inherit"], + ]), + ); + }); + + it.each([ + undefined, + null, + "on", + 1, + true, + [], + Object.assign([], enabled), + {}, + { modelId: "gpt-6-sol", mode: "on" }, + { ...enabled, provider: 1 }, + { ...enabled, provider: "" }, + { provider: "openai-codex", mode: "on" }, + { ...enabled, modelId: null }, + { ...enabled, modelId: "" }, + { provider: "openai-codex", modelId: "gpt-6-sol" }, + { ...enabled, mode: true }, + { ...enabled, mode: "fast" }, + { ...enabled, mode: "ON" }, + { ...enabled, mode: null }, + { provider: "openai-codex", modelId: "gpt-6-sol", extended: true }, + ])("ignores malformed data: %j", (data) => { + expect(readFastPreferences([preference(data)])).toEqual(new Map()); + }); + + it("requires a custom entry with the fast customType", () => { + expect( + readFastPreferences([ + preference(enabled, "unrelated", null, "another-extension"), + { type: "custom_message", customType: FAST_ENTRY, data: enabled }, + { type: "custom", data: enabled }, + ]), + ).toEqual(new Map()); + }); + + it("uses the latest valid entry in branch order, not timestamp order", () => { + const first = preference(enabled, "first"); + const second = preference({ ...enabled, mode: "off" }, "second", first.id); + second.timestamp = "2025-01-01T00:00:00Z"; + const malformed = preference({ ...enabled, mode: "invalid" }, "bad", second.id); + expect(readFastPreferences([first, second, malformed])).toEqual( + new Map([["openai-codex/gpt-6-sol", "off"]]), + ); + expect(readFastPreferences([first, second, preference(enabled, "third", second.id)])).toEqual( + new Map([["openai-codex/gpt-6-sol", "on"]]), + ); + }); + + it("preserves inherit as a tombstone until a later valid choice replaces it", () => { + const branch = [preference(enabled), preference({ ...enabled, mode: "inherit" })]; + const preferences = readFastPreferences(branch); + expect(preferences.has("openai-codex/gpt-6-sol")).toBe(true); + expect(preferences.get("openai-codex/gpt-6-sol")).toBe("inherit"); + expect(readFastPreferences([...branch, preference({ ...enabled, mode: "off" })])).toEqual( + new Map([["openai-codex/gpt-6-sol", "off"]]), + ); + }); + + it("reads only the supplied branch without retaining abandoned choices", () => { + const shared = preference(enabled, "shared"); + const abandoned: SessionEntry[] = [ + shared, + preference({ ...enabled, mode: "inherit" }, "abandoned", shared.id), + preference({ ...enabled, modelId: "gpt-6-astra" }, "abandoned-model", "abandoned"), + ]; + const active: readonly SessionEntry[] = [ + shared, + { type: "session_info", id: "active", parentId: shared.id, timestamp: shared.timestamp }, + ]; + expect(readFastPreferences(abandoned)).toEqual( + new Map([ + ["openai-codex/gpt-6-sol", "inherit"], + ["openai-codex/gpt-6-astra", "on"], + ]), + ); + expect(readFastPreferences(active)).toEqual(new Map([["openai-codex/gpt-6-sol", "on"]])); + expect(readFastPreferences([])).toEqual(new Map()); + }); + + it("keeps fast and context schemas independent even when payload fields overlap", () => { + const context = { provider: enabled.provider, modelId: enabled.modelId, extended: false }; + const branch = [ + preference({ ...enabled, extended: true }), + preference({ ...context, mode: "off" }, "context", null, CONTEXT_ENTRY), + preference(context, "wrong-fast-schema"), + preference(enabled, "wrong-context-schema", null, CONTEXT_ENTRY), + ]; + expect(readFastPreferences(branch)).toEqual(new Map([["openai-codex/gpt-6-sol", "on"]])); + expect(readPreferences(branch)).toEqual(new Map([["openai-codex/gpt-6-sol", false]])); + }); +}); diff --git a/extensions/codex/src/fast-state.ts b/extensions/codex/src/fast-state.ts new file mode 100644 index 0000000..f071608 --- /dev/null +++ b/extensions/codex/src/fast-state.ts @@ -0,0 +1,36 @@ +export const FAST_ENTRY = "eleith-codex-fast"; + +export type FastPreference = "on" | "off" | "inherit"; + +export interface FastChoice { + readonly provider: string; + readonly modelId: string; + readonly mode: FastPreference; +} + +interface Entry { + readonly type: string; + readonly customType?: string; + readonly data?: unknown; +} + +// Replay only the supplied active branch, retaining explicit inherit tombstones. +export function readFastPreferences(branch: readonly Entry[]): Map { + const preferences = new Map(); + for (const entry of branch) { + if (entry.type !== "custom" || entry.customType !== FAST_ENTRY) continue; + if (!isFastChoice(entry.data)) continue; + preferences.set(`${entry.data.provider}/${entry.data.modelId}`, entry.data.mode); + } + return preferences; +} + +function isFastChoice(value: unknown): value is FastChoice { + if (typeof value !== "object" || value === null || Array.isArray(value)) return false; + if (!("provider" in value) || typeof value.provider !== "string" || !value.provider.length) + return false; + if (!("modelId" in value) || typeof value.modelId !== "string" || !value.modelId.length) + return false; + if (!("mode" in value)) return false; + return value.mode === "on" || value.mode === "off" || value.mode === "inherit"; +} diff --git a/extensions/codex/src/fast-status.test.ts b/extensions/codex/src/fast-status.test.ts new file mode 100644 index 0000000..a50f9f1 --- /dev/null +++ b/extensions/codex/src/fast-status.test.ts @@ -0,0 +1,42 @@ +import { expect, it } from "vitest"; +import { formatFastStatus, type FastStatus } from "./fast-status.ts"; + +const suspended: FastStatus = { + preference: "on", + eligibility: "unknown", + requestPolicy: "suspended", + lastRequested: undefined, + lastObserved: undefined, + conflict: false, +}; + +it("labels unknown delivery and explains selected scope, allowance consumption and cost-estimate limits", () => { + const text = formatFastStatus(suspended); + expect(text.split("\n")).toEqual([ + "Codex fast mode (selected model)", + "Preference: on", + "Eligibility: unknown", + "Next request policy: suspended", + "Last requested tier: unknown", + "Last observed tier: unknown", + "Scope: direct legacy Codex; may include retries, compaction and cache warming. Selected scope is not physical-request identity.", + "Priority uses more allowance. Pi's cost estimate may miss priority pricing when response tier metadata is absent or normalized.", + "Observed tier is model-scoped, not a receipt for the last requested tier.", + ]); +}); + +it("keeps requested and observed tiers distinct and adds the conflict note only when needed", () => { + const status: FastStatus = { + ...suspended, + eligibility: "supported", + requestPolicy: "priority", + lastRequested: "priority", + lastObserved: "default", + }; + const text = formatFastStatus(status); + expect(text).toContain("Last requested tier: priority\nLast observed tier: default"); + expect(text).not.toContain("Tier conflict:"); + expect(formatFastStatus({ ...status, conflict: true })).toContain( + "Tier conflict: an existing non-owned tier was left untouched.", + ); +}); diff --git a/extensions/codex/src/fast-status.ts b/extensions/codex/src/fast-status.ts new file mode 100644 index 0000000..f893d43 --- /dev/null +++ b/extensions/codex/src/fast-status.ts @@ -0,0 +1,33 @@ +import type { FastEligibility } from "./capabilities.ts"; +import type { FastPreference } from "./fast-state.ts"; + +export type RequestedTier = "priority" | "default"; +export type ObservedTier = RequestedTier | "auto" | "flex" | "fast"; +export type FastRequestPolicy = RequestedTier | "untouched" | "suspended" | "outside scope"; + +export interface FastStatus { + preference: FastPreference; + eligibility: FastEligibility; + requestPolicy: FastRequestPolicy; + lastRequested: RequestedTier | undefined; + lastObserved: ObservedTier | undefined; + conflict: boolean; +} + +export function formatFastStatus(status: FastStatus): string { + const rows = [ + "Codex fast mode (selected model)", + `Preference: ${status.preference}`, + `Eligibility: ${status.eligibility}`, + `Next request policy: ${status.requestPolicy}`, + `Last requested tier: ${status.lastRequested ?? "unknown"}`, + `Last observed tier: ${status.lastObserved ?? "unknown"}`, + ]; + if (status.conflict) rows.push("Tier conflict: an existing non-owned tier was left untouched."); + rows.push( + "Scope: direct legacy Codex; may include retries, compaction and cache warming. Selected scope is not physical-request identity.", + "Priority uses more allowance. Pi's cost estimate may miss priority pricing when response tier metadata is absent or normalized.", + "Observed tier is model-scoped, not a receipt for the last requested tier.", + ); + return rows.join("\n"); +} diff --git a/extensions/codex/src/fast.test.ts b/extensions/codex/src/fast.test.ts new file mode 100644 index 0000000..f23a387 --- /dev/null +++ b/extensions/codex/src/fast.test.ts @@ -0,0 +1,818 @@ +import type { Api, Model } from "@earendil-works/pi-ai"; +import type { + ExtensionAPI, + ProviderStreamEvent, + SessionEntry, +} from "@earendil-works/pi-coding-agent"; +import { describe, expect, it, vi } from "vitest"; +import { + CodexAccountError, + type CodexModelCatalog, + type fetchCodexModelCatalog, + type readCodexAuthScope, +} from "./client.ts"; +import { FastController, type FastHost } from "./fast.ts"; +import { FAST_ENTRY, type FastPreference } from "./fast-state.ts"; +import { CONTEXT_ENTRY } from "./state.ts"; + +const model: Model = { + id: "gpt-6-sol", + name: "Sol", + provider: "openai-codex", + api: "openai-codex-responses", + baseUrl: "https://chatgpt.com/backend-api", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 272_000, + maxTokens: 128_000, +}; +const otherModel: Model = { ...model, id: "gpt-6-astra", name: "Astra" }; +const catalog: CodexModelCatalog = { + authScope: "scope-a", + data: { models: [{ slug: model.id, service_tiers: [{ id: "default" }, { id: "priority" }] }] }, +}; + +describe("commands and ownership", () => { + it.each(["", " STATUS "])( + "reports %j without confirmation, auth IO, catalog IO or writes", + async (args) => { + const fast = createFastControllerFixture(); + fast.ctx.sessionManager.getBranch.mockReturnValue([choice("on")]); + fast.controller.restore(fast.ctx); + await fast.controller.handle(args, fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenCalledExactlyOnceWith( + expect.stringContaining( + "Preference: on\nEligibility: unknown\nNext request policy: suspended", + ), + "info", + ); + expect(fast.ctx.ui.confirm).not.toHaveBeenCalled(); + expect(fast.readAuthScope).not.toHaveBeenCalled(); + expect(fast.ctx.modelRegistry.getProviderAuth).not.toHaveBeenCalled(); + expect(fast.fetchCatalog).not.toHaveBeenCalled(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + }, + ); + + it("does not discover or persist anything when the cost confirmation is declined", async () => { + const fast = createFastControllerFixture(); + fast.ctx.ui.confirm.mockResolvedValue(false); + await fast.controller.handle("on", fast.ctx); + expect(fast.ctx.ui.confirm).toHaveBeenCalledExactlyOnceWith( + "Enable Codex fast mode?", + expect.stringContaining("use more of your Codex allowance"), + { signal: expect.any(AbortSignal) }, + ); + expect(fast.readAuthScope).not.toHaveBeenCalled(); + expect(fast.fetchCatalog).not.toHaveBeenCalled(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toBeUndefined(); + }); + + it("requires UI to enable, but permits a noninteractive explicit off", async () => { + const fast = createFastControllerFixture(); + fast.ctx.hasUI = false; + await fast.controller.handle("on", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("interactive confirmation"), + "warning", + ); + expect(fast.ctx.ui.confirm).not.toHaveBeenCalled(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + await fast.controller.handle("off", fast.ctx); + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toEqual({ + model: model.id, + service_tier: "default", + }); + expect(fast.readAuthScope).not.toHaveBeenCalled(); + expect(fast.fetchCatalog).not.toHaveBeenCalled(); + }); + + it("rejects invalid arguments and mutations during a run without prompting or IO", async () => { + const fast = createFastControllerFixture(); + await fast.controller.handle("on extra", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + "Usage: /eleith:codex fast [on|off|toggle|status|reset]", + "warning", + ); + fast.ctx.isIdle.mockReturnValue(false); + for (const action of ["on", "off", "toggle", "reset"]) + await fast.controller.handle(action, fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Wait for the current run"), + "warning", + ); + expect(fast.ctx.ui.confirm).not.toHaveBeenCalled(); + expect(fast.readAuthScope).not.toHaveBeenCalled(); + expect(fast.fetchCatalog).not.toHaveBeenCalled(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + }); + + it("reports a missing selection rather than confirming or persisting a choice", async () => { + const fast = createFastControllerFixture(); + fast.ctx.model = undefined; + await fast.controller.handle("on", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith("No model selected.", "warning"); + expect(fast.ctx.ui.confirm).not.toHaveBeenCalled(); + expect(fast.readAuthScope).not.toHaveBeenCalled(); + expect(fast.fetchCatalog).not.toHaveBeenCalled(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + }); + + it("persists on only after cost confirmation and matching catalog discovery complete", async () => { + const fast = createFastControllerFixture(); + const confirmation = deferred(); + const response = deferred(); + const started = deferred(); + fast.ctx.ui.confirm.mockReturnValue(confirmation.promise); + fast.fetchCatalog.mockImplementationOnce(() => { + started.resolve(); + return response.promise; + }); + const enabling = fast.controller.handle(" ON ", fast.ctx); + expect(fast.readAuthScope).not.toHaveBeenCalled(); + expect(fast.fetchCatalog).not.toHaveBeenCalled(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + // A second mutating command must not compete with the confirmation. + await fast.controller.handle("off", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Wait for the current run"), + "warning", + ); + confirmation.resolve(true); + await started.promise; + expect(fast.readAuthScope).toHaveBeenCalledExactlyOnceWith( + fast.ctx.modelRegistry, + expect.any(AbortSignal), + ); + expect(fast.fetchCatalog).toHaveBeenCalledExactlyOnceWith( + fast.ctx.modelRegistry, + expect.any(AbortSignal), + ); + expect(fast.appendEntry).not.toHaveBeenCalled(); + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toBeUndefined(); + response.resolve(catalog); + await enabling; + expect(fast.appendEntry).toHaveBeenCalledExactlyOnceWith(FAST_ENTRY, { + provider: model.provider, + modelId: model.id, + mode: "on", + }); + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toEqual({ + model: model.id, + service_tier: "priority", + }); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Next request policy: priority"), + "info", + ); + }); + + it.each([ + { models: [{ slug: otherModel.id, service_tiers: [{ id: "priority" }] }] }, + { models: [{ slug: model.id, service_tiers: [{ id: "default" }] }] }, + ])( + "does not enable without priority support for the selected catalog model: %j", + async (data) => { + const fast = createFastControllerFixture(); + fast.fetchCatalog.mockResolvedValue({ ...catalog, data }); + await fast.controller.handle("on", fast.ctx); + expect(fast.appendEntry).not.toHaveBeenCalled(); + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toBeUndefined(); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Preference unchanged"), + "warning", + ); + }, + ); + + it("copies a frozen payload for explicit default, preserves fields, then resets with an inherit tombstone", async () => { + const fast = createFastControllerFixture(); + const input = Object.freeze([{ role: "user", content: "hello" }]); + const payload = Object.freeze({ + model: model.id, + input, + stream: true, + metadata: { trace: "keep" }, + max_output_tokens: 42, + }); + await fast.controller.handle("off", fast.ctx); + const transformed = fast.controller.transform(payload, fast.ctx); + expect(transformed).toEqual({ ...payload, service_tier: "default" }); + expect(transformed).not.toBe(payload); + expect(transformed?.input).toBe(input); + expect(transformed?.metadata).toBe(payload.metadata); + expect(payload).not.toHaveProperty("service_tier"); + expect( + fast.controller.transform(Object.freeze({ ...payload, service_tier: "default" }), fast.ctx), + ).toBeUndefined(); + await fast.controller.handle("off", fast.ctx); + expect(fast.appendEntry).toHaveBeenCalledOnce(); + await fast.controller.handle("reset", fast.ctx); + expect(fast.appendEntry).toHaveBeenLastCalledWith(FAST_ENTRY, { + provider: model.provider, + modelId: model.id, + mode: "inherit", + }); + expect(fast.controller.transform(payload, fast.ctx)).toBeUndefined(); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining( + "Preference: inherit\nEligibility: unknown\nNext request policy: untouched", + ), + "info", + ); + await fast.controller.handle("reset", fast.ctx); + expect(fast.appendEntry).toHaveBeenCalledTimes(2); + expect(fast.ctx.ui.confirm).not.toHaveBeenCalled(); + expect(fast.readAuthScope).not.toHaveBeenCalled(); + expect(fast.fetchCatalog).not.toHaveBeenCalled(); + }); + + it("can remove its saved choice after the selection leaves direct scope", async () => { + const fast = createFastControllerFixture(); + await fast.controller.handle("off", fast.ctx); + fast.ctx.modelRegistry.isUsingOAuth.mockReturnValue(false); + await fast.controller.handle("reset", fast.ctx); + expect(fast.appendEntry).toHaveBeenLastCalledWith(FAST_ENTRY, { + provider: model.provider, + modelId: model.id, + mode: "inherit", + }); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Next request policy: outside scope"), + "info", + ); + }); + + it.each(["on", "off"])( + "preserves an existing non-owned flex tier when %s is selected", + async (mode) => { + const fast = createFastControllerFixture(); + const tier = mode === "on" ? "priority" : "default"; + await fast.controller.handle(mode, fast.ctx); + const payload = Object.freeze({ model: model.id, service_tier: "flex", input: [] }); + expect(fast.controller.transform(payload, fast.ctx)).toBeUndefined(); + expect(payload.service_tier).toBe("flex"); + await fast.controller.handle("status", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Tier conflict: an existing non-owned tier was left untouched."), + "info", + ); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Last requested tier: unknown"), + "info", + ); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining(`Next request policy: ${tier}`), + "info", + ); + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toEqual({ + model: model.id, + service_tier: tier, + }); + await fast.controller.handle("status", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.not.stringContaining("Tier conflict:"), + "info", + ); + }, + ); + + it("toggle confirms only transitions to on and leaves off intact if the next confirmation is declined", async () => { + const fast = createFastControllerFixture(); + await fast.controller.handle("toggle", fast.ctx); + expect(fast.ctx.ui.confirm).toHaveBeenCalledOnce(); + expect(fast.controller.transform({ model: model.id }, fast.ctx)?.service_tier).toBe("priority"); + await fast.controller.handle("toggle", fast.ctx); + expect(fast.ctx.ui.confirm).toHaveBeenCalledOnce(); + expect(fast.controller.transform({ model: model.id }, fast.ctx)?.service_tier).toBe("default"); + fast.ctx.ui.confirm.mockResolvedValue(false); + await fast.controller.handle("toggle", fast.ctx); + expect(fast.ctx.ui.confirm).toHaveBeenCalledTimes(2); + expect(fast.appendEntry.mock.calls.map((call) => call[1])).toEqual([ + { provider: model.provider, modelId: model.id, mode: "on" }, + { provider: model.provider, modelId: model.id, mode: "off" }, + ]); + expect(fast.controller.transform({ model: model.id }, fast.ctx)?.service_tier).toBe("default"); + expect(fast.fetchCatalog).toHaveBeenCalledOnce(); + }); +}); + +describe("branch replay and preparation", () => { + it("replays only the current branch, separates model choices, and ignores context choices", async () => { + const fast = createFastControllerFixture(); + fast.ctx.sessionManager.getBranch.mockReturnValue([ + choice("off"), + choice("on"), + choice("off", otherModel), + { + ...choice("off"), + customType: CONTEXT_ENTRY, + data: { provider: model.provider, modelId: model.id, extended: true, mode: "off" }, + }, + ]); + fast.controller.restore(fast.ctx); + expect(fast.readAuthScope).not.toHaveBeenCalled(); + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toBeUndefined(); + await fast.controller.prepare(fast.ctx); + expect(fast.controller.transform({ model: model.id }, fast.ctx)?.service_tier).toBe("priority"); + fast.ctx.model = otherModel; + expect(fast.controller.transform({ model: otherModel.id }, fast.ctx)?.service_tier).toBe( + "default", + ); + fast.ctx.sessionManager.getBranch.mockReturnValue([choice("on"), choice("inherit")]); + fast.controller.restore(fast.ctx); + expect(fast.controller.transform({ model: otherModel.id }, fast.ctx)).toBeUndefined(); + fast.ctx.model = model; + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toBeUndefined(); + await fast.controller.prepare(fast.ctx); + expect(fast.readAuthScope).toHaveBeenCalledOnce(); + expect(fast.fetchCatalog).toHaveBeenCalledOnce(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + expect(fast.ctx.ui.confirm).not.toHaveBeenCalled(); + }); + + it("lets fresh on expire at the TTL without hook IO or erasing the suspended preference", async () => { + const fast = createFastControllerFixture(); + await fast.controller.handle("on", fast.ctx); + fast.clock.now += 5 * 60_000 - 1; + expect(fast.controller.transform({ model: model.id }, fast.ctx)?.service_tier).toBe("priority"); + fast.clock.now += 1; + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toBeUndefined(); + fast.controller.observe( + stream({ type: "response.completed", response: { service_tier: "default" } }), + fast.ctx, + ); + await fast.controller.handle("status", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining( + "Preference: on\nEligibility: unknown\nNext request policy: suspended", + ), + "info", + ); + expect(fast.readAuthScope).toHaveBeenCalledOnce(); + expect(fast.fetchCatalog).toHaveBeenCalledOnce(); + expect(fast.appendEntry).toHaveBeenCalledOnce(); + await fast.controller.prepare(fast.ctx); + expect(fast.fetchCatalog).toHaveBeenCalledTimes(2); + expect(fast.controller.transform({ model: model.id }, fast.ctx)?.service_tier).toBe("priority"); + expect(fast.appendEntry).toHaveBeenCalledOnce(); + expect(fast.ctx.ui.confirm).toHaveBeenCalledOnce(); + }); + + it("prepares only idle saved-on input, checking auth even when the catalog cache is fresh", async () => { + const fast = createFastControllerFixture(); + for (const branch of [[], [choice("off")], [choice("inherit")]]) { + fast.ctx.sessionManager.getBranch.mockReturnValue(branch); + fast.controller.restore(fast.ctx); + await fast.controller.prepare(fast.ctx); + } + expect(fast.readAuthScope).not.toHaveBeenCalled(); + fast.ctx.sessionManager.getBranch.mockReturnValue([choice("on")]); + fast.controller.restore(fast.ctx); + fast.ctx.isIdle.mockReturnValue(false); + await fast.controller.prepare(fast.ctx); + expect(fast.readAuthScope).not.toHaveBeenCalled(); + fast.ctx.isIdle.mockReturnValue(true); + await fast.controller.prepare(fast.ctx); + await fast.controller.prepare(fast.ctx); + expect(fast.readAuthScope).toHaveBeenCalledTimes(2); + expect(fast.fetchCatalog).toHaveBeenCalledOnce(); + expect(fast.controller.transform({ model: model.id }, fast.ctx)?.service_tier).toBe("priority"); + expect(fast.ctx.ui.confirm).not.toHaveBeenCalled(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + }); + + it("invalidates a fresh catalog on auth scope change and recatalogs before resuming", async () => { + const fast = createFastControllerFixture(); + fast.ctx.sessionManager.getBranch.mockReturnValue([choice("on")]); + fast.controller.restore(fast.ctx); + await fast.controller.prepare(fast.ctx); + fast.readAuthScope.mockResolvedValue("scope-b"); + fast.fetchCatalog.mockResolvedValue({ + authScope: "scope-b", + data: { models: [{ slug: model.id, service_tiers: [{ id: "default" }] }] }, + }); + await fast.controller.prepare(fast.ctx); + expect(fast.fetchCatalog).toHaveBeenCalledTimes(2); + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toBeUndefined(); + await fast.controller.handle("status", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining( + "Preference: on\nEligibility: unsupported\nNext request policy: suspended", + ), + "info", + ); + fast.readAuthScope.mockResolvedValue("scope-c"); + fast.fetchCatalog.mockResolvedValue({ ...catalog, authScope: "scope-c" }); + await fast.controller.prepare(fast.ctx); + expect(fast.fetchCatalog).toHaveBeenCalledTimes(3); + expect(fast.controller.transform({ model: model.id }, fast.ctx)?.service_tier).toBe("priority"); + expect(fast.appendEntry).not.toHaveBeenCalled(); + }); + + it("rejects a catalog fetched under a different auth scope, preserving the preference for a retry", async () => { + const fast = createFastControllerFixture(); + fast.ctx.sessionManager.getBranch.mockReturnValue([choice("on")]); + fast.controller.restore(fast.ctx); + fast.fetchCatalog.mockResolvedValueOnce({ ...catalog, authScope: "scope-b" }); + await fast.controller.prepare(fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("authentication changed during discovery"), + "warning", + ); + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toBeUndefined(); + await fast.controller.prepare(fast.ctx); + expect(fast.fetchCatalog).toHaveBeenCalledTimes(2); + expect(fast.controller.transform({ model: model.id }, fast.ctx)?.service_tier).toBe("priority"); + expect(fast.appendEntry).not.toHaveBeenCalled(); + }); + + it("honors 429 backoff without timers or hook retries and resumes only on later input", async () => { + const fast = createFastControllerFixture(); + fast.ctx.sessionManager.getBranch.mockReturnValue([choice("on")]); + fast.controller.restore(fast.ctx); + fast.fetchCatalog.mockRejectedValueOnce( + new CodexAccountError("http", "Codex account service is rate limited.", 60_000), + ); + await fast.controller.prepare(fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("rate limited. Preference unchanged."), + "warning", + ); + await fast.controller.prepare(fast.ctx); + expect(fast.fetchCatalog).toHaveBeenCalledOnce(); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Try again later"), + "warning", + ); + fast.clock.now += 60_000; + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toBeUndefined(); + fast.controller.observe( + stream({ type: "response.completed", response: { service_tier: "priority" } }), + fast.ctx, + ); + await fast.controller.handle("status", fast.ctx); + expect(fast.fetchCatalog).toHaveBeenCalledOnce(); + expect(fast.readAuthScope).toHaveBeenCalledTimes(2); + await fast.controller.prepare(fast.ctx); + expect(fast.fetchCatalog).toHaveBeenCalledTimes(2); + expect(fast.controller.transform({ model: model.id }, fast.ctx)?.service_tier).toBe("priority"); + expect(fast.appendEntry).not.toHaveBeenCalled(); + }); + + it.each(["auth", "catalog"])( + "does not display unexpected %s errors or clear a saved-on preference", + async (source) => { + const fast = createFastControllerFixture(); + fast.ctx.sessionManager.getBranch.mockReturnValue([choice("on")]); + fast.controller.restore(fast.ctx); + await fast.controller.prepare(fast.ctx); + const error = new Error("private failure details must not appear"); + if (source === "auth") fast.readAuthScope.mockRejectedValueOnce(error); + else { + fast.clock.now += 5 * 60_000; + fast.fetchCatalog.mockRejectedValueOnce(error); + } + await fast.controller.prepare(fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + "Fast mode suspended: Could not prepare Codex fast mode. Preference unchanged.", + "warning", + ); + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toBeUndefined(); + await fast.controller.handle("status", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Preference: on"), + "info", + ); + expect(fast.appendEntry).not.toHaveBeenCalled(); + }, + ); +}); + +describe("asynchronous ownership", () => { + it.each(["model", "session", "leaf"])( + "discards confirmation after the %s changes", + async (changed) => { + const fast = createFastControllerFixture(); + const confirmation = deferred(); + fast.ctx.ui.confirm.mockReturnValue(confirmation.promise); + const enabling = fast.controller.handle("on", fast.ctx); + if (changed === "model") fast.ctx.model = otherModel; + if (changed === "session") fast.ctx.sessionManager.getSessionId.mockReturnValue("session-b"); + if (changed === "leaf") fast.ctx.sessionManager.getLeafId.mockReturnValue("leaf-b"); + confirmation.resolve(true); + await enabling; + expect(fast.readAuthScope).not.toHaveBeenCalled(); + expect(fast.fetchCatalog).not.toHaveBeenCalled(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + expect(fast.ctx.ui.notify).not.toHaveBeenCalled(); + }, + ); + + it.each(["model", "session", "leaf"])( + "discards catalog discovery after the %s changes", + async (changed) => { + const fast = createFastControllerFixture(); + const response = deferred(); + const started = deferred(); + fast.fetchCatalog.mockImplementationOnce(() => { + started.resolve(); + return response.promise; + }); + const enabling = fast.controller.handle("on", fast.ctx); + await started.promise; + if (changed === "model") fast.ctx.model = otherModel; + if (changed === "session") fast.ctx.sessionManager.getSessionId.mockReturnValue("session-b"); + if (changed === "leaf") fast.ctx.sessionManager.getLeafId.mockReturnValue("leaf-b"); + response.resolve(catalog); + await enabling; + expect(fast.appendEntry).not.toHaveBeenCalled(); + expect(fast.ctx.ui.notify).not.toHaveBeenCalled(); + expect(fast.controller.transform({ model: fast.ctx.model?.id }, fast.ctx)).toBeUndefined(); + }, + ); + + it("aborts invalidated discovery without letting a late result clear a newer operation", async () => { + const fast = createFastControllerFixture(); + const oldCatalog = deferred(); + const oldStarted = deferred(); + fast.fetchCatalog.mockImplementationOnce(() => { + oldStarted.resolve(); + return oldCatalog.promise; + }); + const oldEnable = fast.controller.handle("on", fast.ctx); + await oldStarted.promise; + const oldSignal = fast.fetchCatalog.mock.calls[0][1]!; + fast.controller.invalidate(); + expect(oldSignal.aborted).toBe(true); + const confirmation = deferred(); + fast.ctx.ui.confirm.mockReturnValueOnce(confirmation.promise); + const newEnable = fast.controller.handle("on", fast.ctx); + oldCatalog.resolve(catalog); + await oldEnable; + expect(fast.appendEntry).not.toHaveBeenCalled(); + await fast.controller.handle("off", fast.ctx); + expect(fast.appendEntry).not.toHaveBeenCalled(); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Wait for the current run"), + "warning", + ); + confirmation.resolve(true); + await newEnable; + expect(fast.appendEntry).toHaveBeenCalledExactlyOnceWith(FAST_ENTRY, { + provider: model.provider, + modelId: model.id, + mode: "on", + }); + }); + + it.each(["resolve", "reject"])( + "shutdown ignores late confirmation %s before touching expired host getters", + async (outcome) => { + const fast = createFastControllerFixture(); + const confirmation = deferred(); + fast.ctx.ui.confirm.mockReturnValue(confirmation.promise); + const enabling = fast.controller.handle("on", fast.ctx); + const signal = fast.ctx.ui.confirm.mock.calls[0][2]!.signal!; + const notify = fast.ctx.ui.notify; + fast.controller.dispose(); + fast.controller.dispose(); + for (const key of ["model", "modelRegistry", "isIdle", "sessionManager", "ui", "hasUI"]) { + Object.defineProperty(fast.ctx, key, { + get: () => { + throw new Error(`Expired ${key}`); + }, + }); + } + if (outcome === "resolve") confirmation.resolve(true); + else confirmation.reject(new Error("private late failure")); + await enabling; + expect(signal.aborted).toBe(true); + expect(notify).not.toHaveBeenCalled(); + expect(fast.readAuthScope).not.toHaveBeenCalled(); + expect(fast.fetchCatalog).not.toHaveBeenCalled(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + await fast.controller.handle("status", fast.ctx); + await fast.controller.prepare(fast.ctx); + fast.controller.restore(fast.ctx); + expect(fast.controller.transform({ model: model.id }, fast.ctx)).toBeUndefined(); + fast.controller.observe(stream({ type: "response.completed" }), fast.ctx); + }, + ); + + it.each(["resolve", "reject"])( + "shutdown ignores late catalog %s without accessing the expired host", + async (outcome) => { + const fast = createFastControllerFixture(); + const response = deferred(); + const started = deferred(); + fast.fetchCatalog.mockImplementationOnce(() => { + started.resolve(); + return response.promise; + }); + const enabling = fast.controller.handle("on", fast.ctx); + await started.promise; + const signal = fast.fetchCatalog.mock.calls[0][1]!; + const notify = fast.ctx.ui.notify; + fast.controller.dispose(); + for (const key of ["model", "modelRegistry", "isIdle", "sessionManager", "ui"]) { + Object.defineProperty(fast.ctx, key, { + get: () => { + throw new Error(`Expired ${key}`); + }, + }); + } + if (outcome === "resolve") response.resolve(catalog); + else + response.reject( + new CodexAccountError("http", "Codex account service is rate limited.", 60_000), + ); + await enabling; + expect(signal.aborted).toBe(true); + expect(notify).not.toHaveBeenCalled(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + }, + ); + + it("shutdown stops a saved-on auth check before any late catalog IO or host access", async () => { + const fast = createFastControllerFixture(); + fast.ctx.sessionManager.getBranch.mockReturnValue([choice("on")]); + fast.controller.restore(fast.ctx); + const scope = deferred(); + fast.readAuthScope.mockReturnValueOnce(scope.promise); + const preparing = fast.controller.prepare(fast.ctx); + const signal = fast.readAuthScope.mock.calls[0][1]!; + const notify = fast.ctx.ui.notify; + fast.controller.dispose(); + for (const key of ["model", "modelRegistry", "isIdle", "sessionManager", "ui"]) { + Object.defineProperty(fast.ctx, key, { + get: () => { + throw new Error(`Expired ${key}`); + }, + }); + } + scope.resolve("scope-a"); + await preparing; + expect(signal.aborted).toBe(true); + expect(fast.fetchCatalog).not.toHaveBeenCalled(); + expect(notify).not.toHaveBeenCalled(); + expect(fast.appendEntry).not.toHaveBeenCalled(); + }); +}); + +describe("provider stream observations", () => { + it("reports an observed default independently of requested priority, without changing the next preference", async () => { + const fast = createFastControllerFixture(); + await fast.controller.handle("on", fast.ctx); + const payload = Object.freeze({ model: model.id, input: [], service_tier: "priority" }); + expect(fast.controller.transform(payload, fast.ctx)).toBeUndefined(); + fast.controller.observe( + stream({ type: "response.completed", response: { service_tier: "default" } }), + fast.ctx, + ); + await fast.controller.handle("status", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining( + "Next request policy: priority\nLast requested tier: priority\nLast observed tier: default", + ), + "info", + ); + expect(fast.controller.transform({ model: model.id }, fast.ctx)?.service_tier).toBe("priority"); + expect(fast.fetchCatalog).toHaveBeenCalledOnce(); + expect(fast.readAuthScope).toHaveBeenCalledOnce(); + }); + + it.each([ + { type: "response.done", response: {} }, + { type: "response.incomplete", response: { service_tier: "unrecognized-tier" } }, + ])("reports missing or unrecognized terminal tier metadata as unknown: %j", async (data) => { + const fast = createFastControllerFixture(); + await fast.controller.handle("off", fast.ctx); + fast.controller.transform({ model: model.id }, fast.ctx); + fast.controller.observe( + stream({ type: "response.completed", response: { service_tier: "default" } }), + fast.ctx, + ); + fast.controller.observe(stream(data), fast.ctx); + await fast.controller.handle("status", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Last requested tier: default\nLast observed tier: unknown"), + "info", + ); + expect(fast.readAuthScope).not.toHaveBeenCalled(); + expect(fast.fetchCatalog).not.toHaveBeenCalled(); + }); + + it("ignores wrong-model and nonterminal events, and does not show delivery from another session", async () => { + const fast = createFastControllerFixture(); + await fast.controller.handle("off", fast.ctx); + fast.controller.transform({ model: model.id }, fast.ctx); + fast.controller.observe( + stream({ type: "response.completed", response: { service_tier: "default" } }), + fast.ctx, + ); + fast.controller.observe( + { + ...stream({ type: "response.completed", response: { service_tier: "priority" } }), + model: otherModel.id, + }, + fast.ctx, + ); + fast.controller.observe( + stream({ type: "response.created", response: { service_tier: "priority" } }), + fast.ctx, + ); + await fast.controller.handle("status", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Last observed tier: default"), + "info", + ); + fast.ctx.sessionManager.getSessionId.mockReturnValue("session-b"); + fast.controller.observe( + stream({ type: "response.completed", response: { service_tier: "priority" } }), + fast.ctx, + ); + await fast.controller.handle("status", fast.ctx); + expect(fast.ctx.ui.notify).toHaveBeenLastCalledWith( + expect.stringContaining("Last requested tier: unknown\nLast observed tier: unknown"), + "info", + ); + }); +}); + +function choice(mode: FastPreference, selected = model): Extract { + return { + type: "custom", + id: `${selected.id}-${mode}`, + parentId: null, + timestamp: "2026-01-01T00:00:00Z", + customType: FAST_ENTRY, + data: { provider: selected.provider, modelId: selected.id, mode }, + }; +} + +function deferred() { + let resolve!: (value: T) => void; + let reject!: (error: unknown) => void; + const promise = new Promise((done, fail) => { + resolve = done; + reject = fail; + }); + return { promise, resolve, reject }; +} + +function createFastControllerFixture() { + const clock = { now: 1_000 }; + const registry = { + getProvider: vi.fn(() => ({ + baseUrl: model.baseUrl, + getModels: () => [model, otherModel], + })), + isUsingOAuth: vi.fn(() => true), + getRegisteredProviderConfig: vi.fn( + () => undefined, + ), + getRegisteredNativeProvider: vi.fn( + () => undefined, + ), + getProviderAuth: vi.fn(async () => undefined), + }; + const ctx = { + model: model as FastHost["model"], + modelRegistry: registry, + hasUI: true as FastHost["hasUI"], + isIdle: vi.fn(() => true), + ui: { + notify: vi.fn(), + confirm: vi.fn(async () => true), + }, + sessionManager: { + getBranch: vi.fn(() => []), + getSessionId: vi.fn(() => "session-a"), + getLeafId: vi.fn(() => "leaf-a"), + }, + } satisfies FastHost; + const appendEntry = vi.fn(); + const fetchCatalog = vi.fn(async () => catalog); + const readAuthScope = vi.fn(async () => "scope-a"); + const controller = new FastController( + { appendEntry }, + "eleith:codex", + fetchCatalog, + readAuthScope, + () => clock.now, + ); + return { controller, ctx, appendEntry, fetchCatalog, readAuthScope, clock }; +} + +function stream(data: unknown): ProviderStreamEvent { + return { + type: "provider_stream_event", + provider: model.provider, + api: model.api, + model: model.id, + data, + }; +} diff --git a/extensions/codex/src/fast.ts b/extensions/codex/src/fast.ts new file mode 100644 index 0000000..2d74b6d --- /dev/null +++ b/extensions/codex/src/fast.ts @@ -0,0 +1,432 @@ +import type { + ExtensionAPI, + ExtensionContext, + ProviderStreamEvent, +} from "@earendil-works/pi-coding-agent"; +import { + CodexAccountError, + fetchCodexModelCatalog, + readCodexAuthScope, + type CodexAccountRegistry, +} from "./client.ts"; +import { codexFastEligibility, type FastEligibility } from "./capabilities.ts"; +import { + FAST_ENTRY, + readFastPreferences, + type FastChoice, + type FastPreference, +} from "./fast-state.ts"; +import { + formatFastStatus, + type FastRequestPolicy, + type ObservedTier, + type RequestedTier, +} from "./fast-status.ts"; +import { modelKey } from "./profiles.ts"; +import { + isDirectCodexSelection, + matchesDirectCodexSelection, + type DirectCodexHost, +} from "./request-scope.ts"; + +const CAPABILITY_TTL_MS = 5 * 60_000; + +export interface FastHost extends DirectCodexHost { + readonly modelRegistry: CodexAccountRegistry & + Pick< + ExtensionContext["modelRegistry"], + "getRegisteredProviderConfig" | "getRegisteredNativeProvider" + >; + readonly hasUI: boolean; + readonly isIdle: ExtensionContext["isIdle"]; + readonly ui: Pick; + readonly sessionManager: Pick< + ExtensionContext["sessionManager"], + "getBranch" | "getSessionId" | "getLeafId" + >; +} + +interface EligibilityPreparation { + abortController: AbortController; + sessionId: string; + leafId: string | null; + modelKey: string; + api: string; + baseUrl: string; +} + +interface CachedEligibility { + modelKey: string; + authScope: string; + eligibility: FastEligibility; + expiresAt: number; +} + +interface ModelTierActivity { + modelKey: string; + sessionId: string; + requested: RequestedTier | undefined; + observed: ObservedTier | undefined; + conflict: boolean; +} + +export class FastController { + private preferences = new Map(); + private cachedEligibility: CachedEligibility | undefined; + private pendingPreparation: EligibilityPreparation | undefined; + private modelTierActivity: ModelTierActivity | undefined; + // Honor endpoint Retry-After across model/branch changes; account switching is not supported. + private retryAt = 0; + private disposed = false; + + constructor( + private readonly pi: Pick, + private readonly commandName: string, + private readonly fetchCatalog: typeof fetchCodexModelCatalog = fetchCodexModelCatalog, + private readonly readAuthScope: typeof readCodexAuthScope = readCodexAuthScope, + private readonly now: () => number = Date.now, + ) {} + + restore(ctx: FastHost): void { + if (this.disposed) return; + this.invalidate(); + this.preferences = readFastPreferences(ctx.sessionManager.getBranch()); + } + + async handle(args: string, ctx: FastHost): Promise { + if (this.disposed) return; + const action = args.trim().toLowerCase(); + if (!action || action === "status") { + this.report(ctx); + return; + } + if (!["on", "off", "toggle", "reset"].includes(action)) { + ctx.ui.notify(`Usage: /${this.commandName} fast [on|off|toggle|status|reset]`, "warning"); + return; + } + if (this.pendingPreparation || !ctx.isIdle()) { + ctx.ui.notify("Wait for the current run or fast-mode change to finish.", "warning"); + return; + } + const model = ctx.model; + if (!model) { + ctx.ui.notify("No model selected.", "warning"); + return; + } + let mode: FastPreference; + if (action === "reset") mode = "inherit"; + else if (action === "off") mode = "off"; + else if (action === "toggle" && this.preferenceFor(modelKey(model)) === "on") mode = "off"; + else mode = "on"; + + // Reset only removes our choice; it is safe even after the selection leaves supported scope. + if (mode === "inherit") { + this.rememberChoice(ctx, mode); + this.report(ctx); + return; + } + if (!isDirectCodexSelection(ctx)) { + ctx.ui.notify( + "Fast controls require a direct official legacy Codex model with Pi OAuth and no provider override.", + "warning", + ); + return; + } + if (mode === "off") { + this.rememberChoice(ctx, mode); + this.report(ctx); + return; + } + if (!ctx.hasUI) { + ctx.ui.notify( + "Enabling fast mode requires interactive confirmation of higher allowance consumption.", + "warning", + ); + return; + } + await this.enable(ctx); + } + + // Only ordinary idle input prepares saved on choices. Never fetch from request/stream hooks. + async prepare(ctx: FastHost): Promise { + if (this.disposed || this.pendingPreparation || !ctx.isIdle()) return; + const model = ctx.model; + if (!model || this.preferenceFor(modelKey(model)) !== "on" || !isDirectCodexSelection(ctx)) + return; + const preparation = this.startPreparation(ctx); + try { + await this.refreshEligibility(ctx, preparation, false); + } catch (error) { + this.reportPreparationFailure(error, ctx, preparation); + } finally { + if (this.pendingPreparation === preparation) this.pendingPreparation = undefined; + } + } + + transform(payload: unknown, ctx: FastHost): Record | undefined { + if (this.disposed) return; + if (!matchesDirectCodexSelection(payload, ctx) || !ctx.model) { + this.modelTierActivity = undefined; + return; + } + const activity: ModelTierActivity = { + modelKey: modelKey(ctx.model), + sessionId: ctx.sessionManager.getSessionId(), + requested: undefined, + observed: undefined, + conflict: false, + }; + this.modelTierActivity = activity; + const policy = this.requestPolicy(ctx); + if (policy !== "priority" && policy !== "default") return; + if (payload.service_tier !== undefined && payload.service_tier !== policy) { + activity.conflict = true; + return; + } + activity.requested = policy; + if (payload.service_tier === policy) return; + return { ...payload, service_tier: policy }; + } + + observe(event: ProviderStreamEvent, ctx: FastHost): void { + if (this.disposed) return; + const activity = this.modelTierActivity; + if (!activity || !isDirectCodexSelection(ctx) || !ctx.model) return; + if ( + event.provider !== ctx.model.provider || + event.api !== ctx.model.api || + event.model !== ctx.model.id + ) + return; + if ( + activity.modelKey !== modelKey(ctx.model) || + activity.sessionId !== ctx.sessionManager.getSessionId() + ) + return; + const data = readObject(event.data); + if (!data || typeof data.type !== "string") return; + if (!["response.completed", "response.done", "response.incomplete"].includes(data.type)) return; + const tier = readObject(data.response)?.service_tier; + activity.observed = readObservedTier(tier); + } + + invalidate(): void { + this.pendingPreparation?.abortController.abort(); + this.pendingPreparation = undefined; + this.cachedEligibility = undefined; + this.modelTierActivity = undefined; + } + + dispose(): void { + if (this.disposed) return; + this.disposed = true; + this.invalidate(); + this.preferences.clear(); + } + + private async enable(ctx: FastHost): Promise { + const preparation = this.startPreparation(ctx); + try { + const confirmed = await ctx.ui.confirm( + "Enable Codex fast mode?", + "Priority requests use more of your Codex allowance. This branch/model choice can also affect retries, compaction and cache warming. Speed and priority processing are not guaranteed.", + { signal: preparation.abortController.signal }, + ); + if (!this.canApplyPreparation(ctx, preparation) || !confirmed) return; + const eligibility = await this.refreshEligibility(ctx, preparation, true); + if (!this.canApplyPreparation(ctx, preparation) || eligibility !== "supported") return; + this.rememberChoice(ctx, "on"); + this.pendingPreparation = undefined; + this.report(ctx); + } catch (error) { + this.reportPreparationFailure(error, ctx, preparation); + } finally { + if (this.pendingPreparation === preparation) this.pendingPreparation = undefined; + } + } + + private async refreshEligibility( + ctx: FastHost, + preparation: EligibilityPreparation, + force: boolean, + ): Promise { + const signal = preparation.abortController.signal; + const authScope = await this.readAuthScope(ctx.modelRegistry, signal); + if (!this.canApplyPreparation(ctx, preparation)) return; + const cached = this.cachedEligibility; + if ( + !force && + cached?.modelKey === preparation.modelKey && + cached.authScope === authScope && + cached.expiresAt > this.now() + ) + return cached.eligibility; + + this.cachedEligibility = undefined; + this.modelTierActivity = undefined; + if (this.now() < this.retryAt) { + ctx.ui.notify( + "Fast mode suspended: the catalog service is rate limited. Try again later.", + "warning", + ); + return; + } + const catalog = await this.fetchCatalog(ctx.modelRegistry, signal); + if (!this.canApplyPreparation(ctx, preparation)) return; + if (catalog.authScope !== authScope) { + ctx.ui.notify( + "Fast mode suspended: authentication changed during discovery. Try fast on again.", + "warning", + ); + return; + } + const model = ctx.model; + if (!model) return; + const eligibility = codexFastEligibility(catalog.data, model.id); + this.cachedEligibility = { + modelKey: preparation.modelKey, + authScope, + eligibility, + expiresAt: this.now() + CAPABILITY_TTL_MS, + }; + if (eligibility !== "supported") { + ctx.ui.notify( + `Fast mode suspended: model eligibility is ${eligibility}. Preference unchanged.`, + "warning", + ); + } + return eligibility; + } + + private startPreparation(ctx: FastHost): EligibilityPreparation { + const model = ctx.model; + if (!model) throw new Error("No model selected"); + const preparation: EligibilityPreparation = { + abortController: new AbortController(), + sessionId: ctx.sessionManager.getSessionId(), + leafId: ctx.sessionManager.getLeafId(), + modelKey: modelKey(model), + api: model.api, + baseUrl: model.baseUrl, + }; + this.pendingPreparation = preparation; + return preparation; + } + + private canApplyPreparation(ctx: FastHost, preparation: EligibilityPreparation): boolean { + // Shutdown/cancellation can expire every host getter; check ownership before touching ctx. + if ( + this.disposed || + this.pendingPreparation !== preparation || + preparation.abortController.signal.aborted + ) + return false; + const model = ctx.model; + return ( + !!model && + modelKey(model) === preparation.modelKey && + model.api === preparation.api && + model.baseUrl === preparation.baseUrl && + ctx.isIdle() && + ctx.sessionManager.getSessionId() === preparation.sessionId && + ctx.sessionManager.getLeafId() === preparation.leafId && + isDirectCodexSelection(ctx) + ); + } + + private reportPreparationFailure( + error: unknown, + ctx: FastHost, + preparation: EligibilityPreparation, + ): void { + if (!this.canApplyPreparation(ctx, preparation)) return; + this.cachedEligibility = undefined; + this.modelTierActivity = undefined; + if (error instanceof CodexAccountError && error.retryAfterMs !== undefined) { + this.retryAt = this.now() + error.retryAfterMs; + } + let message = "Could not prepare Codex fast mode."; + if (error instanceof CodexAccountError) message = error.message; + ctx.ui.notify(`Fast mode suspended: ${message} Preference unchanged.`, "warning"); + } + + private preferenceFor(selectedModelKey: string): FastPreference { + return this.preferences.get(selectedModelKey) ?? "inherit"; + } + + private rememberChoice(ctx: FastHost, mode: FastPreference): void { + const model = ctx.model; + if (!model) return; + const selectedModelKey = modelKey(model); + if (this.preferences.get(selectedModelKey) !== mode) { + const choice: FastChoice = { provider: model.provider, modelId: model.id, mode }; + this.pi.appendEntry(FAST_ENTRY, choice); + this.preferences.set(selectedModelKey, mode); + } + this.modelTierActivity = undefined; + } + + private currentEligibility(selectedModelKey: string): FastEligibility { + const cached = this.cachedEligibility; + if (!cached || cached.modelKey !== selectedModelKey || cached.expiresAt <= this.now()) + return "unknown"; + return cached.eligibility; + } + + private requestPolicy(ctx: FastHost): FastRequestPolicy { + if (!ctx.model || !isDirectCodexSelection(ctx)) return "outside scope"; + const selectedModelKey = modelKey(ctx.model); + const preference = this.preferenceFor(selectedModelKey); + if (preference === "inherit") return "untouched"; + if (preference === "off") return "default"; + if (this.pendingPreparation || this.currentEligibility(selectedModelKey) !== "supported") + return "suspended"; + return "priority"; + } + + private report(ctx: FastHost): void { + const model = ctx.model; + if (!model) { + ctx.ui.notify("No model selected.", "info"); + return; + } + const selectedModelKey = modelKey(model); + const direct = isDirectCodexSelection(ctx); + let activity = this.modelTierActivity; + if ( + !activity || + !direct || + activity.modelKey !== selectedModelKey || + activity.sessionId !== ctx.sessionManager.getSessionId() + ) + activity = undefined; + let eligibility: FastEligibility = "unknown"; + if (direct) eligibility = this.currentEligibility(selectedModelKey); + const status = { + preference: this.preferenceFor(selectedModelKey), + eligibility, + requestPolicy: this.requestPolicy(ctx), + lastRequested: activity?.requested, + lastObserved: activity?.observed, + conflict: activity?.conflict ?? false, + }; + ctx.ui.notify(formatFastStatus(status), "info"); + } +} + +function readObject(value: unknown): Record | undefined { + if (!value || typeof value !== "object" || Array.isArray(value)) return undefined; + return value as Record; +} + +function readObservedTier(value: unknown): ObservedTier | undefined { + if ( + value === "priority" || + value === "default" || + value === "auto" || + value === "flex" || + value === "fast" + ) + return value; + return undefined; +} diff --git a/extensions/codex/src/index.test.ts b/extensions/codex/src/index.test.ts index bd06bae..31a02c9 100644 --- a/extensions/codex/src/index.test.ts +++ b/extensions/codex/src/index.test.ts @@ -3,11 +3,13 @@ import { join } from "node:path"; import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent"; import { afterAll, expect, it, vi } from "vitest"; import codex from "./index.ts"; -import { fetchCodexUsage } from "./client.ts"; +import { fetchCodexModelCatalog, fetchCodexUsage, readCodexAuthScope } from "./client.ts"; vi.mock("./client.ts", async (importOriginal) => ({ ...(await importOriginal()), fetchCodexUsage: vi.fn(), + fetchCodexModelCatalog: vi.fn(), + readCodexAuthScope: vi.fn(), })); const directory = await vi.hoisted(async () => { @@ -28,6 +30,8 @@ afterAll(() => { async function load(prefix = "eleith") { vi.mocked(fetchCodexUsage).mockReset(); + vi.mocked(fetchCodexModelCatalog).mockReset(); + vi.mocked(readCodexAuthScope).mockReset(); writeFileSync(configPath, JSON.stringify({ commandPrefix: prefix })); type Command = Parameters[1]; type Listener = ( @@ -91,6 +95,7 @@ it("registers only the configured native command once after binding", async () = { value: "extend", label: "extend" }, { value: "restore", label: "restore" }, { value: "usage", label: "usage" }, + { value: "fast", label: "fast" }, ]); }); @@ -102,7 +107,7 @@ it("handles no argument as status rather than an invalid default action", async expect(context.pi.setModel).not.toHaveBeenCalled(); await context.command().handler("extend restore", context.ctx); expect(context.notify).toHaveBeenLastCalledWith( - "Usage: /eleith:codex [extend|restore|usage]", + "Usage: /eleith:codex [extend|restore|usage|fast]", "warning", ); }); @@ -121,6 +126,30 @@ it("wires restore/reconcile hooks without provider calls for an unsupported mode expect(context.pi.appendEntry).not.toHaveBeenCalled(); expect(context.pi.on.mock.calls.map(([name]) => name)).not.toContain("before_agent_start"); expect(fetchCodexUsage).not.toHaveBeenCalled(); + expect(fetchCodexModelCatalog).not.toHaveBeenCalled(); + expect(readCodexAuthScope).not.toHaveBeenCalled(); +}); + +it("routes fast subcommands without allowing trailing arguments on other actions", async () => { + const context = await load("personal"); + await context.emit("session_start"); + await context.command().handler("fast status", context.ctx); + expect(context.notify).toHaveBeenLastCalledWith("No model selected.", "info"); + await context.command().handler("fast off extra", context.ctx); + expect(context.notify).toHaveBeenLastCalledWith( + "Usage: /personal:codex fast [on|off|toggle|status|reset]", + "warning", + ); + await context.command().handler("restore extra", context.ctx); + expect(context.notify).toHaveBeenLastCalledWith( + "Usage: /personal:codex [extend|restore|usage|fast]", + "warning", + ); + expect(fetchCodexModelCatalog).not.toHaveBeenCalled(); + expect(readCodexAuthScope).not.toHaveBeenCalled(); + expect(fetchCodexUsage).not.toHaveBeenCalled(); + expect(context.pi.appendEntry).not.toHaveBeenCalled(); + expect(context.pi.setModel).not.toHaveBeenCalled(); }); it("routes usage through the configured command without requiring an idle Codex selection", async () => { diff --git a/extensions/codex/src/index.ts b/extensions/codex/src/index.ts index 6fff236..5930ddf 100644 --- a/extensions/codex/src/index.ts +++ b/extensions/codex/src/index.ts @@ -3,38 +3,53 @@ import { registerCommand } from "@eleith/pi-internal/commands"; import { getSettings } from "@eleith/pi-internal/config"; import { ContextController } from "./controller.ts"; import { UsageController } from "./usage-controller.ts"; +import { FastController } from "./fast.ts"; export default async function codex(pi: ExtensionAPI): Promise { const settings = await getSettings(); const commandName = `${settings.commandPrefix}:codex`; const controller = new ContextController(pi, commandName); const usage = new UsageController(); + const fast = new FastController(pi, commandName); registerCommand(pi, { name: commandName, - description: "Codex context controls and account usage", - actions: ["extend", "restore", "usage"], + description: "Codex context, account usage and fast controls", + actions: ["extend", "restore", "usage", "fast"], + argumentActions: ["fast"], defaultAction: "", - handler(action, ctx) { + handler(action, ctx, argument) { if (action === "usage") return usage.report(ctx); + if (action === "fast") return fast.handle(argument ?? "", ctx); return controller.handle(action, ctx); }, }); - pi.on("session_start", (_event, ctx) => controller.restore(ctx)); + pi.on("session_start", (_event, ctx) => { + fast.restore(ctx); + return controller.restore(ctx); + }); pi.on("session_tree", (_event, ctx) => { usage.cancel(); + fast.restore(ctx); return controller.restore(ctx); }); pi.on("model_select", (_event, ctx) => { usage.cancel(); + fast.invalidate(); return controller.modelSelected(ctx); }); pi.on("session_shutdown", (_event, ctx) => { usage.dispose(); + fast.dispose(); controller.shutdown(ctx); }); pi.on("input", async (event, ctx) => { - if (event.streamingBehavior === undefined && ctx.isIdle()) await controller.reconcile(ctx); + if (event.streamingBehavior === undefined && ctx.isIdle()) { + await controller.reconcile(ctx); + await fast.prepare(ctx); + } }); + pi.on("before_provider_request", (event, ctx) => fast.transform(event.payload, ctx)); + pi.on("provider_stream_event", (event, ctx) => fast.observe(event, ctx)); } diff --git a/extensions/codex/src/request-scope-sdk.test.ts b/extensions/codex/src/request-scope-sdk.test.ts new file mode 100644 index 0000000..a149f8f --- /dev/null +++ b/extensions/codex/src/request-scope-sdk.test.ts @@ -0,0 +1,237 @@ +import { rmSync } from "node:fs"; +import { zstdDecompressSync } from "node:zlib"; +import { afterAll, afterEach, beforeEach, expect, it, vi } from "vitest"; +import { + InMemoryModelsStore, + type Api, + type CredentialStore, + type Model, + type OAuthCredential, +} from "@earendil-works/pi-ai"; +import { + createAgentSession, + DefaultResourceLoader, + ModelRuntime, + SessionManager, + SettingsManager, +} from "@earendil-works/pi-coding-agent"; +import { matchesDirectCodexSelection } from "./request-scope.ts"; + +const directory = await vi.hoisted(async () => { + const { mkdtempSync } = await import("node:fs"); + const { tmpdir } = await import("node:os"); + const { join } = await import("node:path"); + const directory = mkdtempSync(join(tmpdir(), "codex-request-scope-sdk-")); + vi.stubEnv("PI_CODING_AGENT_DIR", directory); + vi.stubEnv("PI_OFFLINE", "1"); + return directory; +}); +beforeEach(blockUnexpectedNetwork); +afterEach(() => { + vi.restoreAllMocks(); + vi.unstubAllGlobals(); +}); +afterAll(() => { + vi.unstubAllEnvs(); + rmSync(directory, { recursive: true, force: true }); +}); + +it("SDK direct-selection payload transforms reach the native Codex SSE request", async () => { + const httpRequest = stubCodexResponse(); + const { runtime, model } = await createOfflineCodexRuntime(httpRequest); + const { session, decisions, errors, closeSession } = await openProbeSession(runtime, model); + try { + await session.prompt("Reply with ok."); + expect(decisions).toEqual([true]); + expect(httpRequest).toHaveBeenCalledOnce(); + const [, init] = httpRequest.mock.calls[0]; + expect(readHttpPayload(init)).toMatchObject({ model: model.id, service_tier: "priority" }); + expect(session.getLastAssistantText()).toBe("ok"); + expect(errors).toEqual([]); + } finally { + await closeSession(); + } +}); + +it("SDK virtual selections remain outside scope even when their dispatch is native Codex", async () => { + const httpRequest = stubCodexResponse(); + const { runtime, model } = await createOfflineCodexRuntime(httpRequest); + runtime.registerVirtualModel({ + provider: "openai-codex", + id: "auto", + name: "Test route", + route: () => ({ model, thinkingLevel: "off" }), + }); + const selected = runtime.getModel("openai-codex", "auto"); + if (!selected) throw new Error("Virtual test model missing"); + const { session, decisions, errors, closeSession } = await openProbeSession(runtime, selected); + try { + await session.prompt("Reply with ok."); + expect(decisions).toEqual([false]); + expect(session.model?.api).toBe("pi-virtual"); + const [, init] = httpRequest.mock.calls[0]; + const payload = readHttpPayload(init); + expect(payload.model).toBe(model.id); + expect(payload).not.toHaveProperty("service_tier"); + expect(session.getLastAssistantText()).toBe("ok"); + expect(errors).toEqual([]); + } finally { + await closeSession(); + } +}); + +async function createOfflineCodexRuntime(httpRequest: typeof fetch) { + const payload = { "https://api.openai.com/auth": { chatgpt_account_id: "scope-test-account" } }; + const access = `e30.${Buffer.from(JSON.stringify(payload)).toString("base64")}.fake-signature`; + const credential: OAuthCredential = { + type: "oauth", + access, + refresh: "must-not-refresh", + expires: Date.now() + 3_600_000, + }; + const credentials: CredentialStore = { + read: async (id) => { + if (id !== "openai-codex") return undefined; + return credential; + }, + list: async () => [{ providerId: "openai-codex", type: "oauth" }], + modify: async () => { + throw new Error("Credential writes forbidden"); + }, + delete: async () => { + throw new Error("Credential deletion forbidden"); + }, + }; + const runtime = await ModelRuntime.create({ + credentials, + modelsPath: null, + modelsStore: new InMemoryModelsStore(), + allowModelNetwork: false, + refreshOnCreate: false, + }); + const refresh = await runtime.refresh({ providers: ["openai-codex"], allowNetwork: false }); + expect(refresh.errors.size).toBe(0); + const model = runtime.getModel("openai-codex", "gpt-6-sol"); + if (!model) throw new Error("Built-in test model missing"); + + const nativeStream = runtime.streamSimple.bind(runtime); + // Force this transport through a public test seam; keep the real provider and extension bridge. + vi.spyOn(runtime, "streamSimple").mockImplementation((selected, context, options) => + nativeStream(selected, context, { + ...options, + transport: "sse", + cacheRetention: "none", + fetch: httpRequest, + }), + ); + return { runtime, model }; +} + +async function openProbeSession(runtime: ModelRuntime, model: Model) { + const decisions: boolean[] = []; + const settingsManager = SettingsManager.inMemory({ + packages: [], + enableInstallTelemetry: false, + enableAnalytics: false, + compaction: { enabled: false }, + retry: { enabled: false }, + cacheWarming: "off", + }); + const resourceLoader = new DefaultResourceLoader({ + cwd: directory, + agentDir: directory, + settingsManager, + noExtensions: true, + noSkills: true, + noPromptTemplates: true, + noThemes: true, + noContextFiles: true, + extensionFactories: [ + (pi) => { + pi.on("before_provider_request", (event, ctx) => { + const payload = event.payload; + const matches = matchesDirectCodexSelection(payload, ctx); + decisions.push(matches); + // Test-only transform; production registration belongs to the fast-command changeset. + if (matches) return { ...payload, service_tier: "priority" }; + }); + }, + ], + }); + await resourceLoader.reload(); + const { session } = await createAgentSession({ + cwd: directory, + agentDir: directory, + resourceLoader, + settingsManager, + modelRuntime: runtime, + model, + thinkingLevel: "off", + sessionManager: SessionManager.inMemory(directory), + noTools: "all", + }); + const errors: unknown[] = []; + await session.bindExtensions({ mode: "print", onError: (error) => errors.push(error) }); + return { + session, + decisions, + errors, + async closeSession() { + await session.extensionRunner.emit({ type: "session_shutdown", reason: "quit" }); + session.dispose(); + }, + }; +} + +function stubCodexResponse() { + const message = { + type: "message", + id: "answer", + role: "assistant", + status: "completed", + content: [{ type: "output_text", text: "ok", annotations: [] }], + }; + const events = [ + { type: "response.created", response: { id: "response" } }, + { type: "response.output_item.added", output_index: 0, item: { ...message, content: [] } }, + { type: "response.output_text.delta", output_index: 0, delta: "ok" }, + { type: "response.output_item.done", output_index: 0, item: message }, + { + type: "response.completed", + response: { id: "response", status: "completed", output: [message] }, + }, + ]; + const body = events.map((event) => `data: ${JSON.stringify(event)}\n\n`).join(""); + return vi.fn(async (url) => { + expect(String(url)).toBe("https://chatgpt.com/backend-api/codex/responses"); + return new Response(body, { headers: { "Content-Type": "text/event-stream" } }); + }); +} + +function blockUnexpectedNetwork() { + vi.stubGlobal( + "fetch", + vi.fn(async () => { + throw new Error("Unexpected network request"); + }), + ); + vi.stubGlobal( + "WebSocket", + class { + constructor() { + throw new Error("Unexpected WebSocket request"); + } + }, + ); +} + +function readHttpPayload(init: RequestInit | undefined): Record { + if (typeof init?.body === "string") return JSON.parse(init.body); + if ( + init?.body instanceof Uint8Array && + new Headers(init.headers).get("content-encoding") === "zstd" + ) { + return JSON.parse(zstdDecompressSync(init.body).toString("utf8")); + } + throw new Error("Unexpected Codex request body"); +} diff --git a/extensions/codex/src/request-scope.test.ts b/extensions/codex/src/request-scope.test.ts new file mode 100644 index 0000000..63f7403 --- /dev/null +++ b/extensions/codex/src/request-scope.test.ts @@ -0,0 +1,79 @@ +import type { Api, Model } from "@earendil-works/pi-ai"; +import { openaiCodexProvider } from "@earendil-works/pi-ai/providers/openai-codex"; +import { expect, it, vi } from "vitest"; +import { matchesDirectCodexSelection, type DirectCodexHost } from "./request-scope.ts"; + +const model: Model = { + id: "gpt-6-sol", + name: "Sol", + provider: "openai-codex", + api: "openai-codex-responses", + baseUrl: "https://chatgpt.com/backend-api", + reasoning: true, + input: ["text"], + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 }, + contextWindow: 272_000, + maxTokens: 128_000, +}; + +type Registry = DirectCodexHost["modelRegistry"]; + +function fixture() { + const registry = { + getProvider: vi.fn(() => ({ baseUrl: model.baseUrl })), + isUsingOAuth: vi.fn(() => true), + getRegisteredProviderConfig: vi.fn(() => undefined), + getRegisteredNativeProvider: vi.fn(() => undefined), + }; + return { model, modelRegistry: registry }; +} + +it("matches a direct selection without modifying the payload or taking ownership of an existing tier", () => { + const payload = Object.freeze({ model: model.id, service_tier: "flex", input: [] }); + expect(matchesDirectCodexSelection(payload, fixture())).toBe(true); + expect(payload).toEqual({ model: model.id, service_tier: "flex", input: [] }); +}); + +it.each([ + undefined, + { ...model, provider: "openai" }, + { ...model, api: "pi-virtual" }, + { ...model, baseUrl: "https://gateway.invalid" }, +])("rejects unsupported selections: %j", (selected) => { + expect(matchesDirectCodexSelection({ model: model.id }, { ...fixture(), model: selected })).toBe( + false, + ); +}); + +it.each([null, [], {}, { model: "different-model" }, { model: 42 }])( + "requires an object payload naming the selected model: %j", + (payload) => { + expect(matchesDirectCodexSelection(payload, fixture())).toBe(false); + }, +); + +it("rejects an effective endpoint override even if the selected model still looks official", () => { + const host = fixture(); + host.modelRegistry.getProvider.mockReturnValue({ baseUrl: "https://gateway.invalid" }); + expect(matchesDirectCodexSelection({ model: model.id }, host)).toBe(false); +}); + +it("excludes registered provider overlays", () => { + const host = fixture(); + host.modelRegistry.getRegisteredProviderConfig.mockReturnValue({ + headers: { "X-Example": "overlay" }, + }); + expect(matchesDirectCodexSelection({ model: model.id }, host)).toBe(false); +}); + +it("excludes registered native replacements, including a re-registered stock provider", () => { + const host = fixture(); + host.modelRegistry.getRegisteredNativeProvider.mockReturnValue(openaiCodexProvider()); + expect(matchesDirectCodexSelection({ model: model.id }, host)).toBe(false); +}); + +it("does not match a selection whose auth snapshot is not OAuth", () => { + const host = fixture(); + host.modelRegistry.isUsingOAuth.mockReturnValue(false); + expect(matchesDirectCodexSelection({ model: model.id }, host)).toBe(false); +}); diff --git a/extensions/codex/src/request-scope.ts b/extensions/codex/src/request-scope.ts new file mode 100644 index 0000000..d9b02d2 --- /dev/null +++ b/extensions/codex/src/request-scope.ts @@ -0,0 +1,42 @@ +import type { Api, Model, Provider } from "@earendil-works/pi-ai"; +import type { ExtensionContext } from "@earendil-works/pi-coding-agent"; + +const PROVIDER = "openai-codex"; +const BASE_URL = "https://chatgpt.com/backend-api"; + +type RequestRegistry = Pick< + ExtensionContext["modelRegistry"], + "isUsingOAuth" | "getRegisteredProviderConfig" | "getRegisteredNativeProvider" +> & { + getProvider(id: string): Pick | undefined; +}; + +export interface DirectCodexHost { + readonly model: Model | undefined; + readonly modelRegistry: RequestRegistry; +} + +// Pi exposes the selection, not physical request identity; this only checks direct-selection scope. +export function matchesDirectCodexSelection( + payload: unknown, + ctx: DirectCodexHost, +): payload is Record { + if (!isDirectCodexSelection(ctx)) return false; + if (!payload || typeof payload !== "object" || Array.isArray(payload)) return false; + return "model" in payload && payload.model === ctx.model?.id; +} + +export function isDirectCodexSelection(ctx: DirectCodexHost): boolean { + const model = ctx.model; + if (!model || model.provider !== PROVIDER || model.api !== "openai-codex-responses") return false; + if (!officialBase(model.baseUrl)) return false; + const registry = ctx.modelRegistry; + if (!officialBase(registry.getProvider(PROVIDER)?.baseUrl)) return false; + if (registry.getRegisteredProviderConfig(PROVIDER) !== undefined) return false; + if (registry.getRegisteredNativeProvider(PROVIDER) !== undefined) return false; + return registry.isUsingOAuth(model); +} + +function officialBase(url: string | undefined): boolean { + return url?.replace(/\/+$/, "") === BASE_URL; +}