feat: centralize AI monthly budget guard#73
Conversation
🤖 Codex PR Review🚫 Merge blocked: 3 serious issue(s) found in high-risk files ⚖️ Codex Review Arbitration🚫 block: At least two current blocking findings remain valid against the PR diff. In 🚫 Blocking IssuesThese issues must be fixed before this PR can be merged: 1. 🟠 [HIGH] Logic in
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cbce3c594f
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| snapshot = ( | ||
| self._anthropic_account_snapshot() | ||
| if provider == "anthropic" | ||
| else self._openai_account_snapshot() | ||
| ) |
There was a problem hiding this comment.
Use month-to-date usage for monthly budget checks
When budget_guard is enabled, this rolling account snapshot is used as the observed usage for a monthly hard-limit decision, but the OpenAI and Anthropic readers default to a 7-day window. In a monthly billing period, spend from earlier than that window is ignored, so a repo can pass the guard after already consuming the monthly budget earlier in the month. Request or validate a snapshot that matches the decision billing period before allowing execution.
Useful? React with 👍 / 👎.
| review_quota = get_quota_manager().check( | ||
| review_repo, | ||
| reviewer_model, | ||
| req.prompt, | ||
| task_class="review", | ||
| budget_guard=_budget_gate_enabled(), | ||
| ) |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c14d0b2a95
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| get_quota_manager().record( | ||
| review_repo, | ||
| reviewer_result.model, | ||
| req.prompt, | ||
| reviewer_result.output if reviewer_result.success else "", | ||
| ) |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| f"{provider}:{provider_scope}:{repo}", | ||
| f"{provider}:{provider_scope}", | ||
| provider, | ||
| "default", |
There was a problem hiding this comment.
Resolve nested budgets before default fallback
When monthly_budgets contains both a global default entry and a provider-specific nested entry such as openai -> project -> repo -> task, this candidate returns the default before the nested traversal below runs. That ignores the more specific repo/task budget and can enforce the wrong monthly hard limit for configured repos; the default fallback should come after the nested provider lookup.
Useful? React with 👍 / 👎.
| hard_limit = _number(decision.get("hard_limit"), float("inf")) | ||
| if hard_limit != float("inf") and ( | ||
| observed + stored_settled + stored_reserved + requested > hard_limit + 1e-12 | ||
| or observed + aggregate_settled + aggregate_reserved + requested > hard_limit + 1e-12 |
There was a problem hiding this comment.
Don’t cap all repos at one repo’s API budget
For API-provider decisions, aggregate_scope is shared across repos for the same provider/project/task, but this reservation check compares that aggregate spend to the current repo’s hard_limit. With two repos each configured for a $10 monthly review budget, one $6 reservation leaves the second repo’s $6 reservation failing here even though both repos are under their own limits. Use a separate provider-project cap for aggregate enforcement, or keep repo budgets isolated from the aggregate reservation check.
Useful? React with 👍 / 👎.
|
|
||
| budget_guard_instance = get_ai_budget_guard() | ||
| snapshot = ( | ||
| self._anthropic_account_snapshot() |
There was a problem hiding this comment.
Fetch costs for scoped Anthropic budgets
When CODEX_AUDIT_SERVICE_ANTHROPIC_ADMIN_WORKSPACE_IDS or CODEX_AUDIT_SERVICE_ANTHROPIC_ADMIN_API_KEY_IDS is configured, read_anthropic_admin_usage() returns filtered message usage but leaves costs as None, so passing that snapshot into the budget guard makes _usage() report usage_cost_missing and blocks every Anthropic request under budget_guard for scoped deployments. Fetch a cost snapshot that matches the Anthropic scope, or avoid using the filtered usage-only snapshot for hard-limit decisions.
Useful? React with 👍 / 👎.
| CODEX_PR_REVIEW_API_FALLBACK_DEFAULT: ${{ vars.CODEX_PR_REVIEW_API_FALLBACK_ENABLED || 'true' }} | ||
| CODEX_PR_REVIEW_DIRECT_API_PRIMARY_DEFAULT: ${{ vars.CODEX_PR_REVIEW_DIRECT_API_PRIMARY_ENABLED || 'true' }} | ||
| CODEX_PR_REVIEW_API_FALLBACK_DEFAULT: ${{ vars.CODEX_PR_REVIEW_API_FALLBACK_ENABLED || 'false' }} | ||
| CODEX_PR_REVIEW_DIRECT_API_PRIMARY_DEFAULT: ${{ vars.CODEX_PR_REVIEW_DIRECT_API_PRIMARY_ENABLED || 'false' }} |
There was a problem hiding this comment.
Preserve direct API primary for self-review
For AIAuditBridge’s own pull_request_target runs, CODEX_PR_REVIEW_REUSABLE_CALL is false, so the shell resolver uses this default rather than the workflow_call input. If the repo has no CODEX_PR_REVIEW_DIRECT_API_PRIMARY_ENABLED variable and no CODEX_AUDIT_SERVICE_URL—the legacy direct-API setup with only OpenAI/Anthropic secrets—this exports CODEX_PR_REVIEW_DIRECT_API_PRIMARY_ENABLED=false, causing run_codex_review_with_fallback() to fail with no backend instead of using the available API keys. Keep this default true for self-review while leaving reusable callers opt-in.
Useful? React with 👍 / 👎.
| if reservation_id: | ||
| payload["_budget_reservation_id"] = reservation_id | ||
| try: | ||
| job = _submit_job(claims, payload) |
There was a problem hiding this comment.
Release async reservations when jobs dedupe
For /v1/ai/execute/jobs, this call happens after a Codex budget reservation is put in payload, but _submit_job() can return an existing active job from the dedupe branch without creating a new job that owns that reservation. In that duplicate-request scenario no worker ever settles or releases the newly reserved ID, so the shared budget ledger keeps phantom Codex headroom consumed until manual cleanup/reset.
Useful? React with 👍 / 👎.
| if any( | ||
| (parsed := _reset_timestamp(value)) is not None and parsed <= self._clock() | ||
| for value in reset_at |
There was a problem hiding this comment.
Keep Codex reservations until both windows reset
Codex calls consume both primary and secondary rate-limit windows, so when only the shorter window’s resets_at has passed while the longer window is still active, this any() branch clears all Codex reservations and settled usage anyway. That lets new work reuse secondary-window headroom that was already consumed, weakening the protected reserve for accounts with staggered reset windows; clear the ledger only for windows that have all reset or track reservations per window.
Useful? React with 👍 / 👎.
| "[quota_or_capacity_failure]" in message | ||
| or "usage limit" in message | ||
| or "daily budget exceeded" in message | ||
| or "deferred_budget" in message |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| snapshot = ( | ||
| self._anthropic_account_snapshot() | ||
| if provider == "anthropic" | ||
| else self._openai_account_snapshot() |
There was a problem hiding this comment.
Use scoped OpenAI cost snapshots for project budgets
When CODEX_AUDIT_SERVICE_OPENAI_PROJECT_SCOPE is set, this path enforces a project-scoped budget but still feeds _openai_account_snapshot() into the guard; inspecting read_openai_admin_usage() shows its cost request populates organization_costs without project/api-key filters. A project can therefore be blocked by unrelated organization spend even though its own monthly budget has capacity, so the snapshot used for this scoped decision must include matching scoped costs.
Useful? React with 👍 / 👎.
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
43b9e31 to
b024478
Compare
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 12b24bc616
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self._save_records_locked() | ||
| self._record_failures.discard(repo or "unknown") |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| if codex_quota is not None: | ||
| estimated_total += float(codex_quota.get("cost_estimate_usd") or 0.0) |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| reasoning_effort = _resolve_codex_reasoning_effort(payload, str(payload.get("task") or TASK_EXECUTE)) | ||
| # Persist a conservative possible-dispatch marker before invoking the | ||
| # external process; crash recovery must not refund ambiguous spend. | ||
| job["dispatch_started"] = True |
There was a problem hiding this comment.
Mark Codex dispatch only after the process starts
If the Codex command is unavailable or _codex_command() raises before subprocess.run() is reached, this pre-call marker makes the exception path settle the budget reservation as spent even though no Codex request consumed rate-limit headroom. In a misconfigured service host, each failed job therefore burns the protected Codex budget until reset instead of releasing the reservation; only set this after the process has actually been launched, or have the adapter report pre-dispatch failures distinctly.
Useful? React with 👍 / 👎.
| matched_result = next( | ||
| (result for result in remaining_results if result.model == reviewer_model), | ||
| None, |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
Co-Authored-By: Codex <noreply@openai.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: bba62d060b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| get_ai_budget_guard().settle( | ||
| reservation_id, | ||
| float(qr.get("cost_estimate_usd") or 0.0), | ||
| ) |
There was a problem hiding this comment.
Handle failed reservation settlement before returning
When budget_guard is enabled and the analyze provider call has already dispatched, settle() can still return False (for example, a transient SQLite/OSError in the shared ledger). This path ignores that return value and proceeds to return the model result, but no persisted job state retains the reservation id for retry, so the ledger can keep the amount in reserved and make later requests see phantom monthly budget consumption until reset. Check the result and persist/retry cleanup or fail closed, as the async job helper does.
Useful? React with 👍 / 👎.
| get_quota_manager().check( | ||
| review_repo, | ||
| reviewer_model, | ||
| req.prompt, | ||
| estimated_output_tokens=DEFAULT_MAX_TOKENS, | ||
| ).get("cost_estimate_usd") or 0.0 |
There was a problem hiding this comment.
Resolve reviewer models for the combined daily precheck
When /v1/ai/review uses a model override or alias that resolve_model() maps to a priced model (for example, an unknown string falling back to Anthropic), this estimate calls check() without budget_guard, so the cost is computed from the raw string and can fall back to the cheaper default estimator. The later per-reviewer checks use the resolved model but do not reserve daily budget, so two reviewers can both run when the remaining daily budget is between this underestimated total and the real total, exceeding the API-key daily cap. Use the same resolved-model cost path for this combined precheck.
Useful? React with 👍 / 👎.
| if reset_event is not None and reset_event <= self._clock() and reset_event > self._codex_reset_events.get(reset_key, 0.0): | ||
| self._codex_reset_events[reset_key] = reset_event |
There was a problem hiding this comment.
Clear Codex reservations after missed reset windows
This reset cleanup only runs when the snapshot's current resets_at is already in the past. If no request arrives during the short period where a cached pre-reset snapshot is still served, the next fresh Codex snapshot normally points at the next future reset, so settled reservations from the previous window are never cleared and later low-volume accounts can remain deferred despite restored rate-limit headroom. Track reset advancement or reservation creation times instead of requiring the current snapshot reset time to be past.
Useful? React with 👍 / 👎.
Co-Authored-By: Codex <noreply@openai.com>
Co-Authored-By: Codex <noreply@openai.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3aa075a3c1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if result.success or bool(getattr(result, "dispatch_started", False)): | ||
| try: | ||
| quota.record(source_repo, resolved_model, req.prompt, result.output if result.success else "") |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| if payload.get("_budget_cleanup_pending") or payload.get("_budget_reservation_id"): | ||
| continue |
There was a problem hiding this comment.
Mark expired uncertain jobs terminal before preserving them
When an expired queued/running async job still has a reservation that finalizes to pending_uncertain, this continue preserves the job file without changing its status. _active_job_count() still counts queued/running files, so enough expired ambiguous jobs can permanently consume all active-job slots and make new submissions fail with the max-active error. Mark the expired job failed/terminal before keeping it for budget reconciliation.
Useful? React with 👍 / 👎.
Co-Authored-By: Codex <noreply@openai.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8f9f9da091
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if not review_result.success and not bool(getattr(review_result, "dispatch_started", False)): | ||
| continue |
This comment was marked as off-topic.
This comment was marked as off-topic.
Sorry, something went wrong.
| for reservation_id, estimated_cost, reviewer_index in review_reservations: | ||
| matched_result = llm_results[reviewer_index] if reviewer_index < len(llm_results) else None | ||
| if matched_result is not None and bool(getattr(matched_result, "dispatch_started", False)): | ||
| get_ai_budget_guard().settle(reservation_id, estimated_cost) |
There was a problem hiding this comment.
Check review reservation settlement before returning
When /v1/ai/review has already dispatched an API reviewer, settle() can still return False (for example, a transient SQLite/OSError in the shared budget ledger), but this path ignores the failure and then drops review_reservations. Because this synchronous review response has no persisted job state to retry cleanup, the reservation can remain in reserved and reduce monthly headroom for later requests until reset; handle a failed settlement by retrying/persisting cleanup or failing closed.
Useful? React with 👍 / 👎.
| if reservation_id and bool(getattr(result, "dispatch_started", False)): | ||
| from service.ai_budget_guard import get_ai_budget_guard | ||
|
|
||
| get_ai_budget_guard().settle(reservation_id, 0.10) |
There was a problem hiding this comment.
Handle sync Codex settlement failures
For the synchronous /v1/ai/execute compatibility path, a successful Codex dispatch can still hit a failed settle() call here (for example, a transient SQLite/OSError in the shared budget ledger), but the return value is ignored and there is no persisted job record to retry cleanup. In that case the reservation can remain in reserved after the response, reducing Codex headroom for later requests until reset; fail closed or persist/retry cleanup when settlement fails.
Useful? React with 👍 / 👎.
Co-Authored-By: Codex <noreply@openai.com>
Summary
ai_budget_decision.v1monthly provider/repo/task budget preflight with fail-closed freshness checksdeferred_budget; disable automatic paid-provider fallback by defaultValidation
python3 -m pytest tests -q(634 passed, 1 skipped, 16 subtests)python3 -m ruff check .actionlintgit diff --checkThis PR is intentionally separate from Contract Oscillation Guard PR #72.