diff --git a/lib/codex-manager/commands/report.ts b/lib/codex-manager/commands/report.ts index a77a2c84..7f7d2e7b 100644 --- a/lib/codex-manager/commands/report.ts +++ b/lib/codex-manager/commands/report.ts @@ -35,6 +35,9 @@ interface ReportCliOptions { json: boolean; explain: boolean; model: string; + maxAccounts?: number; + maxProbes?: number; + cachedOnly: boolean; outPath?: string; } @@ -99,6 +102,9 @@ function printReportUsage(logInfo: (message: string) => void): void { " --json, -j Print machine-readable JSON output", " --explain Print per-account reasoning in text mode", " --model, -m Probe model for live mode (default: gpt-5-codex)", + " --max-accounts N Limit how many enabled accounts live mode can consider", + " --max-probes N Limit how many live quota probes can run", + " --cached-only Skip refreshes and only use already-usable access tokens", " --out Write JSON report to a file path", ].join("\n"), ); @@ -110,6 +116,7 @@ function parseReportArgs(args: string[]): ParsedArgsResult { json: false, explain: false, model: "gpt-5-codex", + cachedOnly: false, }; for (let i = 0; i < args.length; i += 1) { @@ -127,6 +134,10 @@ function parseReportArgs(args: string[]): ParsedArgsResult { options.explain = true; continue; } + if (arg === "--cached-only") { + options.cachedOnly = true; + continue; + } if (arg === "--model" || arg === "-m") { const value = args[i + 1]; if (!value) return { ok: false, message: "Missing value for --model" }; @@ -140,6 +151,44 @@ function parseReportArgs(args: string[]): ParsedArgsResult { options.model = value; continue; } + if (arg === "--max-accounts") { + const value = args[i + 1]; + if (!value) return { ok: false, message: "Missing value for --max-accounts" }; + const parsed = Number.parseInt(value, 10); + if (!Number.isFinite(parsed) || parsed < 1) { + return { ok: false, message: "--max-accounts must be a positive integer" }; + } + options.maxAccounts = parsed; + i += 1; + continue; + } + if (arg.startsWith("--max-accounts=")) { + const parsed = Number.parseInt(arg.slice("--max-accounts=".length), 10); + if (!Number.isFinite(parsed) || parsed < 1) { + return { ok: false, message: "--max-accounts must be a positive integer" }; + } + options.maxAccounts = parsed; + continue; + } + if (arg === "--max-probes") { + const value = args[i + 1]; + if (!value) return { ok: false, message: "Missing value for --max-probes" }; + const parsed = Number.parseInt(value, 10); + if (!Number.isFinite(parsed) || parsed < 1) { + return { ok: false, message: "--max-probes must be a positive integer" }; + } + options.maxProbes = parsed; + i += 1; + continue; + } + if (arg.startsWith("--max-probes=")) { + const parsed = Number.parseInt(arg.slice("--max-probes=".length), 10); + if (!Number.isFinite(parsed) || parsed < 1) { + return { ok: false, message: "--max-probes must be a positive integer" }; + } + options.maxProbes = parsed; + continue; + } if (arg === "--out") { const value = args[i + 1]; if (!value) return { ok: false, message: "Missing value for --out" }; @@ -360,16 +409,34 @@ export async function runReportCommand( const refreshFailures = new Map(); const liveQuotaByIndex = new Map(); const probeErrors: string[] = []; + let consideredLiveAccounts = 0; + let executedLiveProbes = 0; if (storage && options.live) { for (let i = 0; i < storage.accounts.length; i += 1) { + if ( + typeof options.maxAccounts === "number" && + consideredLiveAccounts >= options.maxAccounts + ) { + probeErrors.push( + `live probe account budget reached (${options.maxAccounts})`, + ); + break; + } const account = storage.accounts[i]; if (!account || account.enabled === false) continue; + consideredLiveAccounts += 1; let probeAccessToken = account.accessToken; let probeAccountId = account.accountId ?? extractAccountId(account.accessToken); if (!deps.hasUsableAccessToken(account, now)) { + if (options.cachedOnly) { + probeErrors.push( + `${formatAccountLabel(account, i)}: skipped refresh because --cached-only is enabled`, + ); + continue; + } const refreshResult = await deps.queuedRefresh(account.refreshToken); if (refreshResult.type !== "success") { refreshFailures.set(i, { @@ -444,8 +511,16 @@ export async function runReportCommand( ); continue; } + if ( + typeof options.maxProbes === "number" && + executedLiveProbes >= options.maxProbes + ) { + probeErrors.push(`live probe request budget reached (${options.maxProbes})`); + break; + } try { + executedLiveProbes += 1; const liveQuota = await deps.fetchCodexQuotaSnapshot({ accountId: probeAccountId, accessToken: probeAccessToken, @@ -506,6 +581,13 @@ export async function runReportCommand( capabilities: modelInspection.capabilities, }, liveProbe: options.live, + liveProbeBudget: { + cachedOnly: options.cachedOnly, + maxAccounts: options.maxAccounts ?? null, + maxProbes: options.maxProbes ?? null, + consideredAccounts: consideredLiveAccounts, + executedProbes: executedLiveProbes, + }, accounts: { total: accountCount, enabled: enabledCount, @@ -543,6 +625,22 @@ export async function runReportCommand( logInfo(`Report generated at ${report.generatedAt}`); logInfo(`Storage: ${report.storagePath}`); logInfo(`Model: ${formatModelInspection(modelInspection)}`); + if (options.live) { + const budgetParts = [ + `considered ${consideredLiveAccounts} account(s)`, + `executed ${executedLiveProbes} probe(s)`, + ]; + if (typeof options.maxAccounts === "number") { + budgetParts.push(`max-accounts ${options.maxAccounts}`); + } + if (typeof options.maxProbes === "number") { + budgetParts.push(`max-probes ${options.maxProbes}`); + } + if (options.cachedOnly) { + budgetParts.push("cached-only"); + } + logInfo(`Live probe budget: ${budgetParts.join(", ")}`); + } logInfo( `Accounts: ${report.accounts.total} total (${report.accounts.enabled} enabled, ${report.accounts.disabled} disabled, ${report.accounts.coolingDown} cooling, ${report.accounts.rateLimited} rate-limited)`, ); diff --git a/test/codex-manager-report-command.test.ts b/test/codex-manager-report-command.test.ts index 1eac66d7..7f48ad68 100644 --- a/test/codex-manager-report-command.test.ts +++ b/test/codex-manager-report-command.test.ts @@ -81,6 +81,17 @@ describe("runReportCommand", () => { expect(deps.logError).toHaveBeenCalledWith("Unknown option: --bogus"); }); + it("rejects invalid live probe budget values", async () => { + const deps = createDeps(); + + const result = await runReportCommand(["--max-probes", "0"], deps); + + expect(result).toBe(1); + expect(deps.logError).toHaveBeenCalledWith( + "--max-probes must be a positive integer", + ); + }); + it("writes json report output when requested", async () => { const deps = createDeps(); @@ -97,6 +108,61 @@ describe("runReportCommand", () => { expect(deps.logInfo).toHaveBeenCalledWith( expect.stringContaining('"forecast"'), ); + expect(deps.logInfo).toHaveBeenCalledWith( + expect.stringContaining('"liveProbeBudget"'), + ); + }); + + it("respects live probe account and probe budgets", async () => { + const deps = createDeps({ + loadAccounts: vi.fn(async () => + createStorage([ + { email: "one@example.com", refreshToken: "r1", accessToken: "a1", accountId: "acct-1", expiresAt: 5_000, addedAt: 1, lastUsed: 1, enabled: true }, + { email: "two@example.com", refreshToken: "r2", accessToken: "a2", accountId: "acct-2", expiresAt: 5_000, addedAt: 2, lastUsed: 2, enabled: true }, + { email: "three@example.com", refreshToken: "r3", accessToken: "a3", accountId: "acct-3", expiresAt: 5_000, addedAt: 3, lastUsed: 3, enabled: true }, + ]), + ), + hasUsableAccessToken: vi.fn(() => true), + }); + + const result = await runReportCommand( + ["--live", "--json", "--max-accounts", "2", "--max-probes", "1"], + deps, + ); + + expect(result).toBe(0); + expect(deps.fetchCodexQuotaSnapshot).toHaveBeenCalledTimes(1); + const jsonOutput = JSON.parse( + (deps.logInfo as ReturnType).mock.calls.at(-1)?.[0] ?? "{}", + ) as { liveProbeBudget: { consideredAccounts: number; executedProbes: number }; forecast: { probeErrors: string[] } }; + expect(jsonOutput.liveProbeBudget).toEqual( + expect.objectContaining({ consideredAccounts: 2, executedProbes: 1 }), + ); + expect(jsonOutput.forecast.probeErrors).toEqual( + expect.arrayContaining([ + expect.stringContaining("live probe request budget reached (1)"), + ]), + ); + }); + + it("skips refreshes in cached-only live mode", async () => { + const deps = createDeps({ + hasUsableAccessToken: vi.fn(() => false), + }); + + const result = await runReportCommand(["--live", "--json", "--cached-only"], deps); + + expect(result).toBe(0); + expect(deps.queuedRefresh).not.toHaveBeenCalled(); + expect(deps.fetchCodexQuotaSnapshot).not.toHaveBeenCalled(); + const jsonOutput = JSON.parse( + (deps.logInfo as ReturnType).mock.calls.at(-1)?.[0] ?? "{}", + ) as { forecast: { probeErrors: string[] } }; + expect(jsonOutput.forecast.probeErrors).toEqual( + expect.arrayContaining([ + expect.stringContaining("skipped refresh because --cached-only is enabled"), + ]), + ); }); it("covers live probe refresh failures, missing account ids, and probe errors", async () => {