Skip to content

feat: add live probe budgets and warnings #376

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
ndycode wants to merge 1 commit into from
Closed

feat: add live probe budgets and warnings #376

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
98 changes: 98 additions & 0 deletions lib/codex-manager/commands/report.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,9 @@ interface ReportCliOptions {
json: boolean;
explain: boolean;
model: string;
maxAccounts?: number;
maxProbes?: number;
cachedOnly: boolean;
outPath?: string;
}

Expand Down Expand Up @@ -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"),
);
Expand All @@ -110,6 +116,7 @@ function parseReportArgs(args: string[]): ParsedArgsResult<ReportCliOptions> {
json: false,
explain: false,
model: "gpt-5-codex",
cachedOnly: false,
};

for (let i = 0; i < args.length; i += 1) {
Expand All @@ -127,6 +134,10 @@ function parseReportArgs(args: string[]): ParsedArgsResult<ReportCliOptions> {
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" };
Expand All @@ -140,6 +151,44 @@ function parseReportArgs(args: string[]): ParsedArgsResult<ReportCliOptions> {
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" };
Expand Down Expand Up @@ -360,16 +409,34 @@ export async function runReportCommand(
const refreshFailures = new Map<number, TokenFailure>();
const liveQuotaByIndex = new Map<number, CodexQuotaSnapshot>();
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, {
Expand Down Expand Up @@ -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;
Comment on lines 433 to +523
Copy link
Copy Markdown

@greptile-apps greptile-apps Bot Apr 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Refresh fires before --max-probes budget check

when --max-probes=N and accounts need refreshing, account N+1 executes queuedRefresh (network call + temp-file write on windows) before the probe-budget guard fires and breaks the loop. on windows, an antivirus-held temp handle can trigger EPERM on the next write — per project convention, unnecessary temp files are a real risk.

moving the maxProbes check above the hasUsableAccessToken block prevents the orphaned refresh:

// check probe budget before attempting any refresh
if (
	typeof options.maxProbes === "number" &&
	executedLiveProbes >= options.maxProbes
) {
	probeErrors.push(`live probe request budget reached (${options.maxProbes})`);
	break;
}

let probeAccessToken = account.accessToken;
let probeAccountId = account.accountId ?? extractAccountId(account.accessToken);
if (!deps.hasUsableAccessToken(account, now)) {
	// ... refresh logic
}
Prompt To Fix With AI 
This is a comment left during a code review.
Path: lib/codex-manager/commands/report.ts
Line: 433-523

Comment:
**Refresh fires before `--max-probes` budget check**

when `--max-probes=N` and accounts need refreshing, account N+1 executes `queuedRefresh` (network call + temp-file write on windows) before the probe-budget guard fires and breaks the loop. on windows, an antivirus-held temp handle can trigger EPERM on the next write — per project convention, unnecessary temp files are a real risk.

moving the `maxProbes` check above the `hasUsableAccessToken` block prevents the orphaned refresh:

```typescript
// check probe budget before attempting any refresh
if (
	typeof options.maxProbes === "number" &&
	executedLiveProbes >= options.maxProbes
) {
	probeErrors.push(`live probe request budget reached (${options.maxProbes})`);
	break;
}

let probeAccessToken = account.accessToken;
let probeAccountId = account.accountId ?? extractAccountId(account.accessToken);
if (!deps.hasUsableAccessToken(account, now)) {
	// ... refresh logic
}
```

How can I resolve this? If you propose a fix, please make it concise.

Fix in Codex

const liveQuota = await deps.fetchCodexQuotaSnapshot({
accountId: probeAccountId,
accessToken: probeAccessToken,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)`,
);
Expand Down
66 changes: 66 additions & 0 deletions test/codex-manager-report-command.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();

Expand All @@ -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<typeof vi.fn>).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<typeof vi.fn>).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 () => {
Expand Down