Skip to content

Commit 7fe8e65

Browse files
ndycodendycode
committed
feat: add live probe budgets and warnings
1 parent b417887 commit 7fe8e65

2 files changed

Lines changed: 164 additions & 0 deletions

File tree

‎lib/codex-manager/commands/report.ts‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,9 @@ interface ReportCliOptions {
3535
json: boolean;
3636
explain: boolean;
3737
model: string;
38+
maxAccounts?: number;
39+
maxProbes?: number;
40+
cachedOnly: boolean;
3841
outPath?: string;
3942
}
4043

@@ -99,6 +102,9 @@ function printReportUsage(logInfo: (message: string) => void): void {
99102
" --json, -j Print machine-readable JSON output",
100103
" --explain Print per-account reasoning in text mode",
101104
" --model, -m Probe model for live mode (default: gpt-5-codex)",
105+
" --max-accounts N Limit how many enabled accounts live mode can consider",
106+
" --max-probes N Limit how many live quota probes can run",
107+
" --cached-only Skip refreshes and only use already-usable access tokens",
102108
" --out Write JSON report to a file path",
103109
].join("\n"),
104110
);
@@ -110,6 +116,7 @@ function parseReportArgs(args: string[]): ParsedArgsResult<ReportCliOptions> {
110116
json: false,
111117
explain: false,
112118
model: "gpt-5-codex",
119+
cachedOnly: false,
113120
};
114121

