diff --git a/AGENTS.md b/AGENTS.md index 69e25c35..7c88ce34 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -32,7 +32,7 @@ Package version: 6.1.8 | Installer behavior | `scripts/install-oc-codex-multi-auth.js`, `scripts/install-oc-codex-multi-auth-core.js` | npm bin, config merge, cache cleanup, TUI config enablement | | Plugin orchestration | `index.ts` | OAuth loader, request pipeline, metrics, recovery, `ToolContext` assembly | | TUI quota status | `tui.ts`, `lib/tui-status.ts`, `lib/tui-quota-cache.ts`, `lib/codex-usage.ts` | prompt quota status, quota details, shared quota cache | -| Tool registry | `lib/tools/index.ts` + `lib/tools/codex-*.ts` | 21 registered `codex-*` tools | +| Tool registry | `lib/tools/index.ts` + `lib/tools/codex-*.ts` | 24 registered `codex-*` tools | | OAuth flow + PKCE | `lib/auth/auth.ts`, `lib/auth/server.ts`, `lib/auth/device-code.ts`, `lib/auth/login-runner.ts` | browser/device/manual login, token refresh, workspace selection | | OAuth scopes | `lib/auth/scopes.ts` | connector scope validation and re-auth checks | | Multi-account rotation | `lib/accounts.ts`, `lib/accounts/`, `lib/rotation.ts` | health scoring, cooldowns, token bucket, recovery | diff --git a/README.md b/README.md index 714eac21..8249e646 100644 --- a/README.md +++ b/README.md @@ -180,6 +180,7 @@ If browser launch is blocked, use the alternate login paths in [docs/getting-sta | `codex-limits` | What quota or rate-limit state is visible now? | | `codex-reset` | Do I have a banked rate-limit reset credit, and how do I redeem it? | | `codex-dashboard` | Can I manage accounts from one interactive surface? | +| `codex-pool` | Which accounts are preferred for each model, and how do I change them? | Most of these also run as a **direct CLI** with no agent/model involvement (no token cost) — e.g. `oc-codex-multi-auth warm`, `oc-codex-multi-auth status`, or `npx -y oc-codex-multi-auth@latest warm`. Use `oc-codex-multi-auth warm` to open every enabled account's usage window at the start of a session and stagger the rolling quota cooldowns. Add `--json` for scriptable output. @@ -244,6 +245,58 @@ Primary config files: - `~/.config/opencode/tui.json` - `~/.opencode/openai-codex-auth-config.json` +### Route models to preferred accounts + +Use `modelAccountPools` to assign one or more preferred ChatGPT accounts to a +model. Account references use stable account IDs, so adding, removing, or +reordering accounts does not silently change a model's routing. + +```json +{ + "modelAccountPools": { + "gpt-5.6-sol": [ + "org-example-account-id", + "00000000-0000-0000-0000-000000000000" + ], + "gpt-5.6-terra": [ + "org-another-account-id" + ] + } +} +``` + +Save this configuration in `~/.opencode/openai-codex-auth-config.json`, then +restart OpenCode. Model matching is case-insensitive and uses the effective +model after request model normalization. + +Use `codex-pool` to manage these mappings with ordinary 1-based account +numbers. The tool resolves those numbers and writes stable IDs to disk: + +```text +codex-pool +codex-pool action="set" model="gpt-5.6-sol" accounts=[7,8] +codex-pool action="add" model="gpt-5.6-sol" accounts=[9] +codex-pool action="remove" model="gpt-5.6-sol" accounts=[7] +codex-pool action="clear" model="gpt-5.6-sol" +``` + +Add `dryRun=true` to preview a mutation. Use `format="json"` for structured +output; stable IDs remain redacted unless `includeSensitive=true` is also set. +Restart OpenCode after an applied mutation. The plugin configuration is global +while account storage is per-project by default, so a reference unresolved in +the current project is reported but never automatically deleted. + +Routing behavior: + +- A mapped model uses only healthy, selectable accounts in its preferred pool. +- Existing rotation strategy, quota, cooldown, and token-health rules still apply within the preferred pool. +- If every preferred account is unavailable, disabled, unknown, cooling down, or rate-limited, routing automatically falls back to the healthy general account pool. +- An unmapped model or an empty account list uses the general account pool directly. +- `codex-status`, `codex-dashboard`, and routing diagnostics report the account-pool mode as `preferred`, `general`, or `general-fallback`. + +Account IDs are local account metadata but should still be treated as private +configuration. Do not publish a populated configuration file. + Selected runtime/environment overrides: | Variable | Effect | diff --git a/docs/architecture.md b/docs/architecture.md index 49da45b7..1594f84a 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -67,13 +67,13 @@ Native mode keeps the host payload shape whenever possible. Legacy mode applies ### 4. Tool registry -`lib/tools/index.ts` builds the OpenCode tool map from 21 per-file factories under `lib/tools/`. +`lib/tools/index.ts` builds the OpenCode tool map from 24 per-file factories under `lib/tools/`. Common groups: - setup: `codex-setup`, `codex-help`, `codex-next` -- daily account use: `codex-list`, `codex-switch`, `codex-status`, `codex-limits`, `codex-reset` -- account metadata: `codex-label`, `codex-tag`, `codex-note`, `codex-remove`, `codex-refresh` +- daily account use: `codex-list`, `codex-switch`, `codex-warm`, `codex-status`, `codex-limits`, `codex-reset` +- account metadata and routing: `codex-label`, `codex-tag`, `codex-note`, `codex-pool`, `codex-remove`, `codex-refresh` - diagnostics and resilience: `codex-health`, `codex-metrics`, `codex-doctor`, `codex-diag`, `codex-diff` - backup and secrets: `codex-export`, `codex-import`, `codex-keychain` - interactive surface: `codex-dashboard` diff --git a/docs/configuration.md b/docs/configuration.md index 1cbf0c0b..c26dfed2 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -155,6 +155,9 @@ advanced settings go in `~/.opencode/openai-codex-auth-config.json`: "fastSession": false, "fastSessionStrategy": "hybrid", "fastSessionMaxInputItems": 30, + "modelAccountPools": { + "gpt-5.6-sol": ["org-example-account-id"] + }, "retryProfile": "balanced", "retryBudgetOverrides": { "network": 2, @@ -193,6 +196,7 @@ The sample above intentionally sets `"retryAllAccountsMaxRetries": 3` as a bound | `fastSession` | `false` | forces low-latency settings per request (`reasoningEffort=none/low`, `reasoningSummary=auto`, `textVerbosity=low`) | | `fastSessionStrategy` | `hybrid` | `hybrid` speeds simple turns and keeps full-depth for complex prompts; `always` forces fast mode every turn | | `fastSessionMaxInputItems` | `30` | max input items kept when fast mode is applied | +| `modelAccountPools` | `{}` | optional map of effective model IDs to preferred stable account IDs; selection uses the configured pool while it has a healthy account, then falls back to the general account pool | | `retryProfile` | `balanced` | retry budget profile for request classes (`conservative`, `balanced`, `aggressive`) | | `retryBudgetOverrides` | `{}` | optional per-class budget overrides (`authRefresh`, `network`, `server`, `rateLimitShort`, `rateLimitGlobal`, `emptyResponse`) | | `perProjectAccounts` | `true` | each project gets its own account storage | @@ -212,6 +216,18 @@ The sample above intentionally sets `"retryAllAccountsMaxRetries": 3` as a bound | `fetchTimeoutMs` | `60000` | upstream fetch timeout in ms | | `streamStallTimeoutMs` | `45000` | max time to wait for next SSE chunk before aborting | +Use `codex-pool action="set" model="gpt-5.6-sol" accounts=[7,8]` to manage a +pool with 1-based account numbers while persisting stable IDs. The tool also +supports `status` (default), `add`, `remove`, `clear`, `dryRun=true`, and JSON +output. Restart OpenCode after an applied mutation. Because this config is +global while account storage is per-project by default, references unavailable +in the current project are reported but not automatically pruned. + +Model keys are matched case-insensitively after request model normalization. +Empty lists, unmapped models, and pools with no selectable accounts use the +general account pool. Routing diagnostics expose the resulting mode as +`preferred`, `general`, or `general-fallback`. + ### Beginner Safe Mode Behavior when `beginnerSafeMode` is enabled (`true` or `CODEX_AUTH_BEGINNER_SAFE_MODE=1`), the plugin applies a safer retry profile automatically: diff --git a/docs/development/ARCHITECTURE.md b/docs/development/ARCHITECTURE.md index c7f0fed5..8651c82d 100644 --- a/docs/development/ARCHITECTURE.md +++ b/docs/development/ARCHITECTURE.md @@ -45,7 +45,7 @@ index.ts |- ToolContext construction v lib/tools/index.ts - |- registers 23 OpenCode tools + |- registers 24 OpenCode tools |- each tool delegates to lib/tools/codex-*.ts Request path @@ -90,7 +90,7 @@ tui.ts | Storage | `lib/storage.ts`, `lib/storage/` | V3 JSON storage, atomic writes, migrations, per-project paths, backups, import/export, keychain opt-in, flagged accounts | | Request bridge | `lib/request/fetch-helpers.ts`, `lib/request/request-transformer.ts`, `lib/request/response-handler.ts`, `lib/request/retry-budget.ts`, `lib/request/rate-limit-backoff.ts` | URL/body/header shaping, Codex invariants, SSE conversion, retry budgets, backoff, error mapping | | Model/prompt mapping | `lib/prompts/codex.ts`, `lib/prompts/opencode-codex.ts`, `lib/prompts/codex-opencode-bridge.ts`, `lib/request/helpers/model-map.ts` | model-family detection, Codex instructions cache, OpenCode prompt adaptation, fallback aliases | -| Tool registry | `lib/tools/index.ts`, `lib/tools/codex-*.ts` | 23 OpenCode tools for setup, account switching, status, health, quota resets, diagnostics, backup, keychain, and recovery | +| Tool registry | `lib/tools/index.ts`, `lib/tools/codex-*.ts` | 24 OpenCode tools for setup, account switching, status, health, quota resets, diagnostics, backup, keychain, and recovery | | Runtime support | `lib/runtime.ts`, `lib/circuit-breaker.ts`, `lib/proactive-refresh.ts`, `lib/parallel-probe.ts`, `lib/recovery/`, `lib/shutdown.ts` | pure runtime helpers, failure isolation, refresh scheduling, health probing, session recovery, cleanup | | UI helpers | `lib/ui/` | terminal formatting, auth menu, select/confirm prompts, theme/color handling, beginner checklist | | Config templates | `config/opencode-modern.json`, `config/opencode-legacy.json`, `config/minimal-opencode.json`, `config/README.md` | copy-paste OpenCode provider templates and model catalog guidance | @@ -169,7 +169,7 @@ Legacy mode exists for compatibility with older OpenCode/AI SDK payload behavior ## Tool Registry Architecture -The plugin exposes 23 OpenCode tools through `lib/tools/index.ts`. `index.ts` builds one `ToolContext` from plugin-closure state and helper functions, then passes it to `createToolRegistry(ctx)`. +The plugin exposes 24 OpenCode tools through `lib/tools/index.ts`. `index.ts` builds one `ToolContext` from plugin-closure state and helper functions, then passes it to `createToolRegistry(ctx)`. Why this shape exists: @@ -183,8 +183,8 @@ Tool groups: | Group | Tools | | --- | --- | | Setup and help | `codex-setup`, `codex-help`, `codex-next` | -| Daily account use | `codex-list`, `codex-switch`, `codex-status`, `codex-limits`, `codex-reset`, `codex-dashboard` | -| Account metadata | `codex-label`, `codex-tag`, `codex-note`, `codex-remove`, `codex-refresh` | +| Daily account use | `codex-list`, `codex-switch`, `codex-warm`, `codex-status`, `codex-limits`, `codex-reset`, `codex-dashboard` | +| Account metadata and routing | `codex-label`, `codex-tag`, `codex-note`, `codex-pool`, `codex-remove`, `codex-refresh` | | Diagnostics | `codex-health`, `codex-metrics`, `codex-doctor`, `codex-diag`, `codex-diff` | | Backup/secrets | `codex-export`, `codex-import`, `codex-keychain` | diff --git a/docs/development/CONFIG_FIELDS.md b/docs/development/CONFIG_FIELDS.md index 5703555e..e578050b 100644 --- a/docs/development/CONFIG_FIELDS.md +++ b/docs/development/CONFIG_FIELDS.md @@ -171,6 +171,36 @@ Examples: This normalization is why legacy aliases and snapshot-like IDs can still route to a stable family while preserving the user-facing config surface. +## `modelAccountPools` + +The plugin runtime config at `~/.opencode/openai-codex-auth-config.json` can +map effective model IDs to preferred stable account IDs: + +```json +{ + "modelAccountPools": { + "gpt-5.6-sol": ["org-example-account-id"] + } +} +``` + +The request pipeline resolves the pool after model normalization. All rotation +strategies restrict selection to healthy accounts in the preferred pool while +one is available. If the configured IDs are unknown or every preferred account +is disabled, cooling down, rate-limited, or locally depleted, selection falls +back to the general account pool. Empty lists and unmapped models use the +general pool directly. + +`codex-pool` is the supported mutation surface. It accepts 1-based account +numbers for `set`, `add`, and `remove`, but resolves and atomically persists +only stable account IDs. `clear` removes a model mapping, and every mutation +supports a dry-run preview. Writes preserve unrelated raw config fields and +refuse to replace malformed JSON or an invalid existing pool. + +The config file is global while account storage is per-project by default. +Consequently, status may report unresolved references for the current project; +the tool does not automatically prune them because they may be valid elsewhere. + ## Verification Notes Use these commands when validating config fields: diff --git a/index.ts b/index.ts index 78f10798..5c008bf2 100644 --- a/index.ts +++ b/index.ts @@ -77,6 +77,7 @@ import { getEmptyResponseRetryDelayMs, getPidOffsetEnabled, getRotationStrategy, + getModelAccountPool, getFetchTimeoutMs, getStreamStallTimeoutMs, getCodexTuiV2, @@ -316,6 +317,8 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { fallbackFrom: snapshot?.fallbackFrom ?? null, fallbackTo: snapshot?.fallbackTo ?? null, fallbackReason: snapshot?.fallbackReason ?? null, + accountPoolMode: snapshot?.accountPoolMode ?? null, + configuredAccountPoolSize: snapshot?.configuredAccountPoolSize ?? 0, selectionExplainability: serializeSelectionExplainability( options.selectionExplainability ?? snapshot?.explainability ?? [], ), @@ -372,6 +375,8 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { lines.push(` Fallback from: ${formatRoutingValue(routing.fallbackFrom)}`); lines.push(` Fallback to: ${formatRoutingValue(routing.fallbackTo)}`); lines.push(` Fallback reason: ${formatRoutingValue(routing.fallbackReason)}`); + lines.push(` Account pool: ${formatRoutingValue(routing.accountPoolMode)}`); + lines.push(` Configured pool size: ${routing.configuredAccountPoolSize}`); if (options.includeExplainability) { lines.push(" Selection explainability:"); if (routing.selectionExplainability.length === 0) { @@ -468,6 +473,22 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { routing.fallbackReason ? "warning" : "muted", ), ); + lines.push( + formatUiKeyValue( + ui, + "Account pool", + formatRoutingValue(routing.accountPoolMode), + routing.accountPoolMode === "general-fallback" ? "warning" : "muted", + ), + ); + lines.push( + formatUiKeyValue( + ui, + "Configured pool size", + String(routing.configuredAccountPoolSize), + "muted", + ), + ); if (options.includeExplainability) { lines.push(""); lines.push(...formatUiSection(ui, "Selection explainability")); @@ -1923,8 +1944,9 @@ export const OpenAIOAuthPlugin: Plugin = async ({ client }: PluginInput) => { const attempted = new Set(); let restartAccountTraversalWithFallback = false; let restartAccountTraversalAfterWorkspaceDeactivation = false; + const preferredAccountIds = getModelAccountPool(pluginConfig, model); -while (attempted.size < Math.max(1, accountCount)) { + while (attempted.size < Math.max(1, accountCount)) { const selectionExplainability = accountManager.getSelectionExplainability( modelFamily, model, @@ -1943,8 +1965,15 @@ while (attempted.size < Math.max(1, accountCount)) { fallbackFrom, fallbackTo, fallbackReason, + configuredAccountPoolSize: preferredAccountIds.length, }; - const account = accountManager.getAccountForStrategy(rotationStrategy, modelFamily, model, { pidOffsetEnabled }); + const account = accountManager.getAccountForStrategy( + rotationStrategy, + modelFamily, + model, + { pidOffsetEnabled }, + preferredAccountIds, + ); if (!account || attempted.has(account.index)) { break; } @@ -1961,8 +1990,16 @@ while (attempted.size < Math.max(1, accountCount)) { fallbackApplied, fallbackFrom, fallbackTo, - fallbackReason, - }; + fallbackReason, + accountPoolMode: + preferredAccountIds.length === 0 + ? "general" + : account.accountId !== undefined && + preferredAccountIds.includes(account.accountId) + ? "preferred" + : "general-fallback", + configuredAccountPoolSize: preferredAccountIds.length, + }; } // Log account selection for debugging rotation logDebug( diff --git a/lib/AGENTS.md b/lib/AGENTS.md index 822ce681..16a6faa2 100644 --- a/lib/AGENTS.md +++ b/lib/AGENTS.md @@ -41,7 +41,7 @@ lib/ ├── storage.ts # V3 JSON storage facade ├── storage/ # atomic writes, paths, migrations, keychain, backup/import/export ├── table-formatter.ts # CLI table formatting -├── tools/ # 21 codex-* tool factories + registry +├── tools/ # 24 codex-* tool factories + registry ├── tui-quota-cache.ts # shared quota snapshot cache ├── tui-status.ts # prompt quota status formatting ├── types.ts # TypeScript interfaces diff --git a/lib/accounts.ts b/lib/accounts.ts index 706b2ad8..0597435f 100644 --- a/lib/accounts.ts +++ b/lib/accounts.ts @@ -225,14 +225,28 @@ export class AccountManager { family: ModelFamily, model?: string | null, options?: HybridSelectionOptions, + preferredAccountIds?: readonly string[], ): ManagedAccount | null { switch (strategy) { case "sticky": - return this.rotation.getCurrentOrNextForFamilySticky(family, model); + return this.rotation.getCurrentOrNextForFamilySticky( + family, + model, + preferredAccountIds, + ); case "round-robin": - return this.rotation.getCurrentOrNextForFamily(family, model); + return this.rotation.getCurrentOrNextForFamily( + family, + model, + preferredAccountIds, + ); default: - return this.rotation.getCurrentOrNextForFamilyHybrid(family, model, options); + return this.rotation.getCurrentOrNextForFamilyHybrid( + family, + model, + options, + preferredAccountIds, + ); } } diff --git a/lib/accounts/rotation.ts b/lib/accounts/rotation.ts index e41cb675..3b161660 100644 --- a/lib/accounts/rotation.ts +++ b/lib/accounts/rotation.ts @@ -53,12 +53,44 @@ export class AccountRotation { return getTokenTracker().hasToken(account.index, quotaKey); } + private getPreferredSelectableIndices( + accountIds: readonly string[] | undefined, + family: ModelFamily, + model?: string | null, + ): ReadonlySet | null { + if (!accountIds?.length) return null; + const preferredIds = new Set(accountIds); + const indices = this.state.accounts + .filter( + (account) => + account && + account.accountId !== undefined && + preferredIds.has(account.accountId) && + this.isSelectable(account, family, model), + ) + .map((account) => account.index); + return indices.length > 0 ? new Set(indices) : null; + } + + private isInSelectionPool( + account: ManagedAccount, + preferredIndices: ReadonlySet | null, + ): boolean { + return preferredIndices === null || preferredIndices.has(account.index); + } + getCurrentOrNextForFamily( family: ModelFamily, model?: string | null, + preferredAccountIds?: readonly string[], ): ManagedAccount | null { const count = this.state.accounts.length; if (count === 0) return null; + const preferredIndices = this.getPreferredSelectableIndices( + preferredAccountIds, + family, + model, + ); const cursor = this.state.cursorByFamily[family]; @@ -66,6 +98,7 @@ export class AccountRotation { const idx = (cursor + i) % count; const account = this.state.accounts[idx]; if (!account) continue; + if (!this.isInSelectionPool(account, preferredIndices)) continue; if (!this.isSelectable(account, family, model)) continue; this.state.cursorByFamily[family] = (idx + 1) % count; @@ -101,9 +134,15 @@ export class AccountRotation { family: ModelFamily, model?: string | null, options?: HybridSelectionOptions, + preferredAccountIds?: readonly string[], ): ManagedAccount | null { const count = this.state.accounts.length; if (count === 0) return null; + const preferredIndices = this.getPreferredSelectableIndices( + preferredAccountIds, + family, + model, + ); const currentIndex = this.state.currentAccountIndexByFamily[family]; if (currentIndex >= 0 && currentIndex < count) { @@ -111,7 +150,10 @@ export class AccountRotation { if (currentAccount) { if (currentAccount.enabled === false) { // Fall through to hybrid selection. - } else if (this.isSelectable(currentAccount, family, model)) { + } else if ( + this.isInSelectionPool(currentAccount, preferredIndices) && + this.isSelectable(currentAccount, family, model) + ) { currentAccount.lastUsed = nowMs(); return currentAccount; } @@ -126,6 +168,7 @@ export class AccountRotation { .map((account): AccountWithMetrics | null => { if (!account) return null; if (account.enabled === false) return null; + if (!this.isInSelectionPool(account, preferredIndices)) return null; return { index: account.index, isAvailable: this.isSelectable(account, family, model), @@ -172,15 +215,25 @@ export class AccountRotation { getCurrentOrNextForFamilySticky( family: ModelFamily, model?: string | null, + preferredAccountIds?: readonly string[], ): ManagedAccount | null { const count = this.state.accounts.length; if (count === 0) return null; + const preferredIndices = this.getPreferredSelectableIndices( + preferredAccountIds, + family, + model, + ); // Prefer the account we are already pinned to while it still has quota. const currentIndex = this.state.currentAccountIndexByFamily[family]; if (currentIndex >= 0 && currentIndex < count) { const currentAccount = this.state.accounts[currentIndex]; - if (currentAccount && this.isSelectable(currentAccount, family, model)) { + if ( + currentAccount && + this.isInSelectionPool(currentAccount, preferredIndices) && + this.isSelectable(currentAccount, family, model) + ) { currentAccount.lastUsed = nowMs(); return currentAccount; } @@ -191,6 +244,7 @@ export class AccountRotation { for (let idx = 0; idx < count; idx++) { const account = this.state.accounts[idx]; if (!account) continue; + if (!this.isInSelectionPool(account, preferredIndices)) continue; if (!this.isSelectable(account, family, model)) continue; this.state.currentAccountIndexByFamily[family] = idx; diff --git a/lib/config.ts b/lib/config.ts index df64826a..f482e94f 100644 --- a/lib/config.ts +++ b/lib/config.ts @@ -1,6 +1,8 @@ -import { readFileSync, existsSync } from "node:fs"; -import { join } from "node:path"; +import { readFileSync, existsSync, promises as fs } from "node:fs"; +import { dirname, join } from "node:path"; import { homedir } from "node:os"; +import { randomUUID } from "node:crypto"; +import { lock } from "proper-lockfile"; import type { PluginConfig } from "./types.js"; import { normalizeRetryBudgetValue, @@ -9,6 +11,7 @@ import { } from "./request/retry-budget.js"; import { logWarn } from "./logger.js"; import { stripEffortSuffix } from "./request/helpers/effort-suffix.js"; +import { renameWithWindowsRetry } from "./storage/atomic-write.js"; import { PluginConfigSchema, getValidationErrors, @@ -26,6 +29,18 @@ const RETRY_PROFILES = new Set(["conservative", "balanced", "aggressive"]); export type UnsupportedCodexPolicy = "strict" | "fallback"; +export type ModelAccountPoolMutation = "set" | "add" | "remove" | "clear"; + +export interface ModelAccountPoolMutationResult { + model: string; + previousAccountIds: string[]; + accountIds: string[]; + changed: boolean; + dryRun: boolean; +} + +let modelAccountPoolMutationQueue: Promise = Promise.resolve(); + /** * Default plugin configuration * CODEX_MODE is enabled by default for better Codex CLI parity @@ -139,6 +154,143 @@ export function loadPluginConfig(): PluginConfig { } } +/** + * Update one model pool while preserving every unrelated raw config key. + * Account indexes are deliberately resolved by the caller; only stable IDs + * cross this persistence boundary. + */ +export function updateModelAccountPool( + model: string, + mutation: ModelAccountPoolMutation, + accountIds: readonly string[] = [], + options: { dryRun?: boolean } = {}, +): Promise { + const pending = modelAccountPoolMutationQueue.then(async () => { + await fs.mkdir(dirname(CONFIG_PATH), { recursive: true, mode: 0o700 }); + const release = await lock(CONFIG_PATH, { + realpath: false, + stale: 10_000, + update: 2_000, + retries: { + retries: 20, + factor: 1.2, + minTimeout: 25, + maxTimeout: 200, + randomize: true, + }, + }); + try { + return await performModelAccountPoolMutation( + model, + mutation, + accountIds, + options, + ); + } finally { + await release(); + } + }); + modelAccountPoolMutationQueue = pending.then( + () => undefined, + () => undefined, + ); + return pending; +} + +async function performModelAccountPoolMutation( + model: string, + mutation: ModelAccountPoolMutation, + accountIds: readonly string[], + options: { dryRun?: boolean }, +): Promise { + const normalizedModel = model.trim().toLowerCase(); + if (!normalizedModel) throw new Error("Model is required."); + + let rawConfig: Record = {}; + try { + const content = await fs.readFile(CONFIG_PATH, "utf-8"); + const parsed = JSON.parse(stripUtf8Bom(content)) as unknown; + if (!isRecord(parsed) || Array.isArray(parsed)) { + throw new Error(`Plugin config at ${CONFIG_PATH} is not a JSON object.`); + } + rawConfig = parsed; + } catch (error) { + if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error; + } + + const poolResult = PluginConfigSchema.safeParse({ + modelAccountPools: rawConfig.modelAccountPools, + }); + if (!poolResult.success) { + throw new Error( + `Existing modelAccountPools configuration is invalid: ${getValidationErrors( + PluginConfigSchema, + { modelAccountPools: rawConfig.modelAccountPools }, + )[0] ?? "validation failed"}`, + ); + } + + const pools = { ...(poolResult.data.modelAccountPools ?? {}) }; + const matchingKeys = Object.keys(pools).filter( + (key) => key.trim().toLowerCase() === normalizedModel, + ); + const previousAccountIds = Array.from( + new Set(matchingKeys.flatMap((key) => pools[key] ?? [])), + ); + for (const key of matchingKeys) delete pools[key]; + + const normalizedAccountIds = Array.from( + new Set(accountIds.map((id) => id.trim()).filter(Boolean)), + ); + let nextAccountIds: string[]; + if (mutation === "set") { + nextAccountIds = normalizedAccountIds; + } else if (mutation === "add") { + nextAccountIds = Array.from( + new Set([...previousAccountIds, ...normalizedAccountIds]), + ); + } else if (mutation === "remove") { + const removedIds = new Set(normalizedAccountIds); + nextAccountIds = previousAccountIds.filter((id) => !removedIds.has(id)); + } else { + nextAccountIds = []; + } + + if (nextAccountIds.length > 0) pools[normalizedModel] = nextAccountIds; + const changed = + matchingKeys.length !== (nextAccountIds.length > 0 ? 1 : 0) || + previousAccountIds.length !== nextAccountIds.length || + previousAccountIds.some((id, index) => id !== nextAccountIds[index]) || + (matchingKeys[0] !== undefined && matchingKeys[0] !== normalizedModel); + + if (changed && options.dryRun !== true) { + if (Object.keys(pools).length > 0) { + rawConfig.modelAccountPools = pools; + } else { + delete rawConfig.modelAccountPools; + } + + const tempPath = `${CONFIG_PATH}.${process.pid}.${randomUUID()}.tmp`; + try { + await fs.writeFile(tempPath, `${JSON.stringify(rawConfig, null, 2)}\n`, { + encoding: "utf-8", + mode: 0o600, + }); + await renameWithWindowsRetry(tempPath, CONFIG_PATH); + } finally { + await fs.rm(tempPath, { force: true }).catch(() => undefined); + } + } + + return { + model: normalizedModel, + previousAccountIds, + accountIds: nextAccountIds, + changed, + dryRun: options.dryRun === true, + }; +} + /** * Salvage the subset of user-supplied config keys that individually pass * schema validation. Used when the top-level parse fails so callers still @@ -369,6 +521,22 @@ export function getRotationStrategy(pluginConfig: PluginConfig): RotationStrateg return "hybrid"; } +export function getModelAccountPool( + pluginConfig: PluginConfig, + model?: string | null, +): string[] { + if (!model) return []; + const normalizedModel = model.trim().toLowerCase(); + for (const [configuredModel, accountIds] of Object.entries( + pluginConfig.modelAccountPools ?? {}, + )) { + if (configuredModel.trim().toLowerCase() === normalizedModel) { + return [...new Set(accountIds.map((id) => id.trim()).filter(Boolean))]; + } + } + return []; +} + export function getFastSessionMaxInputItems(pluginConfig: PluginConfig): number { return resolveNumberSetting( "CODEX_AUTH_FAST_SESSION_MAX_INPUT_ITEMS", diff --git a/lib/runtime.ts b/lib/runtime.ts index b559dc65..7f70497b 100644 --- a/lib/runtime.ts +++ b/lib/runtime.ts @@ -84,6 +84,8 @@ export type SelectionSnapshot = { fallbackFrom: string | null; fallbackTo: string | null; fallbackReason: string | null; + accountPoolMode?: "general" | "preferred" | "general-fallback"; + configuredAccountPoolSize?: number; }; export type SerializedSelectionExplainability = { @@ -113,6 +115,8 @@ export type RoutingVisibilitySnapshot = { fallbackFrom: string | null; fallbackTo: string | null; fallbackReason: string | null; + accountPoolMode: "general" | "preferred" | "general-fallback" | null; + configuredAccountPoolSize: number; selectionExplainability: SerializedSelectionExplainability[]; }; diff --git a/lib/schemas.ts b/lib/schemas.ts index c16c6d48..6d1b45b1 100644 --- a/lib/schemas.ts +++ b/lib/schemas.ts @@ -22,6 +22,10 @@ export const PluginConfigSchema = z.object({ fastSession: z.boolean().optional(), fastSessionStrategy: z.enum(["hybrid", "always"]).optional(), rotationStrategy: z.enum(["hybrid", "sticky", "round-robin"]).optional(), + modelAccountPools: z.record( + z.string().min(1), + z.array(z.string().min(1)), + ).optional(), fastSessionMaxInputItems: z.number().min(8).max(200).optional(), retryProfile: z.enum(["conservative", "balanced", "aggressive"]).optional(), retryBudgetOverrides: z.object({ diff --git a/lib/tools/AGENTS.md b/lib/tools/AGENTS.md index 27606b50..4d2e1f26 100644 --- a/lib/tools/AGENTS.md +++ b/lib/tools/AGENTS.md @@ -1,6 +1,6 @@ # lib/tools/ -Per-tool modules for the 23 `codex-*` tools registered by the plugin. +Per-tool modules for the 24 `codex-*` tools registered by the plugin. ## Status @@ -26,6 +26,7 @@ lib/tools/ codex-next.ts codex-label.ts codex-tag.ts + codex-pool.ts codex-note.ts codex-dashboard.ts codex-health.ts diff --git a/lib/tools/codex-help.ts b/lib/tools/codex-help.ts index 373b221e..f121aa8b 100644 --- a/lib/tools/codex-help.ts +++ b/lib/tools/codex-help.ts @@ -21,7 +21,7 @@ export function createCodexHelpTool(ctx: ToolContext): ToolDefinition { .string() .optional() .describe( - "Optional topic: setup, switch, health, backup, dashboard.", + "Optional topic: setup, switch, pools, health, backup, dashboard.", ), }, async execute({ topic }) { @@ -41,6 +41,17 @@ export function createCodexHelpTool(ctx: ToolContext): ToolDefinition { "6) Start requests and monitor: codex-dashboard", ], }, + { + key: "pools", + title: "Model account pools", + lines: [ + "Show all pools: codex-pool", + 'Assign accounts: codex-pool action="set" model="gpt-5.6-sol" accounts=[7,8]', + 'Add an account: codex-pool action="add" model="gpt-5.6-sol" accounts=[9]', + 'Remove an account: codex-pool action="remove" model="gpt-5.6-sol" accounts=[7]', + 'Clear a pool: codex-pool action="clear" model="gpt-5.6-sol"', + ], + }, { key: "switch", title: "Daily account operations", diff --git a/lib/tools/codex-pool.ts b/lib/tools/codex-pool.ts new file mode 100644 index 00000000..dd840605 --- /dev/null +++ b/lib/tools/codex-pool.ts @@ -0,0 +1,254 @@ +/** `codex-pool` tool - manage model-specific account pools by account number. */ + +import { tool, type ToolDefinition } from "@opencode-ai/plugin/tool"; +import { + loadPluginConfig, + updateModelAccountPool, + type ModelAccountPoolMutation, +} from "../config.js"; +import { normalizeToolOutputFormat, renderJsonOutput } from "../runtime.js"; +import { loadAccounts, type AccountStorageV3 } from "../storage.js"; +import type { ToolContext } from "./index.js"; + +type CodexPoolAction = "status" | ModelAccountPoolMutation; + +type CodexPoolArgs = { + action?: string; + model?: string; + accounts?: number[]; + dryRun?: boolean; + format?: string; + includeSensitive?: boolean; +}; + +function normalizePoolAction(action?: string): CodexPoolAction { + if (action === undefined || action === "status") return "status"; + if ( + action === "set" || + action === "add" || + action === "remove" || + action === "clear" + ) { + return action; + } + throw new Error( + `Invalid action "${action}". Expected "status", "set", "add", "remove", or "clear".`, + ); +} + +function normalizeModel(model?: string): string | undefined { + const normalized = model?.trim().toLowerCase(); + return normalized || undefined; +} + +function resolveAccountIds( + storage: AccountStorageV3 | null, + accountNumbers: readonly number[] | undefined, +): string[] { + if (!accountNumbers?.length) { + throw new Error("At least one account number is required."); + } + if (!storage || storage.accounts.length === 0) { + throw new Error("No Codex accounts are configured."); + } + + const ids: string[] = []; + for (const accountNumber of Array.from(new Set(accountNumbers))) { + if ( + !Number.isInteger(accountNumber) || + accountNumber < 1 || + accountNumber > storage.accounts.length + ) { + throw new Error( + `Invalid account number ${accountNumber}. Expected 1-${storage.accounts.length}.`, + ); + } + const account = storage.accounts[accountNumber - 1]; + const accountId = account?.accountId?.trim(); + if (!accountId) { + throw new Error( + `Account ${accountNumber} has no stable account ID and cannot be assigned to a model pool.`, + ); + } + ids.push(accountId); + } + return ids; +} + +function buildPoolSnapshot( + model: string, + accountIds: readonly string[], + storage: AccountStorageV3 | null, + ctx: ToolContext, + includeSensitive: boolean, +): { + model: string; + configuredCount: number; + accounts: Record[]; + unresolvedCount: number; + unresolvedAccountIds?: string[]; +} { + const accounts: Record[] = []; + const unresolvedAccountIds: string[] = []; + const storedAccounts = storage?.accounts ?? []; + const maskEmail = ctx.resolveMaskEmail(); + + for (const accountId of accountIds) { + const index = storedAccounts.findIndex( + (account) => account.accountId?.trim() === accountId, + ); + const account = index >= 0 ? storedAccounts[index] : undefined; + if (!account) { + unresolvedAccountIds.push(accountId); + continue; + } + accounts.push({ + ...ctx.buildJsonAccountIdentity(index, { + includeSensitive, + account, + label: ctx.formatCommandAccountLabel(account, index, { maskEmail }), + }), + enabled: account.enabled !== false, + }); + } + + return { + model, + configuredCount: accountIds.length, + accounts, + unresolvedCount: unresolvedAccountIds.length, + ...(includeSensitive ? { unresolvedAccountIds } : {}), + }; +} + +function renderPoolStatusText( + pools: ReturnType[], + ctx: ToolContext, + storage: AccountStorageV3 | null, +): string { + if (pools.length === 0) return "No model account pools configured."; + const lines = ["Model account pools"]; + const maskEmail = ctx.resolveMaskEmail(); + for (const pool of pools) { + lines.push("", pool.model); + for (const identity of pool.accounts) { + const index = identity.zeroBasedIndex; + if (typeof index !== "number") continue; + const account = storage?.accounts[index]; + if (account) { + lines.push( + ` ${ctx.formatCommandAccountLabel(account, index, { maskEmail })}${ + account.enabled === false ? " [disabled]" : "" + }`, + ); + } + } + if (pool.unresolvedCount > 0) { + lines.push( + ` ${pool.unresolvedCount} account reference${pool.unresolvedCount === 1 ? " is" : "s are"} not present in this project`, + ); + } + } + return lines.join("\n"); +} + +export function createCodexPoolTool(ctx: ToolContext): ToolDefinition { + return tool({ + description: + "Inspect and manage model-specific account pools. Account numbers are 1-based inputs; stable account IDs are persisted.", + args: { + action: tool.schema + .string() + .optional() + .describe('"status" (default), "set", "add", "remove", or "clear".'), + model: tool.schema + .string() + .optional() + .describe("Effective model ID, such as gpt-5.6-sol."), + accounts: tool.schema + .array(tool.schema.number()) + .optional() + .describe("1-based account numbers used by set, add, and remove."), + dryRun: tool.schema + .boolean() + .optional() + .describe("Preview a mutation without changing configuration."), + format: tool.schema + .string() + .optional() + .describe('Output format: "text" (default) or "json".'), + includeSensitive: tool.schema + .boolean() + .optional() + .describe("Include stable account IDs in JSON output (default: false)."), + }, + async execute(args: CodexPoolArgs) { + const action = normalizePoolAction(args.action); + const format = normalizeToolOutputFormat(args.format); + const model = normalizeModel(args.model); + const storage = await loadAccounts(); + const includeSensitive = args.includeSensitive === true; + + if (action === "status") { + const configuredPools = loadPluginConfig().modelAccountPools ?? {}; + const entries = Object.entries(configuredPools).filter( + ([configuredModel]) => + model === undefined || configuredModel.trim().toLowerCase() === model, + ); + const pools = entries.map(([configuredModel, ids]) => + buildPoolSnapshot( + configuredModel.trim().toLowerCase(), + ids, + storage, + ctx, + includeSensitive, + ), + ); + if (format === "json") { + return renderJsonOutput({ action, pools }); + } + if (pools.length === 0 && model) { + return `No model account pool configured for ${model}.`; + } + return renderPoolStatusText(pools, ctx, storage); + } + + if (!model) throw new Error(`Model is required for action "${action}".`); + const accountIds = + action === "clear" + ? [] + : resolveAccountIds(storage, args.accounts); + const result = await updateModelAccountPool(model, action, accountIds, { + dryRun: args.dryRun, + }); + const pool = buildPoolSnapshot( + result.model, + result.accountIds, + storage, + ctx, + includeSensitive, + ); + const applied = result.changed && !result.dryRun; + const payload = { + action, + model: result.model, + changed: result.changed, + applied, + dryRun: result.dryRun, + restartRequired: applied, + previousConfiguredCount: result.previousAccountIds.length, + pool, + }; + if (format === "json") return renderJsonOutput(payload); + + const verb = result.dryRun ? "Previewed" : applied ? "Updated" : "No change to"; + const lines = [ + `${verb} model account pool: ${result.model}`, + `Previous accounts: ${result.previousAccountIds.length}`, + `Current accounts: ${result.accountIds.length}`, + ]; + if (applied) lines.push("Restart OpenCode to apply this routing change."); + return lines.join("\n"); + }, + }); +} diff --git a/lib/tools/index.ts b/lib/tools/index.ts index 4b5b72b4..3ef0644c 100644 --- a/lib/tools/index.ts +++ b/lib/tools/index.ts @@ -45,6 +45,7 @@ import { createCodexMetricsTool } from "./codex-metrics.js"; import { createCodexDoctorTool } from "./codex-doctor.js"; import { createCodexLabelTool } from "./codex-label.js"; import { createCodexTagTool } from "./codex-tag.js"; +import { createCodexPoolTool } from "./codex-pool.js"; import { createCodexNoteTool } from "./codex-note.js"; import { createCodexDashboardTool } from "./codex-dashboard.js"; import { createCodexHealthTool } from "./codex-health.js"; @@ -82,7 +83,7 @@ export type MutableRef = { current: T }; * * The factory `createTool(ctx)` returns a standard `tool({...})` * result. Keeping the surface in one type lets us evolve it without - * threading dozens of arguments through 21 call sites. + * threading dozens of arguments through 24 call sites. * * The type lists the closure state and helpers used across the current * registry. Each tool only destructures the subset it uses. @@ -238,6 +239,7 @@ export function createToolRegistry(ctx: ToolContext): CodexToolRegistry { "codex-next": createCodexNextTool(ctx), "codex-label": createCodexLabelTool(ctx), "codex-tag": createCodexTagTool(ctx), + "codex-pool": createCodexPoolTool(ctx), "codex-note": createCodexNoteTool(ctx), "codex-dashboard": createCodexDashboardTool(ctx), "codex-health": createCodexHealthTool(ctx), diff --git a/package-lock.json b/package-lock.json index 0276c3a0..31a40776 100644 --- a/package-lock.json +++ b/package-lock.json @@ -15,6 +15,7 @@ "@opentui/core": "0.2.6", "@opentui/solid": "0.2.6", "hono": "4.12.26", + "proper-lockfile": "^4.1.2", "solid-js": "^1.9.11", "web-tree-sitter": "^0.25.10", "zod": "^4.3.6" @@ -26,6 +27,7 @@ "@fast-check/vitest": "^0.2.4", "@opencode-ai/sdk": "^1.2.10", "@types/node": "^25.3.0", + "@types/proper-lockfile": "^4.1.4", "@typescript-eslint/eslint-plugin": "^8.56.0", "@typescript-eslint/parser": "^8.56.0", "@vitest/coverage-v8": "^4.0.18", @@ -2192,6 +2194,23 @@ "undici-types": "~7.19.0" } }, + "node_modules/@types/proper-lockfile": { + "version": "4.1.4", + "resolved": "https://registry.npmjs.org/@types/proper-lockfile/-/proper-lockfile-4.1.4.tgz", + "integrity": "sha512-uo2ABllncSqg9F1D4nugVl9v93RmjxF6LJzQLMLDdPaXCUIDPeOJ21Gbqi43xNKzBi/WQ0Q0dICqufzQbMjipQ==", + "dev": true, + "license": "MIT", + "dependencies": { + "@types/retry": "*" + } + }, + "node_modules/@types/retry": { + "version": "0.12.5", + "resolved": "https://registry.npmjs.org/@types/retry/-/retry-0.12.5.tgz", + "integrity": "sha512-3xSjTp3v03X/lSQLkczaN9UIEwJMoMCA1+Nb5HfbJEQWogdeQIyVtTvxPXDQjZ5zws8rFQfVfRdz03ARihPJgw==", + "dev": true, + "license": "MIT" + }, "node_modules/@typescript-eslint/eslint-plugin": { "version": "8.59.0", "resolved": "https://registry.npmjs.org/@typescript-eslint/eslint-plugin/-/eslint-plugin-8.59.0.tgz", @@ -3628,6 +3647,12 @@ "node": ">=10.13.0" } }, + "node_modules/graceful-fs": { + "version": "4.2.11", + "resolved": "https://registry.npmjs.org/graceful-fs/-/graceful-fs-4.2.11.tgz", + "integrity": "sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ==", + "license": "ISC" + }, "node_modules/has-flag": { "version": "4.0.0", "resolved": "https://registry.npmjs.org/has-flag/-/has-flag-4.0.0.tgz", @@ -4517,6 +4542,23 @@ "node": ">= 0.8.0" } }, + "node_modules/proper-lockfile": { + "version": "4.1.2", + "resolved": "https://registry.npmjs.org/proper-lockfile/-/proper-lockfile-4.1.2.tgz", + "integrity": "sha512-TjNPblN4BwAWMXU8s9AEz4JmQxnD1NNL7bNOY/AKUzyamc379FWASUhc/K1pL2noVb+XmZKLL68cjzLsiOAMaA==", + "license": "MIT", + "dependencies": { + "graceful-fs": "^4.2.4", + "retry": "^0.12.0", + "signal-exit": "^3.0.2" + } + }, + "node_modules/proper-lockfile/node_modules/signal-exit": { + "version": "3.0.7", + "resolved": "https://registry.npmjs.org/signal-exit/-/signal-exit-3.0.7.tgz", + "integrity": "sha512-wnD2ZE+l+SPC/uoS0vXeE9L1+0wuaMqKlfz9AMUo38JsyLSBWSFcHR1Rri62LZc12vLr1gb3jl7iwQhgwpAbGQ==", + "license": "ISC" + }, "node_modules/punycode": { "version": "2.3.1", "resolved": "https://registry.npmjs.org/punycode/-/punycode-2.3.1.tgz", @@ -4543,6 +4585,15 @@ ], "license": "MIT" }, + "node_modules/retry": { + "version": "0.12.0", + "resolved": "https://registry.npmjs.org/retry/-/retry-0.12.0.tgz", + "integrity": "sha512-9LkiTwjUh6rT555DtE9rTX+BKByPfrMzEAtnlEtdEwr3Nkffwiihqe2bWADg+OQRjt9gl6ICdmB/ZFDCGAtSow==", + "license": "MIT", + "engines": { + "node": ">= 4" + } + }, "node_modules/reselect": { "version": "4.1.8", "resolved": "https://registry.npmjs.org/reselect/-/reselect-4.1.8.tgz", diff --git a/package.json b/package.json index 61a20db8..d3ae9f9c 100644 --- a/package.json +++ b/package.json @@ -105,6 +105,7 @@ "@fast-check/vitest": "^0.2.4", "@opencode-ai/sdk": "^1.2.10", "@types/node": "^25.3.0", + "@types/proper-lockfile": "^4.1.4", "@typescript-eslint/eslint-plugin": "^8.56.0", "@typescript-eslint/parser": "^8.56.0", "@vitest/coverage-v8": "^4.0.18", @@ -124,6 +125,7 @@ "@opentui/core": "0.2.6", "@opentui/solid": "0.2.6", "hono": "4.12.26", + "proper-lockfile": "^4.1.2", "solid-js": "^1.9.11", "web-tree-sitter": "^0.25.10", "zod": "^4.3.6" diff --git a/test/doc-parity.test.ts b/test/doc-parity.test.ts index 59c25451..0a41c3fb 100644 --- a/test/doc-parity.test.ts +++ b/test/doc-parity.test.ts @@ -253,16 +253,24 @@ describe("runtime documentation parity", () => { ).sort(); expect(registeredTools).toEqual(toolFiles); - expect(registeredTools).toHaveLength(23); + expect(registeredTools).toHaveLength(24); const docsExpectations: Array<[string, string[]]> = [ [ "docs/development/ARCHITECTURE.md", [ - "23 OpenCode tools", + "24 OpenCode tools", + "`codex-list`, `codex-switch`, `codex-warm`", "every registered `codex-*` tool is its own file under `lib/tools/`", ], ], + [ + "docs/architecture.md", + [ + "24 per-file factories", + "`codex-list`, `codex-switch`, `codex-warm`", + ], + ], [ "docs/development/TESTING.md", [ @@ -273,7 +281,7 @@ describe("runtime documentation parity", () => { [ "lib/tools/AGENTS.md", [ - "23 `codex-*` tools", + "24 `codex-*` tools", "codex-keychain.ts", ], ], diff --git a/test/index-retry.test.ts b/test/index-retry.test.ts index 7af8c4ee..d0df0299 100644 --- a/test/index-retry.test.ts +++ b/test/index-retry.test.ts @@ -11,6 +11,7 @@ vi.mock("@opencode-ai/plugin/tool", () => { number: () => makeSchema(), boolean: () => makeSchema(), string: () => makeSchema(), + array: () => makeSchema(), }; return { tool }; diff --git a/test/index.test.ts b/test/index.test.ts index 28f3bd70..b341ce6f 100644 --- a/test/index.test.ts +++ b/test/index.test.ts @@ -15,6 +15,7 @@ vi.mock("@opencode-ai/plugin/tool", () => { number: () => makeSchema(), boolean: () => makeSchema(), string: () => makeSchema(), + array: () => makeSchema(), }; return { tool }; @@ -132,6 +133,7 @@ vi.mock("../lib/config.js", () => ({ getEmptyResponseRetryDelayMs: () => 1000, getPidOffsetEnabled: () => false, getRotationStrategy: () => "hybrid", + getModelAccountPool: vi.fn(() => []), getFetchTimeoutMs: () => 60000, getStreamStallTimeoutMs: () => 45000, getCodexTuiV2: () => false, @@ -532,6 +534,14 @@ type PluginType = { "codex-next": OptionalToolExecute<{ format?: string }>; "codex-label": ToolExecute<{ index?: number; label: string }>; "codex-tag": ToolExecute<{ index?: number; tags: string }>; + "codex-pool": OptionalToolExecute<{ + action?: string; + model?: string; + accounts?: number[]; + dryRun?: boolean; + format?: string; + includeSensitive?: boolean; + }>; "codex-note": ToolExecute<{ index?: number; note: string }>; "codex-dashboard": OptionalToolExecute<{ format?: string; includeSensitive?: boolean }>; "codex-health": OptionalToolExecute<{ format?: string; includeSensitive?: boolean }>; @@ -594,6 +604,7 @@ describe("OpenAIOAuthPlugin", () => { expect(plugin.tool["codex-next"]).toBeDefined(); expect(plugin.tool["codex-label"]).toBeDefined(); expect(plugin.tool["codex-tag"]).toBeDefined(); + expect(plugin.tool["codex-pool"]).toBeDefined(); expect(plugin.tool["codex-note"]).toBeDefined(); expect(plugin.tool["codex-dashboard"]).toBeDefined(); expect(plugin.tool["codex-health"]).toBeDefined(); @@ -2996,6 +3007,39 @@ describe("OpenAIOAuthPlugin fetch handler", () => { expect(response.status).toBe(200); }); + it.each([ + { pool: [], mode: "general", size: 0 }, + { pool: ["acc-1"], mode: "preferred", size: 1 }, + { pool: ["unavailable-account"], mode: "general-fallback", size: 1 }, + ] as const)( + "reports $mode account pool routing diagnostics", + async ({ pool, mode, size }) => { + const configModule = await import("../lib/config.js"); + vi.mocked(configModule.getModelAccountPool).mockReturnValueOnce([...pool]); + globalThis.fetch = vi.fn().mockResolvedValue( + new Response(JSON.stringify({ content: "test" }), { status: 200 }), + ); + + const { plugin, sdk } = await setupPlugin(); + const response = await sdk.fetch!("https://api.openai.com/v1/chat", { + method: "POST", + body: JSON.stringify({ model: "gpt-5.1" }), + }); + + expect(response.status).toBe(200); + const metrics = parseJsonOutput<{ + routingVisibility: { + accountPoolMode: string | null; + configuredAccountPoolSize: number; + }; + }>(await plugin.tool["codex-metrics"].execute({ format: "json" })); + expect(metrics.routingVisibility).toMatchObject({ + accountPoolMode: mode, + configuredAccountPoolSize: size, + }); + }, + ); + it("records TUI quota cache from successful Codex response headers", async () => { const previousStateDir = process.env.OPENCODE_STATE_DIR; const stateDir = await mkdtemp(join(tmpdir(), "tui-quota-index-")); diff --git a/test/model-pool-config.test.ts b/test/model-pool-config.test.ts new file mode 100644 index 00000000..2c9ecccf --- /dev/null +++ b/test/model-pool-config.test.ts @@ -0,0 +1,180 @@ +import { afterAll, beforeEach, describe, expect, it, vi } from "vitest"; +import { promises as fs } from "node:fs"; +import { join } from "node:path"; +import { lock } from "proper-lockfile"; + +const testHome = vi.hoisted( + () => `/tmp/oc-codex-model-pool-config-${process.pid}`, +); + +vi.mock("node:os", async () => { + const actual = await vi.importActual("node:os"); + return { ...actual, homedir: () => testHome }; +}); + +import { updateModelAccountPool } from "../lib/config.js"; + +const configDir = join(testHome, ".opencode"); +const configPath = join(configDir, "openai-codex-auth-config.json"); + +async function readConfig(): Promise> { + return JSON.parse(await fs.readFile(configPath, "utf-8")) as Record< + string, + unknown + >; +} + +describe("model account pool config mutation", () => { + beforeEach(async () => { + await fs.rm(testHome, { recursive: true, force: true }); + await fs.mkdir(configDir, { recursive: true }); + }); + + afterAll(async () => { + await fs.rm(testHome, { recursive: true, force: true }); + }); + + it("sets a canonical model key and preserves unrelated configuration", async () => { + await fs.writeFile( + configPath, + JSON.stringify({ + rotationStrategy: "sticky", + customFutureField: { keep: true }, + modelAccountPools: { + " GPT-5.6-SOL ": ["old-one"], + "gpt-5.6-sol": ["old-two"], + "gpt-5.6-terra": ["terra-one"], + }, + }), + ); + + const result = await updateModelAccountPool( + " GPT-5.6-SOL ", + "set", + ["new-one", "new-one", "new-two"], + ); + const saved = await readConfig(); + + expect(result).toMatchObject({ + model: "gpt-5.6-sol", + previousAccountIds: ["old-one", "old-two"], + accountIds: ["new-one", "new-two"], + changed: true, + }); + expect(saved).toMatchObject({ + rotationStrategy: "sticky", + customFutureField: { keep: true }, + modelAccountPools: { + "gpt-5.6-sol": ["new-one", "new-two"], + "gpt-5.6-terra": ["terra-one"], + }, + }); + }); + + it("adds, removes, and clears stable IDs while preserving order", async () => { + await fs.writeFile( + configPath, + JSON.stringify({ modelAccountPools: { model: ["one", "two"] } }), + ); + + await updateModelAccountPool("model", "add", ["two", "three"]); + expect((await readConfig()).modelAccountPools).toEqual({ + model: ["one", "two", "three"], + }); + + await updateModelAccountPool("model", "remove", ["two"]); + expect((await readConfig()).modelAccountPools).toEqual({ + model: ["one", "three"], + }); + + await updateModelAccountPool("model", "clear"); + expect(await readConfig()).not.toHaveProperty("modelAccountPools"); + }); + + it("serializes concurrent mutations so updates are not lost", async () => { + await fs.writeFile( + configPath, + JSON.stringify({ modelAccountPools: { model: ["one"] } }), + ); + + await Promise.all([ + updateModelAccountPool("model", "add", ["two"]), + updateModelAccountPool("model", "add", ["three"]), + ]); + + expect((await readConfig()).modelAccountPools).toEqual({ + model: ["one", "two", "three"], + }); + }); + + it("waits for a foreign lock before reading and preserves the latest config", async () => { + await fs.writeFile( + configPath, + JSON.stringify({ modelAccountPools: { model: ["one"] } }), + ); + const releaseForeignLock = await lock(configPath, { realpath: false }); + const pendingMutation = updateModelAccountPool("model", "add", ["two"]); + + await new Promise((resolve) => setTimeout(resolve, 75)); + await fs.writeFile( + configPath, + JSON.stringify({ + rotationStrategy: "sticky", + modelAccountPools: { model: ["one", "external"] }, + }), + ); + await releaseForeignLock(); + await pendingMutation; + + expect(await readConfig()).toMatchObject({ + rotationStrategy: "sticky", + modelAccountPools: { model: ["one", "external", "two"] }, + }); + }); + + it("previews a change without creating or modifying the config file", async () => { + const result = await updateModelAccountPool("model", "set", ["one"], { + dryRun: true, + }); + + expect(result).toMatchObject({ + accountIds: ["one"], + changed: true, + dryRun: true, + }); + await expect(fs.stat(configPath)).rejects.toMatchObject({ code: "ENOENT" }); + }); + + it("does not rewrite an unchanged pool", async () => { + const original = `${JSON.stringify({ + modelAccountPools: { model: ["one"] }, + })}\n\n`; + await fs.writeFile(configPath, original); + + const result = await updateModelAccountPool("model", "set", ["one"]); + + expect(result.changed).toBe(false); + expect(await fs.readFile(configPath, "utf-8")).toBe(original); + }); + + it("refuses to overwrite malformed JSON", async () => { + await fs.writeFile(configPath, "{ malformed"); + + await expect( + updateModelAccountPool("model", "set", ["one"]), + ).rejects.toThrow(); + expect(await fs.readFile(configPath, "utf-8")).toBe("{ malformed"); + const release = await lock(configPath, { realpath: false }); + await release(); + }); + + it("refuses to overwrite an invalid existing pool", async () => { + const original = JSON.stringify({ modelAccountPools: { model: [1] } }); + await fs.writeFile(configPath, original); + + await expect( + updateModelAccountPool("model", "set", ["one"]), + ).rejects.toThrow("Existing modelAccountPools configuration is invalid"); + expect(await fs.readFile(configPath, "utf-8")).toBe(original); + }); +}); diff --git a/test/rotation-strategy.test.ts b/test/rotation-strategy.test.ts index cc0dca82..7265b1c9 100644 --- a/test/rotation-strategy.test.ts +++ b/test/rotation-strategy.test.ts @@ -1,7 +1,7 @@ import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; import { AccountManager } from "../lib/accounts.js"; import { resetTrackers } from "../lib/rotation.js"; -import { getRotationStrategy } from "../lib/config.js"; +import { getModelAccountPool, getRotationStrategy } from "../lib/config.js"; import type { PluginConfig } from "../lib/types.js"; import type { AccountStorageV3 } from "../lib/storage.js"; import type { ModelFamily } from "../lib/prompts/codex.js"; @@ -30,6 +30,7 @@ function makeStorage(count: number): AccountStorageV3 { return { version: 3, accounts: Array.from({ length: count }, (_v, idx) => ({ + accountId: `account-id-${idx + 1}`, email: `account${idx + 1}@example.com`, refreshToken: `fake_refresh_token_${idx + 1}_for_testing_only`, addedAt: now - (count - idx) * 1000, @@ -82,6 +83,24 @@ describe("getRotationStrategy (#183 config)", () => { }); }); +describe("model account pool config", () => { + it("matches model names case-insensitively and removes duplicate IDs", () => { + const config = { + modelAccountPools: { + "GPT-5.6-SOL": [" account-id-2 ", "account-id-2", "account-id-3"], + }, + } as PluginConfig; + expect(getModelAccountPool(config, "gpt-5.6-sol")).toEqual([ + "account-id-2", + "account-id-3", + ]); + }); + + it("returns an empty pool for unmapped models", () => { + expect(getModelAccountPool({ modelAccountPools: {} } as PluginConfig, "gpt-5.5")).toEqual([]); + }); +}); + describe("sticky selection (#183, drain-first)", () => { let manager: AccountManager; @@ -176,4 +195,45 @@ describe("strategy dispatcher (#183)", () => { viaDispatcher?.index, ); }); + + it.each(["sticky", "round-robin", "hybrid"] as const)( + "%s prefers an assigned healthy account", + (strategy) => { + const selected = manager.getAccountForStrategy( + strategy, + FAMILY, + "gpt-5.6-sol", + undefined, + ["account-id-2"], + ); + expect(selected?.accountId).toBe("account-id-2"); + }, + ); + + it.each(["sticky", "round-robin", "hybrid"] as const)( + "%s falls back to the general pool when assigned accounts are unavailable", + (strategy) => { + manager.setAccountEnabled(1, false); + const selected = manager.getAccountForStrategy( + strategy, + FAMILY, + "gpt-5.6-sol", + undefined, + ["account-id-2"], + ); + expect(selected).not.toBeNull(); + expect(selected?.accountId).not.toBe("account-id-2"); + }, + ); + + it("falls back to the general pool when configured IDs are unknown", () => { + const selected = manager.getAccountForStrategy( + "sticky", + FAMILY, + "gpt-5.6-sol", + undefined, + ["unknown-account-id"], + ); + expect(selected?.index).toBe(0); + }); }); diff --git a/test/tools-codex-pool.test.ts b/test/tools-codex-pool.test.ts new file mode 100644 index 00000000..11a529c8 --- /dev/null +++ b/test/tools-codex-pool.test.ts @@ -0,0 +1,212 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +import type { AccountStorageV3 } from "../lib/storage.js"; +import type { ToolContext } from "../lib/tools/index.js"; +import { createCodexPoolTool } from "../lib/tools/codex-pool.js"; + +vi.mock("../lib/storage.js", () => ({ + loadAccounts: vi.fn(), +})); + +vi.mock("../lib/config.js", () => ({ + loadPluginConfig: vi.fn(), + updateModelAccountPool: vi.fn(), +})); + +import { loadPluginConfig, updateModelAccountPool } from "../lib/config.js"; +import { loadAccounts } from "../lib/storage.js"; + +const storage: AccountStorageV3 = { + version: 3, + activeIndex: 0, + accounts: [ + { + email: "one@example.com", + accountId: "account-one", + accountLabel: "Primary", + refreshToken: "refresh-one", + addedAt: 1, + lastUsed: 1, + }, + { + email: "two@example.com", + accountId: "account-two", + refreshToken: "refresh-two", + addedAt: 2, + lastUsed: 2, + enabled: false, + }, + ], +}; + +function buildCtx(): ToolContext { + return { + resolveMaskEmail: () => true, + formatCommandAccountLabel: (account, index) => + `Account ${index + 1}${account?.accountLabel ? ` (${account.accountLabel})` : ""}`, + buildJsonAccountIdentity: (index, options) => ({ + index: index + 1, + zeroBasedIndex: index, + ...(options?.includeSensitive + ? { accountId: options.account?.accountId ?? null } + : {}), + }), + } as unknown as ToolContext; +} + +describe("codex-pool tool", () => { + beforeEach(() => { + vi.mocked(loadAccounts).mockReset(); + vi.mocked(loadPluginConfig).mockReset(); + vi.mocked(updateModelAccountPool).mockReset(); + vi.mocked(loadAccounts).mockResolvedValue(storage); + vi.mocked(loadPluginConfig).mockReturnValue({ + modelAccountPools: { + "gpt-5.6-sol": ["account-one", "missing-account"], + }, + }); + }); + + it("reports configured and unresolved accounts without exposing IDs", async () => { + const output = await createCodexPoolTool(buildCtx()).execute( + { action: "status", format: "json" }, + {} as never, + ); + const parsed = JSON.parse(output as string) as { + pools: Array<{ + configuredCount: number; + accounts: Array>; + unresolvedCount: number; + unresolvedAccountIds?: string[]; + }>; + }; + + expect(parsed.pools[0]).toMatchObject({ + configuredCount: 2, + unresolvedCount: 1, + }); + expect(parsed.pools[0]?.accounts[0]).toMatchObject({ index: 1 }); + expect(parsed.pools[0]?.accounts[0]).not.toHaveProperty("accountId"); + expect(parsed.pools[0]).not.toHaveProperty("unresolvedAccountIds"); + }); + + it("includes stable IDs only when sensitive JSON is requested", async () => { + const output = await createCodexPoolTool(buildCtx()).execute( + { format: "json", includeSensitive: true }, + {} as never, + ); + const parsed = JSON.parse(output as string) as { + pools: Array<{ + accounts: Array>; + unresolvedAccountIds: string[]; + }>; + }; + + expect(parsed.pools[0]?.accounts[0]).toMatchObject({ + accountId: "account-one", + }); + expect(parsed.pools[0]?.unresolvedAccountIds).toEqual(["missing-account"]); + }); + + it("resolves unique 1-based account numbers to stable IDs", async () => { + vi.mocked(updateModelAccountPool).mockResolvedValue({ + model: "gpt-5.6-sol", + previousAccountIds: [], + accountIds: ["account-two", "account-one"], + changed: true, + dryRun: false, + }); + + const output = await createCodexPoolTool(buildCtx()).execute( + { + action: "set", + model: " GPT-5.6-SOL ", + accounts: [2, 1, 2], + }, + {} as never, + ); + + expect(updateModelAccountPool).toHaveBeenCalledWith( + "gpt-5.6-sol", + "set", + ["account-two", "account-one"], + { dryRun: undefined }, + ); + expect(output).toContain("Restart OpenCode"); + }); + + it("passes dry runs through without requesting a restart", async () => { + vi.mocked(updateModelAccountPool).mockResolvedValue({ + model: "gpt-5.6-sol", + previousAccountIds: ["account-one"], + accountIds: ["account-one", "account-two"], + changed: true, + dryRun: true, + }); + + const output = await createCodexPoolTool(buildCtx()).execute( + { + action: "add", + model: "gpt-5.6-sol", + accounts: [2], + dryRun: true, + }, + {} as never, + ); + + expect(updateModelAccountPool).toHaveBeenCalledWith( + "gpt-5.6-sol", + "add", + ["account-two"], + { dryRun: true }, + ); + expect(output).not.toContain("Restart OpenCode"); + }); + + it("clears a pool without requiring account storage", async () => { + vi.mocked(loadAccounts).mockResolvedValue(null); + vi.mocked(updateModelAccountPool).mockResolvedValue({ + model: "gpt-5.6-sol", + previousAccountIds: ["account-one"], + accountIds: [], + changed: true, + dryRun: false, + }); + + await createCodexPoolTool(buildCtx()).execute( + { action: "clear", model: "gpt-5.6-sol" }, + {} as never, + ); + + expect(updateModelAccountPool).toHaveBeenCalledWith( + "gpt-5.6-sol", + "clear", + [], + { dryRun: undefined }, + ); + }); + + it("rejects invalid account numbers before writing", async () => { + await expect( + createCodexPoolTool(buildCtx()).execute( + { action: "remove", model: "gpt-5.6-sol", accounts: [3] }, + {} as never, + ), + ).rejects.toThrow("Expected 1-2"); + expect(updateModelAccountPool).not.toHaveBeenCalled(); + }); + + it("rejects accounts that do not have a stable ID", async () => { + vi.mocked(loadAccounts).mockResolvedValue({ + ...storage, + accounts: [{ ...storage.accounts[0], accountId: undefined }], + }); + + await expect( + createCodexPoolTool(buildCtx()).execute( + { action: "set", model: "gpt-5.6-sol", accounts: [1] }, + {} as never, + ), + ).rejects.toThrow("has no stable account ID"); + }); +});