115122
for (let i = 0; i < args.length; i += 1) {
@@ -127,6 +134,10 @@ function parseReportArgs(args: string[]): ParsedArgsResult<ReportCliOptions> {
127134
options.explain = true;
128135
continue;
129136
}
137+
if (arg === "--cached-only") {
138+
options.cachedOnly = true;
139+
continue;
140+
}
130141
if (arg === "--model" || arg === "-m") {
131142
const value = args[i + 1];
132143
if (!value) return { ok: false, message: "Missing value for --model" };
@@ -140,6 +151,44 @@ function parseReportArgs(args: string[]): ParsedArgsResult<ReportCliOptions> {
140151
options.model = value;
141152
continue;
142153
}
154+
if (arg === "--max-accounts") {
155+
const value = args[i + 1];
156+
if (!value) return { ok: false, message: "Missing value for --max-accounts" };
157+
const parsed = Number.parseInt(value, 10);
158+
if (!Number.isFinite(parsed) || parsed < 1) {
159+
return { ok: false, message: "--max-accounts must be a positive integer" };
160+
}
161+
options.maxAccounts = parsed;
162+
i += 1;
163+
continue;
164+
}
165+
if (arg.startsWith("--max-accounts=")) {
166+
const parsed = Number.parseInt(arg.slice("--max-accounts=".length), 10);
167+
if (!Number.isFinite(parsed) || parsed < 1) {
168+
return { ok: false, message: "--max-accounts must be a positive integer" };
169+
}
170+
options.maxAccounts = parsed;
171+
continue;
172+
}
173+
if (arg === "--max-probes") {
174+
const value = args[i + 1];
175+
if (!value) return { ok: false, message: "Missing value for --max-probes" };
176+
const parsed = Number.parseInt(value, 10);
177+
if (!Number.isFinite(parsed) || parsed < 1) {
178+
return { ok: false, message: "--max-probes must be a positive integer" };
179+
}
180+
options.maxProbes = parsed;
181+
i += 1;
182+
continue;
183+
}
184+
if (arg.startsWith("--max-probes=")) {
185+
const parsed = Number.parseInt(arg.slice("--max-probes=".length), 10);
186+
if (!Number.isFinite(parsed) || parsed < 1) {
187+
return { ok: false, message: "--max-probes must be a positive integer" };
188+
}
189+
options.maxProbes = parsed;
190+
continue;
191+
}
143192
if (arg === "--out") {
144193
const value = args[i + 1];
145194
if (!value) return { ok: false, message: "Missing value for --out" };
@@ -242,16 +291,34 @@ export async function runReportCommand(
242291
const refreshFailures = new Map<number, TokenFailure>();
243292
const liveQuotaByIndex = new Map<number, CodexQuotaSnapshot>();
244293
const probeErrors: string[] = [];
294+
let consideredLiveAccounts = 0;
295+
let executedLiveProbes = 0;
245296

246297
if (storage && options.live) {
247298
for (let i = 0; i < storage.accounts.length; i += 1) {
299+
if (
300+
typeof options.maxAccounts === "number" &&
301+
consideredLiveAccounts >= options.maxAccounts
302+
) {
303+
probeErrors.push(
304+
`live probe account budget reached (${options.maxAccounts})`,
305+
);
306+
break;
307+
}
248308
const account = storage.accounts[i];
249309
if (!account || account.enabled === false) continue;
310+
consideredLiveAccounts += 1;
250311

251312
let probeAccessToken = account.accessToken;
252313
let probeAccountId =
253314
account.accountId ?? extractAccountId(account.accessToken);
254315
if (!deps.hasUsableAccessToken(account, now)) {
316+
if (options.cachedOnly) {
317+
probeErrors.push(
318+
`${formatAccountLabel(account, i)}: skipped refresh because --cached-only is enabled`,
319+
);
320+
continue;
321+
}
255322
const refreshResult = await deps.queuedRefresh(account.refreshToken);
256323
if (refreshResult.type !== "success") {
257324
refreshFailures.set(i, {
@@ -327,8 +394,16 @@ export async function runReportCommand(
327394
);
328395
continue;
329396
}
397+
if (
398+
typeof options.maxProbes === "number" &&
399+
executedLiveProbes >= options.maxProbes
400+
) {
401+
probeErrors.push(`live probe request budget reached (${options.maxProbes})`);
402+
break;
403+
}
330404

331405
try {
406+
executedLiveProbes += 1;
332407
const liveQuota = await deps.fetchCodexQuotaSnapshot({
333408
accountId: probeAccountId,
334409
accessToken: probeAccessToken,
@@ -389,6 +464,13 @@ export async function runReportCommand(
389464
capabilities: modelInspection.capabilities,
390465
},
391466
liveProbe: options.live,
467+
liveProbeBudget: {
468+
cachedOnly: options.cachedOnly,
469+
maxAccounts: options.maxAccounts ?? null,
470+
maxProbes: options.maxProbes ?? null,
471+
consideredAccounts: consideredLiveAccounts,
472+
executedProbes: executedLiveProbes,
473+
},
392474
accounts: {
393475
total: accountCount,
394476
enabled: enabledCount,
@@ -427,6 +509,22 @@ export async function runReportCommand(
427509
logInfo(`Report generated at ${report.generatedAt}`);
428510
logInfo(`Storage: ${report.storagePath}`);
429511
logInfo(`Model: ${formatModelInspection(modelInspection)}`);
512+
if (options.live) {
513+
const budgetParts = [
514+
`considered ${consideredLiveAccounts} account(s)`,
515+
`executed ${executedLiveProbes} probe(s)`,
516+
];
517+
if (typeof options.maxAccounts === "number") {
518+
budgetParts.push(`max-accounts ${options.maxAccounts}`);
519+
}
520+
if (typeof options.maxProbes === "number") {
521+
budgetParts.push(`max-probes ${options.maxProbes}`);
522+
}
523+
if (options.cachedOnly) {
524+
budgetParts.push("cached-only");
525+
}
526+
logInfo(`Live probe budget: ${budgetParts.join(", ")}`);
527+
}
430528
logInfo(
431529
`Accounts: ${report.accounts.total} total (${report.accounts.enabled} enabled, ${report.accounts.disabled} disabled, ${report.accounts.coolingDown} cooling, ${report.accounts.rateLimited} rate-limited)`,
432530
);

‎test/codex-manager-report-command.test.ts‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,17 @@ describe("runReportCommand", () => {
8181
expect(deps.logError).toHaveBeenCalledWith("Unknown option: --bogus");
8282
});
8383

84+
it("rejects invalid live probe budget values", async () => {
85+
const deps = createDeps();
86+
87+
const result = await runReportCommand(["--max-probes", "0"], deps);
88+
89+
expect(result).toBe(1);
90+
expect(deps.logError).toHaveBeenCalledWith(
91+
"--max-probes must be a positive integer",
92+
);
93+
});
94+
8495
it("writes json report output when requested", async () => {
8596
const deps = createDeps();
8697

@@ -97,6 +108,61 @@ describe("runReportCommand", () => {
97108
expect(deps.logInfo).toHaveBeenCalledWith(
98109
expect.stringContaining('"forecast"'),
99110
);
111+
expect(deps.logInfo).toHaveBeenCalledWith(
112+
expect.stringContaining('"liveProbeBudget"'),
113+
);
114+
});
115+
116+
it("respects live probe account and probe budgets", async () => {
117+
const deps = createDeps({
118+
loadAccounts: vi.fn(async () =>
119+
createStorage([
120+
{ email: "one@example.com", refreshToken: "r1", accessToken: "a1", accountId: "acct-1", expiresAt: 5_000, addedAt: 1, lastUsed: 1, enabled: true },
121+
{ email: "two@example.com", refreshToken: "r2", accessToken: "a2", accountId: "acct-2", expiresAt: 5_000, addedAt: 2, lastUsed: 2, enabled: true },
122+
{ email: "three@example.com", refreshToken: "r3", accessToken: "a3", accountId: "acct-3", expiresAt: 5_000, addedAt: 3, lastUsed: 3, enabled: true },
123+
]),
124+
),
125+
hasUsableAccessToken: vi.fn(() => true),
126+
});
127+
128+
const result = await runReportCommand(
129+
["--live", "--json", "--max-accounts", "2", "--max-probes", "1"],
130+
deps,
131+
);
132+
133+
expect(result).toBe(0);
134+
expect(deps.fetchCodexQuotaSnapshot).toHaveBeenCalledTimes(1);
135+
const jsonOutput = JSON.parse(
136+
(deps.logInfo as ReturnType<typeof vi.fn>).mock.calls.at(-1)?.[0] ?? "{}",
137+
) as { liveProbeBudget: { consideredAccounts: number; executedProbes: number }; forecast: { probeErrors: string[] } };
138+
expect(jsonOutput.liveProbeBudget).toEqual(
139+
expect.objectContaining({ consideredAccounts: 2, executedProbes: 1 }),
140+
);
141+
expect(jsonOutput.forecast.probeErrors).toEqual(
142+
expect.arrayContaining([
143+
expect.stringContaining("live probe request budget reached (1)"),
144+
]),
145+
);
146+
});
147+
148+
it("skips refreshes in cached-only live mode", async () => {
149+
const deps = createDeps({
150+
hasUsableAccessToken: vi.fn(() => false),
151+
});
152+
153+
const result = await runReportCommand(["--live", "--json", "--cached-only"], deps);
154+
155+
expect(result).toBe(0);
156+
expect(deps.queuedRefresh).not.toHaveBeenCalled();
157+
expect(deps.fetchCodexQuotaSnapshot).not.toHaveBeenCalled();
158+
const jsonOutput = JSON.parse(
159+
(deps.logInfo as ReturnType<typeof vi.fn>).mock.calls.at(-1)?.[0] ?? "{}",
160+
) as { forecast: { probeErrors: string[] } };
161+
expect(jsonOutput.forecast.probeErrors).toEqual(
162+
expect.arrayContaining([
163+
expect.stringContaining("skipped refresh because --cached-only is enabled"),
164+
]),
165+
);
100166
});
101167

102168
it("covers live probe refresh failures, missing account ids, and probe errors", async () => {

0 commit comments

Comments
 (0)