diff --git a/.github/workflows/codex_pr_review.yml b/.github/workflows/codex_pr_review.yml index f3f9d3a9..fed16616 100644 --- a/.github/workflows/codex_pr_review.yml +++ b/.github/workflows/codex_pr_review.yml @@ -110,7 +110,7 @@ jobs: CODEX_PR_REVIEW_REUSABLE_CALL: ${{ github.repository != 'QuantStrategyLab/AIAuditBridge' && 'true' || 'false' }} CODEX_PR_REVIEW_API_FALLBACK_INPUT: ${{ inputs.api_fallback_enabled }} CODEX_PR_REVIEW_DIRECT_API_PRIMARY_INPUT: ${{ inputs.direct_api_primary_enabled }} - CODEX_PR_REVIEW_API_FALLBACK_DEFAULT: ${{ vars.CODEX_PR_REVIEW_API_FALLBACK_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 || 'true' }} working-directory: source run: | diff --git a/docs/ai_autonomy_architecture.md b/docs/ai_autonomy_architecture.md index d04509d6..6444cf3d 100644 --- a/docs/ai_autonomy_architecture.md +++ b/docs/ai_autonomy_architecture.md @@ -424,3 +424,24 @@ trusted review comment 只保存最近固定轮数、固定字节上限且脱敏 2. 再补一份“自治边界表”，把自动 / 人工 / 禁止 三类动作列清楚。 3. 然后补指标持久化和 dashboard 视图。 4. 最后再决定是否扩大 auto-merge 覆盖面。 + +## 7. AI Budget Guard v1（中央执行前门） + +所有 `research`、`optimization`、`review` 和 `auto_fix` 入口在启动前必须生成 +`ai_budget_decision.v1` 并取得 `decision=allow`。这是 AIAuditBridge 的中央执行门， +不是各仓库自行复制的 branch rule；确定性监控、kill-switch、暂停和回滚不依赖 AI 额度。 + +- API-key 预算按 provider、key/project scope、repo、task class 和明确的 billing period + （默认 `UTC` 日历月）记账。未配置月度预算时可用上限为 `0`；内部 hard limit 取用户月度预算 + 与 provider/project limit 的 `80%` 两者较小值。 +- usage/cost 快照缺失、过期或无法解析时 fail closed。每次启动先做原子 reservation，并发 + reservation 不能共同突破 hard limit。reservation 的终态由 provider dispatch 证据决定：明确未 + dispatch/provider 未接收时才 release；明确已 dispatch、accepted 或 billable 时按 actual 或保守 + estimate settle；发送边界崩溃、timeout 或重启后无法证明时持久化为 `pending_uncertain`，继续占用 + 额度并 fail closed，直到 provider usage、idempotency 或显式 reconciliation 给出证据。不得静默 + release，也不得把不确定状态直接永久 settle。 +- Codex 同时读取 primary/secondary rate-limit window，使用更紧张的剩余比例；research、 + maintenance/review/auto-fix、incident 的门槛分别是 `30%/20%/10%`，并永久保留最后 `10%`。 + 快照不可用时不启动 research 或 auto-fix。 +- Codex 或 API 额度不足的状态是 `deferred_budget`，禁止自动切换 OpenAI/Anthropic、换 key 或 + 推荐付费低成本模型；只有人工显式批准才允许一次性 fallback。reset 后可安全恢复延后队列。 diff --git a/docs/async_service_deployment.md b/docs/async_service_deployment.md index 4a1a6b53..f2d0979a 100644 --- a/docs/async_service_deployment.md +++ b/docs/async_service_deployment.md @@ -64,8 +64,6 @@ CODEX_AUDIT_SERVICE_AUDIENCE=quant-codex-audit \ CODEX_AUDIT_SERVICE_MODEL=gpt-5.4 \ CODEX_AUDIT_SERVICE_REASONING_EFFORT=auto \ CODEX_AUDIT_SERVICE_CODEX_ACCOUNT_USAGE=1 \ -CODEX_AUDIT_SERVICE_OPENAI_USAGE_WINDOW_DAYS=7 \ -CODEX_AUDIT_SERVICE_ANTHROPIC_USAGE_WINDOW_DAYS=7 \ CODEX_AUDIT_SERVICE_JOB_DIR=/var/lib/codex-audit-bridge/jobs \ bash scripts/deploy_codex_audit_service.sh deploy ``` diff --git a/scripts/run_codex_pr_review.py b/scripts/run_codex_pr_review.py index d5adbea2..073fd553 100644 --- a/scripts/run_codex_pr_review.py +++ b/scripts/run_codex_pr_review.py @@ -527,15 +527,21 @@ def _review_capacity_is_unavailable(exc: ReviewError) -> bool: "[quota_or_capacity_failure]" in message or "usage limit" in message or "daily budget exceeded" in message + or "deferred_budget" in message or "codex service job failed [unknown_failure]: codex exec failed" in message ) def _api_fallback_enabled() -> bool: - return parse_bool(env_value("CODEX_PR_REVIEW_API_FALLBACK_ENABLED", "true")) + # Codex quota exhaustion must become deferred_budget; never consume a + # paid API key unless a human explicitly enables the exception. + return parse_bool(env_value("CODEX_PR_REVIEW_API_FALLBACK_ENABLED", "false")) def _direct_api_primary_enabled() -> bool: + # Preserve the legacy API-only primary path when no Codex service URL is + # configured; only the secondary fallback after Codex quota failure is off + # by default. return parse_bool(env_value("CODEX_PR_REVIEW_DIRECT_API_PRIMARY_ENABLED", "true")) diff --git a/service/adapters/codex_adapter.py b/service/adapters/codex_adapter.py index 55f102fe..20f64dab 100644 --- a/service/adapters/codex_adapter.py +++ b/service/adapters/codex_adapter.py @@ -18,6 +18,20 @@ SECRET_ENV_MARKERS = ("TOKEN", "SECRET", "PASSWORD", "PRIVATE_KEY", "CREDENTIAL", "API_KEY", "ADMIN_KEY") CODEX_REASONING_EFFORTS = frozenset({"minimal", "low", "medium", "high", "xhigh"}) +_LOCAL_CODEX_FAILURE_MARKERS = ( + "authentication", + "bootstrap", + "command not found", + "configuration", + "invalid option", + "invalid sandbox", + "no such file", + "not logged in", + "permission denied", + "unrecognized option", + "unknown option", + "unsupported", +) @dataclass(frozen=True) @@ -25,6 +39,14 @@ class CodexResult: success: bool output: str = "" error: str = "" + dispatch_started: bool = False + dispatch_uncertain: bool = False + + +def _codex_dispatch_is_uncertain(detail: str) -> bool: + """Classify known local CLI failures as definitely not dispatched.""" + text = detail.lower() + return not any(marker in text for marker in _LOCAL_CODEX_FAILURE_MARKERS) def _codex_env() -> dict[str, str]: @@ -122,17 +144,21 @@ def execute( with tempfile.TemporaryDirectory() as tmp: output_last_message = Path(tmp) / "codex-final-message.md" + try: + command = _codex_command( + output_last_message, + sandbox=sandbox, + model=model, + reasoning_effort=reasoning_effort, + output_schema=output_schema, + cwd=cwd, + images=images, + ) + except (RuntimeError, ValueError) as exc: + return CodexResult(success=False, error=f"codex command configuration failed: {exc}") try: completed = subprocess.run( - _codex_command( - output_last_message, - sandbox=sandbox, - model=model, - reasoning_effort=reasoning_effort, - output_schema=output_schema, - cwd=cwd, - images=images, - ), + command, input=prompt, text=True, stdout=subprocess.PIPE, @@ -142,14 +168,22 @@ def execute( env=_codex_env(), ) except subprocess.TimeoutExpired as exc: - return CodexResult(success=False, error=f"codex exec timed out after {timeout}s: {exc}") + return CodexResult( + success=False, + error=f"codex exec timed out after {timeout}s: {exc}", + dispatch_uncertain=True, + ) except FileNotFoundError as exc: return CodexResult(success=False, error=f"codex command not found: {exc}") if completed.returncode != 0: detail = (completed.stdout[-4000:] + completed.stderr[-4000:]).strip() - return CodexResult(success=False, error=f"codex exec failed (rc={completed.returncode})" + (f":\n{detail}" if detail else "")) + return CodexResult( + success=False, + error=f"codex exec failed (rc={completed.returncode})" + (f":\n{detail}" if detail else ""), + dispatch_uncertain=_codex_dispatch_is_uncertain(detail), + ) if output_last_message.exists() and output_last_message.read_text(encoding="utf-8").strip(): - return CodexResult(success=True, output=output_last_message.read_text(encoding="utf-8")) - return CodexResult(success=True, output=completed.stdout) + return CodexResult(success=True, output=output_last_message.read_text(encoding="utf-8"), dispatch_started=True) + return CodexResult(success=True, output=completed.stdout, dispatch_started=True) diff --git a/service/adapters/llm_adapter.py b/service/adapters/llm_adapter.py index 9031c852..0b6643b6 100644 --- a/service/adapters/llm_adapter.py +++ b/service/adapters/llm_adapter.py @@ -55,6 +55,8 @@ class LlmResult: success: bool = True error: str = "" latency_seconds: float = 0.0 + dispatch_started: bool = False + dispatch_uncertain: bool = False class LlmAdapterError(RuntimeError): @@ -105,6 +107,23 @@ def resolve_model(model: str) -> tuple[str, str]: return PROVIDER_ANTHROPIC, DEFAULT_ANTHROPIC_MODEL +def _dispatch_is_uncertain(error: str) -> bool: + """Keep capacity only when provider acceptance cannot be disproven.""" + text = error.lower() + if "api_key is not configured" in text: + return False + if any(f"http {status}" in text for status in range(400, 500)): + return False + return not any(marker in text for marker in ( + "invalid model", + "unsupported model", + "validation", + "malformed", + "missing required", + "not configured", + )) + + # ── OpenAI ───────────────────────────────────────────────────────────── @@ -273,7 +292,13 @@ def complete( output = _anthropic_completion(resolved_model, system, user, max_tokens=max_tokens, timeout=timeout) else: output = _openai_completion(resolved_model, system, user, max_tokens=max_tokens, timeout=timeout) - return LlmResult(provider=provider, model=resolved_model, output=output, latency_seconds=time.time() - started) + return LlmResult( + provider=provider, + model=resolved_model, + output=output, + latency_seconds=time.time() - started, + dispatch_started=True, + ) except LlmAdapterError as exc: return LlmResult( provider=provider, @@ -282,6 +307,8 @@ def complete( success=False, error=str(exc), latency_seconds=time.time() - started, + dispatch_started=False, + dispatch_uncertain=_dispatch_is_uncertain(str(exc)), ) def parallel_review( @@ -299,18 +326,24 @@ def parallel_review( """ import concurrent.futures - results: list[LlmResult] = [] + results: list[LlmResult | None] = [None] * len(reviewers) with concurrent.futures.ThreadPoolExecutor(max_workers=min(len(reviewers), 4)) as pool: futures = { - pool.submit(self.complete, model=model, system=system, user=user, max_tokens=max_tokens, timeout=timeout): (label, model) - for label, model in reviewers + pool.submit(self.complete, model=model, system=system, user=user, max_tokens=max_tokens, timeout=timeout): (index, label, model) + for index, (label, model) in enumerate(reviewers) } for f in concurrent.futures.as_completed(futures): + index, label, model = futures[f] try: - results.append(f.result()) + results[index] = f.result() except Exception as exc: - label, model = futures[f] - results.append( - LlmResult(provider=label, model=model, output="", success=False, error=str(exc)) + results[index] = ( + LlmResult( + provider=label, + model=model, + output="", + success=False, + error=str(exc), + ) ) - return results + return [result for result in results if result is not None] diff --git a/service/ai_budget_guard.py b/service/ai_budget_guard.py new file mode 100644 index 00000000..63e5a9b8 --- /dev/null +++ b/service/ai_budget_guard.py @@ -0,0 +1,857 @@ +"""Fail-closed monthly AI budget and Codex rate-limit gate. + +The guard is deliberately independent from model selection: a low budget never +silently selects another paid provider. Callers must reserve before starting +work and settle or release the reservation when the work finishes. +""" + +from __future__ import annotations + +import json +import math +import os +import sqlite3 +import threading +import time +import uuid +from dataclasses import dataclass, replace +from datetime import UTC, datetime +from pathlib import Path +from typing import Any + +DECISION_SCHEMA = "ai_budget_decision.v1" +DEFAULT_BILLING_TIMEZONE = "UTC" +DEFAULT_MAX_USAGE_AGE_SECONDS = 86_400 +CODEX_RESERVE_RATIO = 0.10 +CODEX_RESERVATION_RATIO = 0.10 +CODEX_THRESHOLDS = { + "research": 0.30, + "optimization": 0.20, + "maintenance": 0.20, + "review": 0.20, + "auto_fix": 0.20, + "incident": 0.10, +} +API_TASK_CLASSES = frozenset(CODEX_THRESHOLDS) + + +def _now() -> float: + return time.time() + + +def _number(value: Any, default: float = 0.0) -> float: + try: + return float(value) + except (TypeError, ValueError): + return default + + +def _timestamp(snapshot: dict[str, Any]) -> float | None: + for key in ("observed_at", "updated_at", "as_of", "timestamp"): + value = snapshot.get(key) + if isinstance(value, str): + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp() + except ValueError: + continue + if value is not None: + parsed = _number(value, -1) + if parsed >= 0: + return parsed + return None + + +def _reset_timestamp(value: Any) -> float | None: + if isinstance(value, str): + try: + return datetime.fromisoformat(value.replace("Z", "+00:00")).timestamp() + except ValueError: + return None + parsed = _number(value, -1.0) + return parsed if parsed >= 0 else None + + +def _period(now: float, timezone_name: str) -> str: + # Billing periods are calendar months. Invalid timezone configuration is + # not allowed to change the period silently; UTC is the safe fallback. + try: + from zoneinfo import ZoneInfo + + zone = ZoneInfo(timezone_name or DEFAULT_BILLING_TIMEZONE) + except Exception: # noqa: BLE001 + zone = UTC + return datetime.fromtimestamp(now, zone).strftime("%Y-%m") + + +@dataclass(frozen=True) +class Reservation: + reservation_id: str + scope: str + amount: float + task_class: str + created_at: float + baseline_usage: float | None = None + period: str = "" + aggregate_scope: str = "" + state: str = "reserved" + + +class AIBudgetGuard: + """In-process atomic reservation ledger with a versioned decision contract.""" + + def __init__(self, config: dict[str, Any] | None = None, *, clock: Any = _now) -> None: + self._clock = clock + self._lock = threading.RLock() + self._reservations: dict[str, Reservation] = {} + self._reserved: dict[str, float] = {} + self._settled: dict[str, float] = {} + self._settled_baseline: dict[str, float] = {} + self._codex_reset_events: dict[str, float] = {} + self._config = config if isinstance(config, dict) else self._load_config() + self._timezone = str(self._config.get("billing_timezone") or DEFAULT_BILLING_TIMEZONE) + self._max_age = max(1, int(_number(self._config.get("usage_max_age_seconds"), DEFAULT_MAX_USAGE_AGE_SECONDS))) + self._settled_period = self.period() + ledger_path = str(self._config.get("ledger_path") or os.environ.get("CODEX_AUDIT_SERVICE_AI_BUDGET_LEDGER_PATH", "")).strip() + if not ledger_path: + job_dir = os.environ.get("CODEX_AUDIT_SERVICE_JOB_DIR", "").strip() + ledger_path = str(Path(job_dir) / "ai_budget_ledger.sqlite3") if job_dir else "" + self._ledger_path = ledger_path + allow_in_memory = bool(self._config.get("allow_in_memory_ledger")) or os.environ.get( + "CODEX_AUDIT_SERVICE_ALLOW_IN_MEMORY_LEDGER", "" + ).strip().lower() in {"1", "true", "yes"} + self._ledger_error = not self._ledger_path and not allow_in_memory + self._init_ledger() + + def _init_ledger(self) -> None: + if not self._ledger_path: + return + try: + path = Path(self._ledger_path) + path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) + with sqlite3.connect(path, timeout=30) as db: + db.execute( + "CREATE TABLE IF NOT EXISTS ai_budget_ledger (scope TEXT NOT NULL, period TEXT NOT NULL, reserved REAL NOT NULL, settled REAL NOT NULL, baseline REAL, updated_at REAL NOT NULL, PRIMARY KEY(scope, period))" + ) + db.execute( + "CREATE TABLE IF NOT EXISTS ai_budget_reservations (reservation_id TEXT PRIMARY KEY, scope TEXT NOT NULL, aggregate_scope TEXT NOT NULL, period TEXT NOT NULL, amount REAL NOT NULL, task_class TEXT NOT NULL, created_at REAL NOT NULL, baseline_usage REAL, state TEXT NOT NULL DEFAULT 'reserved')" + ) + columns = { + row[1] + for row in db.execute("PRAGMA table_info(ai_budget_reservations)").fetchall() + } + if "aggregate_scope" not in columns: + db.execute( + "ALTER TABLE ai_budget_reservations ADD COLUMN aggregate_scope TEXT NOT NULL DEFAULT ''" + ) + db.execute( + "UPDATE ai_budget_reservations SET aggregate_scope=scope WHERE aggregate_scope=''" + ) + if "state" not in columns: + db.execute( + "ALTER TABLE ai_budget_reservations ADD COLUMN state TEXT NOT NULL DEFAULT 'reserved'" + ) + except (OSError, sqlite3.Error): + self._ledger_error = True + + def _load_scope_locked(self, scope: str) -> None: + if not self._ledger_path: + return + try: + with sqlite3.connect(self._ledger_path, timeout=30) as db: + row = db.execute( + "SELECT period,reserved,settled,baseline FROM ai_budget_ledger WHERE scope=? AND period=?", + (scope, self._settled_period), + ).fetchone() + if row is None: + return + period, reserved, settled, baseline = row + if period != self._settled_period: + return + self._reserved[scope] = max(0.0, float(reserved)) + self._settled[scope] = max(0.0, float(settled)) + if baseline is not None: + self._settled_baseline[scope] = float(baseline) + except (OSError, sqlite3.Error, TypeError, ValueError): + self._ledger_error = True + return + + def _save_scope_locked(self, scope: str) -> None: + if not self._ledger_path: + return + try: + with sqlite3.connect(self._ledger_path, timeout=30) as db: + db.execute( + "INSERT INTO ai_budget_ledger(scope,period,reserved,settled,baseline,updated_at) VALUES(?,?,?,?,?,?) " + "ON CONFLICT(scope,period) DO UPDATE SET reserved=excluded.reserved,settled=excluded.settled,baseline=excluded.baseline,updated_at=excluded.updated_at", + ( + scope, + self._settled_period, + self._reserved.get(scope, 0.0), + self._settled.get(scope, 0.0), + self._settled_baseline.get(scope), + self._clock(), + ), + ) + except (OSError, sqlite3.Error): + self._ledger_error = True + return + + @staticmethod + def _load_config() -> dict[str, Any]: + path = os.environ.get("CODEX_AUDIT_SERVICE_AI_BUDGET_CONFIG", "").strip() + if not path: + return {} + try: + raw = json.loads(Path(path).read_text(encoding="utf-8")) + except (OSError, ValueError): + return {"_config_error": "budget_config_unavailable"} + return raw if isinstance(raw, dict) else {"_config_error": "budget_config_invalid"} + + @property + def billing_timezone(self) -> str: + return self._timezone + + def period(self, now: float | None = None) -> str: + return _period(self._clock() if now is None else now, self._timezone) + + def _scope(self, provider: str, provider_scope: str, repo: str, task_class: str, period: str) -> str: + return "/".join((provider or "unknown", provider_scope or "default", repo or "unknown", task_class or "maintenance", period)) + + @staticmethod + def _aggregate_scope(provider: str, provider_scope: str, period: str) -> str: + # Aggregate reservations share one provider/project cap across all + # task classes; task-specific thresholds stay in the decision only. + return "/".join((provider or "unknown", provider_scope or "default", "_aggregate", period)) + + def _budget_entry(self, provider: str, provider_scope: str, repo: str, task_class: str) -> dict[str, Any] | None: + budgets = self._config.get("monthly_budgets") + if not isinstance(budgets, dict): + return None + candidates = ( + f"{provider}:{provider_scope}:{repo}:{task_class}", + f"{provider}:{provider_scope}:{repo}", + f"{provider}:{provider_scope}", + provider, + "default", + ) + for key in candidates: + value = budgets.get(key) + if isinstance(value, (int, float)): + return {"user_monthly_budget_usd": float(value)} + if isinstance(value, dict): + if any(field in value for field in ("user_monthly_budget_usd", "monthly_budget_usd", "provider_project_limit_usd")): + return value + provider_node = budgets.get(provider) + if isinstance(provider_node, dict): + scope_node = provider_node.get(provider_scope) or provider_node.get("default") + if isinstance(scope_node, dict): + repo_node = scope_node.get(repo) or scope_node.get("default") + if isinstance(repo_node, dict): + task_node = repo_node.get(task_class) or repo_node.get("default") + if isinstance(task_node, (int, float)): + return {"user_monthly_budget_usd": float(task_node)} + if isinstance(task_node, dict): + return task_node + return None + + def _usage(self, snapshot: dict[str, Any] | None) -> tuple[float, bool, str]: + if not isinstance(snapshot, dict): + return 0.0, False, "usage_snapshot_missing" + observed = _timestamp(snapshot) + if observed is None or self._clock() - observed > self._max_age: + return 0.0, False, "usage_snapshot_stale" + value = snapshot.get("used_usd", snapshot.get("cost_usd", snapshot.get("total_cost"))) + if value is None: + for nested_key in ("organization_costs", "costs", "usage"): + nested = snapshot.get(nested_key) + if isinstance(nested, dict): + value = nested.get("used_usd", nested.get("cost_usd", nested.get("total_cost"))) + if value is not None: + break + if value is None: + return 0.0, False, "usage_cost_missing" + if isinstance(value, bool): + return 0.0, False, "usage_cost_invalid" + try: + parsed = float(value) + except (TypeError, ValueError): + return 0.0, False, "usage_cost_invalid" + if not math.isfinite(parsed) or parsed < 0: + return 0.0, False, "usage_cost_invalid" + return parsed, True, "" + + def _reconcile_settled_locked(self, scope: str, observed: float) -> float: + settled = self._settled.get(scope, 0.0) + baseline = self._settled_baseline.get(scope) + if settled <= 0 or baseline is None or observed <= baseline: + return settled + observed_advance = observed - baseline + if observed_advance >= settled: + self._settled.pop(scope, None) + self._settled_baseline.pop(scope, None) + return 0.0 + remaining = settled - observed_advance + self._settled[scope] = remaining + self._settled_baseline[scope] = observed + return remaining + + def _decision(self, *, task_class: str, provider_scope: str, period: str, observed: Any, + reserved: float, hard_limit: Any, remaining: Any, freshness: str, + decision: str, reasons: list[str], reset_at: Any = None, + aggregate_observed: Any = None, aggregate_hard_limit: Any = None, + aggregate_reserved: float = 0.0, + auto_fallback_allowed: bool = False, reservation_scope: str = "", aggregate_scope: str = "") -> dict[str, Any]: + return { + "schema": DECISION_SCHEMA, + "task_class": task_class, + "provider_scope": provider_scope, + "reservation_scope": reservation_scope or provider_scope, + "aggregate_scope": aggregate_scope, + "period": period, + "observed_usage": observed, + "aggregate_observed_usage": aggregate_observed if aggregate_observed is not None else observed, + "reserved_usage": round(reserved, 8), + "hard_limit": hard_limit, + "aggregate_hard_limit": aggregate_hard_limit if aggregate_hard_limit is not None else hard_limit, + "aggregate_reserved_usage": round(aggregate_reserved, 8), + "remaining_after_reservation": remaining, + "usage_freshness": freshness, + "decision": decision, + "reason_codes": reasons, + "reset_at": reset_at, + "billing_timezone": self._timezone, + "auto_fallback_allowed": auto_fallback_allowed, + } + + def preflight( + self, + *, + task_class: str, + provider: str, + provider_scope: str = "default", + repo: str = "", + estimated_cost_usd: float = 0.0, + usage_snapshot: dict[str, Any] | None = None, + codex_snapshot: dict[str, Any] | None = None, + human_approved_fallback: bool = False, + ) -> dict[str, Any]: + task = str(task_class or "maintenance").strip().lower() + provider_name = str(provider or "").strip().lower() + current_period = self.period() + if self._ledger_error: + return self._decision( + task_class=task, provider_scope=provider_scope, period=current_period, + observed={}, reserved=0.0, hard_limit=0.0, remaining=0.0, + freshness="unavailable", decision="block", reasons=["shared_budget_ledger_unavailable"], + ) + scope = self._scope(provider_name, provider_scope, repo, task, current_period) + aggregate_scope = self._aggregate_scope(provider_name, provider_scope, current_period) + amount = max(0.0, _number(estimated_cost_usd)) + if provider_name in {"codex", "codex-cli", "codex_account"}: + return self._preflight_codex(task, provider_scope, current_period, codex_snapshot, scope, aggregate_scope) + + entry = self._budget_entry(provider_name, provider_scope, repo, task) + if entry is None: + return self._decision( + task_class=task, provider_scope=provider_scope, period=current_period, + observed=0.0, reserved=0.0, hard_limit=0.0, remaining=0.0, + freshness="unknown", decision="defer", reasons=["monthly_budget_not_configured"], + ) + has_provider_project_limit = entry.get("provider_project_limit_usd") is not None + used, fresh, freshness_reason = self._usage(usage_snapshot) + if not fresh: + return self._decision( + task_class=task, provider_scope=provider_scope, period=current_period, + observed=used, reserved=0.0, hard_limit=0.0, remaining=0.0, + freshness="stale", decision="block", reasons=[freshness_reason], + ) + user_limit = max(0.0, _number(entry.get("user_monthly_budget_usd", entry.get("monthly_budget_usd")))) + has_provider_project_limit = entry.get("provider_project_limit_usd") is not None + provider_limit = _number(entry.get("provider_project_limit_usd"), user_limit) + hard_limit = min(user_limit, provider_limit * 0.80) if has_provider_project_limit else user_limit * 0.80 + aggregate_hard_limit = max(0.0, provider_limit * 0.80) if has_provider_project_limit else hard_limit + reservation_aggregate_scope = ( + self._aggregate_scope(provider_name, provider_scope, current_period) + if has_provider_project_limit else scope + ) + effective_used = used + with self._lock: + if self._settled_period != current_period: + self._reserved.clear() + self._settled.clear() + self._settled_baseline.clear() + self._reservations.clear() + self._settled_period = current_period + self._load_scope_locked(scope) + if reservation_aggregate_scope != scope: + self._load_scope_locked(reservation_aggregate_scope) + reserved = self._reserved.get(scope, 0.0) + settled = self._reconcile_settled_locked(scope, effective_used if reservation_aggregate_scope == scope else 0.0) + aggregate_reserved = self._reserved.get(reservation_aggregate_scope, 0.0) + aggregate_settled = settled if reservation_aggregate_scope == scope else self._reconcile_settled_locked(reservation_aggregate_scope, effective_used) + repo_remaining = hard_limit - settled - reserved - amount + aggregate_accounted = effective_used + aggregate_settled + aggregate_remaining = aggregate_hard_limit - aggregate_accounted - aggregate_reserved - amount + remaining = min(repo_remaining, aggregate_remaining) + if remaining < 0: + reason = "monthly_hard_limit_reached" if remaining <= 0 else "monthly_budget_insufficient" + return self._decision( + task_class=task, provider_scope=provider_scope, period=current_period, + observed=0.0, reserved=reserved, hard_limit=round(hard_limit, 8), + remaining=round(max(0.0, remaining), 8), freshness="fresh", decision="defer", + reasons=[item for item in (reason, "api_fallback_requires_human_approval" if not human_approved_fallback else "") if item], + aggregate_observed=effective_used, aggregate_hard_limit=round(aggregate_hard_limit, 8), + aggregate_reserved=aggregate_reserved, + ) + return self._decision( + task_class=task, provider_scope=provider_scope, period=current_period, + observed=0.0, reserved=reserved, hard_limit=round(hard_limit, 8), + remaining=round(remaining, 8), freshness="fresh", decision="allow", reasons=[], + aggregate_observed=effective_used, aggregate_hard_limit=round(aggregate_hard_limit, 8), + aggregate_reserved=aggregate_reserved, + auto_fallback_allowed=human_approved_fallback, + reservation_scope=scope, aggregate_scope=reservation_aggregate_scope, + ) + + def _preflight_codex(self, task: str, provider_scope: str, period: str, + snapshot: dict[str, Any] | None, scope: str, aggregate_scope: str) -> dict[str, Any]: + if not isinstance(snapshot, dict): + return self._decision( + task_class=task, provider_scope=provider_scope, period=period, + observed={}, reserved=0.0, hard_limit=CODEX_RESERVE_RATIO, + remaining=0.0, freshness="unavailable", decision="defer", + reasons=["codex_rate_limit_snapshot_unavailable"], + ) + observed_at = _timestamp(snapshot) + if observed_at is None or self._clock() - observed_at > self._max_age: + return self._decision( + task_class=task, provider_scope=provider_scope, period=period, + observed=snapshot, reserved=0.0, hard_limit=CODEX_RESERVE_RATIO, + remaining=0.0, freshness="stale", decision="defer", + reasons=["codex_rate_limit_snapshot_stale"], + ) + limits = snapshot.get("rate_limits") if isinstance(snapshot.get("rate_limits"), dict) else snapshot + ratios: list[float] = [] + reset_at: list[Any] = [] + for name in ("primary", "secondary"): + window = limits.get(name) + if not isinstance(window, dict): + return self._decision( + task_class=task, provider_scope=provider_scope, period=period, + observed=snapshot, reserved=0.0, hard_limit=CODEX_RESERVE_RATIO, + remaining=0.0, freshness="invalid", decision="defer", + reasons=["codex_rate_limit_window_missing"], + ) + remaining = window.get("remaining_percent") + if remaining is None and window.get("used_percent") is not None: + raw_used = window.get("used_percent") + try: + parsed_used = float(raw_used) + except (TypeError, ValueError): + parsed_used = float("nan") + if not math.isfinite(parsed_used): + return self._decision( + task_class=task, provider_scope=provider_scope, period=period, + observed=snapshot, reserved=0.0, hard_limit=CODEX_RESERVE_RATIO, + remaining=0.0, freshness="invalid", decision="defer", + reasons=["codex_rate_limit_window_invalid"], + ) + remaining = 100 - parsed_used + if remaining is None: + return self._decision( + task_class=task, provider_scope=provider_scope, period=period, + observed=snapshot, reserved=0.0, hard_limit=CODEX_RESERVE_RATIO, + remaining=0.0, freshness="invalid", decision="defer", + reasons=["codex_rate_limit_window_invalid"], + ) + try: + parsed_remaining = float(remaining) + except (TypeError, ValueError): + parsed_remaining = float("nan") + if not math.isfinite(parsed_remaining): + return self._decision( + task_class=task, provider_scope=provider_scope, period=period, + observed=snapshot, reserved=0.0, hard_limit=CODEX_RESERVE_RATIO, + remaining=0.0, freshness="invalid", decision="defer", + reasons=["codex_rate_limit_window_invalid"], + ) + ratios.append(max(0.0, min(1.0, parsed_remaining / 100.0))) + reset_at.append(window.get("resets_at")) + tightest = min(ratios) + with self._lock: + if self._settled_period != period: + self._reserved.clear() + self._settled.clear() + self._settled_baseline.clear() + self._reservations.clear() + self._settled_period = period + self._load_scope_locked(scope) + parsed_resets = [ + parsed + for value in reset_at + if (parsed := _reset_timestamp(value)) is not None + ] + reset_event = min(parsed_resets) if len(parsed_resets) == len(reset_at) and parsed_resets else None + reset_key = aggregate_scope + 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 + scope_prefix = f"codex/{provider_scope or 'default'}/" + for key in list(self._reserved): + if key.startswith(scope_prefix): + self._reserved.pop(key, None) + for key in list(self._settled): + if key.startswith(scope_prefix): + self._settled.pop(key, None) + for key in list(self._settled_baseline): + if key.startswith(scope_prefix): + self._settled_baseline.pop(key, None) + for reservation_id, reservation in list(self._reservations.items()): + if reservation.scope.startswith(scope_prefix) or reservation.aggregate_scope.startswith(scope_prefix): + self._reservations.pop(reservation_id, None) + if self._ledger_path: + try: + with sqlite3.connect(self._ledger_path, timeout=30) as db: + db.execute( + "DELETE FROM ai_budget_reservations WHERE scope LIKE ? OR aggregate_scope LIKE ?", + (f"{scope_prefix}%", f"{scope_prefix}%"), + ) + db.execute("DELETE FROM ai_budget_ledger WHERE scope LIKE ?", (f"{scope_prefix}%",)) + except (OSError, sqlite3.Error): + pass + settled = self._settled.get(scope, 0.0) + reserved = self._reserved.get(scope, 0.0) + threshold = CODEX_THRESHOLDS.get(task, CODEX_THRESHOLDS["maintenance"]) + headroom = max(0.0, tightest - CODEX_RESERVE_RATIO) + remaining = max(0.0, headroom - settled - reserved) + allowed = tightest >= threshold and remaining >= CODEX_RESERVATION_RATIO + return self._decision( + task_class=task, provider_scope=provider_scope, period=period, + observed={"primary_remaining_ratio": ratios[0], "secondary_remaining_ratio": ratios[1]}, + reserved=round(settled + reserved, 8), hard_limit=round(headroom, 8), + remaining=round(remaining, 8), freshness="fresh", + decision="allow" if allowed else "defer", + reasons=[] if allowed else [ + "codex_rate_limit_below_task_threshold" + if tightest < threshold + else "codex_reservation_headroom_insufficient" + ], + reset_at=reset_at, + reservation_scope=scope, + aggregate_scope=aggregate_scope, + ) + + def reserve(self, decision: dict[str, Any], amount: float | None = None) -> Reservation | None: + if not isinstance(decision, dict) or decision.get("decision") != "allow": + return None + if amount is None: + return None + scope = str(decision.get("reservation_scope") or decision.get("provider_scope") or "default") + aggregate_scope = str(decision.get("aggregate_scope") or scope) + # Provider scope in a decision is intentionally the stable ledger key; + # callers may pass an explicit amount for API reservations. + requested = max(0.0, _number(amount, _number(decision.get("remaining_after_reservation")))) + observed_value = decision.get("observed_usage") + aggregate_observed_value = decision.get("aggregate_observed_usage") + baseline_usage = _number(aggregate_observed_value, _number(observed_value, -1.0)) + reservation = Reservation( + uuid.uuid4().hex, + scope, + requested, + str(decision.get("task_class") or "maintenance"), + self._clock(), + baseline_usage if baseline_usage >= 0 else None, + self.period(), + aggregate_scope, + ) + with self._lock: + current_period = self.period() + decision_period = str(decision.get("period") or current_period) + if decision_period != current_period: + return None + if self._settled_period != current_period: + self._reserved.clear() + self._settled.clear() + self._settled_baseline.clear() + self._reservations.clear() + self._settled_period = current_period + if self._ledger_path: + try: + with sqlite3.connect(self._ledger_path, timeout=30) as db: + db.execute("BEGIN IMMEDIATE") + row = db.execute( + "SELECT reserved,settled,baseline FROM ai_budget_ledger WHERE scope=? AND period=?", + (scope, self._settled_period), + ).fetchone() + aggregate_row = db.execute( + "SELECT reserved,settled,baseline FROM ai_budget_ledger WHERE scope=? AND period=?", + (aggregate_scope, self._settled_period), + ).fetchone() + stored_reserved = max(0.0, float(row[0])) if row else 0.0 + stored_settled = max(0.0, float(row[1])) if row else 0.0 + stored_baseline = float(row[2]) if row and row[2] is not None else None + aggregate_reserved = max(0.0, float(aggregate_row[0])) if aggregate_row else 0.0 + aggregate_settled = max(0.0, float(aggregate_row[1])) if aggregate_row else 0.0 + aggregate_baseline = float(aggregate_row[2]) if aggregate_row and aggregate_row[2] is not None else None + observed_value = decision.get("observed_usage") + observed = float(observed_value) if isinstance(observed_value, (int, float)) and not isinstance(observed_value, bool) else 0.0 + aggregate_observed_value = decision.get("aggregate_observed_usage") + aggregate_observed = float(aggregate_observed_value) if isinstance(aggregate_observed_value, (int, float)) and not isinstance(aggregate_observed_value, bool) else observed + if aggregate_baseline is not None and aggregate_observed >= aggregate_baseline + aggregate_settled: + aggregate_settled = 0.0 + aggregate_baseline = None + if stored_baseline is not None and observed >= stored_baseline + stored_settled: + stored_settled = 0.0 + stored_baseline = None + hard_limit = _number(decision.get("hard_limit"), float("inf")) + aggregate_hard_limit = _number(decision.get("aggregate_hard_limit"), hard_limit) + if hard_limit != float("inf") and ( + observed + stored_settled + stored_reserved + requested > hard_limit + 1e-12 + or aggregate_observed + aggregate_settled + aggregate_reserved + requested > aggregate_hard_limit + 1e-12 + ): + return None + new_reserved = stored_reserved + requested + db.execute( + "INSERT INTO ai_budget_ledger(scope,period,reserved,settled,baseline,updated_at) VALUES(?,?,?,?,?,?) " + "ON CONFLICT(scope,period) DO UPDATE SET reserved=excluded.reserved,settled=excluded.settled,baseline=excluded.baseline,updated_at=excluded.updated_at", + (scope, self._settled_period, new_reserved, stored_settled, stored_baseline, self._clock()), + ) + if aggregate_scope != scope: + db.execute( + "INSERT INTO ai_budget_ledger(scope,period,reserved,settled,baseline,updated_at) VALUES(?,?,?,?,?,?) " + "ON CONFLICT(scope,period) DO UPDATE SET reserved=excluded.reserved,settled=excluded.settled,baseline=excluded.baseline,updated_at=excluded.updated_at", + (aggregate_scope, self._settled_period, aggregate_reserved + requested, aggregate_settled, aggregate_baseline, self._clock()), + ) + db.execute( + "INSERT INTO ai_budget_reservations(reservation_id,scope,aggregate_scope,period,amount,task_class,created_at,baseline_usage,state) VALUES(?,?,?,?,?,?,?,?,?)", + (reservation.reservation_id, reservation.scope, reservation.aggregate_scope, reservation.period, reservation.amount, reservation.task_class, reservation.created_at, reservation.baseline_usage, reservation.state), + ) + db.commit() + self._load_scope_locked(scope) + if aggregate_scope != scope: + self._load_scope_locked(aggregate_scope) + self._reservations[reservation.reservation_id] = reservation + return reservation + except (OSError, sqlite3.Error, TypeError, ValueError): + return None + self._load_scope_locked(scope) + hard_limit = _number(decision.get("hard_limit"), float("inf")) + aggregate_hard_limit = _number(decision.get("aggregate_hard_limit"), hard_limit) + observed_value = decision.get("observed_usage") + if isinstance(observed_value, (int, float)) and not isinstance(observed_value, bool): + observed = float(observed_value) + settled = self._reconcile_settled_locked(scope, observed) + else: + observed = 0.0 + settled = self._settled.get(scope, 0.0) + aggregate_observed = _number(decision.get("aggregate_observed_usage"), observed) + aggregate_settled = self._reconcile_settled_locked(aggregate_scope, aggregate_observed) + aggregate_reserved = self._reserved.get(aggregate_scope, 0.0) + if hard_limit != float("inf") and ( + observed + settled + self._reserved.get(scope, 0.0) + requested > hard_limit + 1e-12 + or aggregate_observed + aggregate_settled + aggregate_reserved + requested > aggregate_hard_limit + 1e-12 + ): + return None + self._reservations[reservation.reservation_id] = reservation + self._reserved[scope] = self._reserved.get(scope, 0.0) + requested + if aggregate_scope != scope: + self._reserved[aggregate_scope] = aggregate_reserved + requested + self._save_scope_locked(scope) + if aggregate_scope != scope: + self._save_scope_locked(aggregate_scope) + return reservation + + def mark_uncertain(self, reservation: Reservation | str) -> bool: + """Persist an ambiguous provider dispatch without refunding its reservation. + + This is deliberately neither release nor settle: only an explicit + provider/idempotency reconciliation can decide the final outcome. + """ + rid = reservation.reservation_id if isinstance(reservation, Reservation) else str(reservation) + with self._lock: + if self._ledger_path: + try: + with sqlite3.connect(self._ledger_path, timeout=30) as db: + cursor = db.execute( + "UPDATE ai_budget_reservations SET state='pending_uncertain' WHERE reservation_id=?", + (rid,), + ) + if cursor.rowcount != 1: + return False + item = self._reservations.get(rid) + if item is not None: + self._reservations[rid] = replace(item, state="pending_uncertain") + return True + except (OSError, sqlite3.Error): + return False + item = self._reservations.get(rid) + if item is None: + return False + self._reservations[rid] = replace(item, state="pending_uncertain") + return True + + def reconcile_pending_uncertain( + self, + reservation: Reservation | str, + *, + dispatched: bool | None, + actual_cost: float = 0.0, + ) -> bool: + """Resolve a pending dispatch only from explicit provider evidence.""" + if dispatched is None: + return False + rid = reservation.reservation_id if isinstance(reservation, Reservation) else str(reservation) + if self._ledger_path: + try: + with sqlite3.connect(self._ledger_path, timeout=30) as db: + row = db.execute( + "SELECT state FROM ai_budget_reservations WHERE reservation_id=?", (rid,) + ).fetchone() + if row is None or row[0] != "pending_uncertain": + return False + except (OSError, sqlite3.Error): + return False + else: + item = self._reservations.get(rid) + if item is None or item.state != "pending_uncertain": + return False + return self.settle(rid, actual_cost) if dispatched else self.release(rid) + + def release(self, reservation: Reservation | str) -> bool: + rid = reservation.reservation_id if isinstance(reservation, Reservation) else str(reservation) + with self._lock: + if self._ledger_path: + try: + with sqlite3.connect(self._ledger_path, timeout=30) as db: + db.execute("BEGIN IMMEDIATE") + row = db.execute( + "SELECT scope,aggregate_scope,period,amount FROM ai_budget_reservations WHERE reservation_id=?", + (rid,), + ).fetchone() + if row is None: + db.rollback() + return False + scope, aggregate_scope, period, amount = row + db.execute( + "UPDATE ai_budget_ledger SET reserved=MAX(0,reserved-?),updated_at=? WHERE scope=? AND period=?", + (float(amount), self._clock(), scope, period), + ) + if aggregate_scope != scope: + db.execute( + "UPDATE ai_budget_ledger SET reserved=MAX(0,reserved-?),updated_at=? WHERE scope=? AND period=?", + (float(amount), self._clock(), aggregate_scope, period), + ) + db.execute("DELETE FROM ai_budget_reservations WHERE reservation_id=?", (rid,)) + self._reservations.pop(rid, None) + self._load_scope_locked(scope) + if aggregate_scope != scope: + self._load_scope_locked(aggregate_scope) + return True + except (OSError, sqlite3.Error, TypeError, ValueError): + return False + item = self._reservations.pop(rid, None) + if item is None: + return False + if item.period != self.period(): + return True + self._reserved[item.scope] = max(0.0, self._reserved.get(item.scope, 0.0) - item.amount) + if item.aggregate_scope != item.scope: + self._reserved[item.aggregate_scope] = max(0.0, self._reserved.get(item.aggregate_scope, 0.0) - item.amount) + self._save_scope_locked(item.scope) + if item.aggregate_scope != item.scope: + self._save_scope_locked(item.aggregate_scope) + return True + + def settle(self, reservation: Reservation | str, actual_cost: float) -> bool: + rid = reservation.reservation_id if isinstance(reservation, Reservation) else str(reservation) + with self._lock: + if self._ledger_path: + try: + amount = max(0.0, _number(actual_cost)) + with sqlite3.connect(self._ledger_path, timeout=30) as db: + db.execute("BEGIN IMMEDIATE") + row = db.execute( + "SELECT scope,aggregate_scope,period,amount,baseline_usage FROM ai_budget_reservations WHERE reservation_id=?", + (rid,), + ).fetchone() + if row is None: + db.rollback() + return False + scope, aggregate_scope, period, reserved_amount, baseline_usage = row + ledger = db.execute( + "SELECT period,settled,baseline FROM ai_budget_ledger WHERE scope=? AND period=?", + (scope, period), + ).fetchone() + period = ledger[0] if ledger else period + settled = float(ledger[1]) if ledger else 0.0 + baseline = ledger[2] if ledger else None + if baseline is None: + baseline = baseline_usage + db.execute( + "INSERT INTO ai_budget_ledger(scope,period,reserved,settled,baseline,updated_at) VALUES(?,?,?,?,?,?) " + "ON CONFLICT(scope,period) DO UPDATE SET reserved=MAX(0,ai_budget_ledger.reserved-?),settled=ai_budget_ledger.settled+?,baseline=COALESCE(ai_budget_ledger.baseline,excluded.baseline),updated_at=excluded.updated_at", + (scope, period, max(0.0, float(reserved_amount)), settled + amount, baseline, self._clock(), float(reserved_amount), amount), + ) + if aggregate_scope != scope: + aggregate_row = db.execute( + "SELECT settled,baseline FROM ai_budget_ledger WHERE scope=? AND period=?", + (aggregate_scope, period), + ).fetchone() + aggregate_baseline = aggregate_row[1] if aggregate_row and aggregate_row[1] is not None else baseline_usage + db.execute( + "UPDATE ai_budget_ledger SET reserved=MAX(0,reserved-?),settled=settled+?,baseline=COALESCE(baseline,?),updated_at=? WHERE scope=? AND period=?", + (float(reserved_amount), amount, aggregate_baseline, self._clock(), aggregate_scope, period), + ) + db.execute("DELETE FROM ai_budget_reservations WHERE reservation_id=?", (rid,)) + self._reservations.pop(rid, None) + self._load_scope_locked(scope) + if aggregate_scope != scope: + self._load_scope_locked(aggregate_scope) + return True + except (OSError, sqlite3.Error, TypeError, ValueError): + return False + item = self._reservations.pop(rid, None) + if item is None: + return False + if item.period != self.period(): + return True + self._reserved[item.scope] = max(0.0, self._reserved.get(item.scope, 0.0) - item.amount) + if item.aggregate_scope != item.scope: + self._reserved[item.aggregate_scope] = max(0.0, self._reserved.get(item.aggregate_scope, 0.0) - item.amount) + amount = max(0.0, _number(actual_cost)) + if amount: + self._settled.setdefault(item.scope, 0.0) + if item.baseline_usage is not None: + self._settled_baseline.setdefault(item.scope, item.baseline_usage) + self._settled[item.scope] += amount + if item.aggregate_scope != item.scope: + if item.baseline_usage is not None: + self._settled_baseline.setdefault(item.aggregate_scope, item.baseline_usage) + self._settled[item.aggregate_scope] = self._settled.get(item.aggregate_scope, 0.0) + amount + self._save_scope_locked(item.scope) + if item.aggregate_scope != item.scope: + self._save_scope_locked(item.aggregate_scope) + return True + + +_guard = AIBudgetGuard() + + +def get_ai_budget_guard() -> AIBudgetGuard: + return _guard + + +def ai_budget_preflight(**kwargs: Any) -> dict[str, Any]: + return _guard.preflight(**kwargs) + + +__all__ = [ + "AIBudgetGuard", + "API_TASK_CLASSES", + "CODEX_RESERVE_RATIO", + "CODEX_RESERVATION_RATIO", + "CODEX_THRESHOLDS", + "DECISION_SCHEMA", + "Reservation", + "ai_budget_preflight", + "get_ai_budget_guard", +] diff --git a/service/ai_gateway_service.py b/service/ai_gateway_service.py index e8a5921e..4c4be782 100644 --- a/service/ai_gateway_service.py +++ b/service/ai_gateway_service.py @@ -44,7 +44,7 @@ parse_execute_request, parse_review_request, ) -from service.adapters.llm_adapter import LlmAdapter, resolve_model +from service.adapters.llm_adapter import DEFAULT_MAX_TOKENS, LlmAdapter, resolve_model from service.adapters.codex_adapter import CodexAdapter from service.autonomy import ( ACTION_AUTO_PR, @@ -111,6 +111,15 @@ def try_record_platform_execution(*_args, **_kwargs): TRUSTED_AUTOMATION_PROOF_PATH_ENV = "CODEX_AUDIT_SERVICE_TRUSTED_AUTOMATION_PROOF_PATH" DASHBOARD_REPOSITORIES_ENV = "CODEX_AUDIT_SERVICE_DASHBOARD_REPOSITORIES" + +def _budget_gate_enabled() -> bool: + """Keep the explicit unauthenticated local test mode backward compatible.""" + return not ( + os.environ.get("CODEX_AUDIT_SERVICE_AUTH", "").strip().lower() == "none" + and os.environ.get("CODEX_AUDIT_SERVICE_ALLOW_NO_AUTH_FOR_LOCAL_TESTS", "").strip().lower() + in {"1", "true", "yes", "on"} + ) + # sandbox allowlist — restrict what callers can request ALLOWED_SANDBOXES: frozenset[str] = frozenset( os.environ.get("CODEX_AUDIT_SERVICE_ALLOWED_SANDBOXES", "read-only").replace(" ", "").split(",") @@ -471,6 +480,11 @@ def _cleanup_expired_jobs() -> None: except Exception: expires_at = 0 if expires_at and expires_at < now: + reservation_id = str(payload.get("_budget_reservation_id") or "") + if reservation_id: + _finalize_budget_reservation(payload, 0.10) + if payload.get("_budget_cleanup_pending") or payload.get("_budget_reservation_id"): + continue try: path.unlink() except FileNotFoundError: @@ -494,6 +508,7 @@ def _recover_orphaned_jobs() -> int: job["failure_category"] = "service_restart" _write_job(job) _record_job_automation_run(job) + _finalize_budget_reservation(job, 0.10) _audit_log( "job_failed", job_id=job.get("job_id"), @@ -517,9 +532,64 @@ def _mark_stale_job_failed(job: dict[str, Any]) -> dict[str, Any]: job["failure_category"] = "stale_job_timeout" _write_job(job) _record_job_automation_run(job) + _finalize_budget_reservation(job, 0.10) return job +def _settle_budget_reservation(payload: dict[str, Any], actual_cost: float = 0.0) -> None: + reservation_id = str(payload.get("_budget_reservation_id") or "") + if not reservation_id: + return + from service.ai_budget_guard import get_ai_budget_guard + + if not get_ai_budget_guard().settle(reservation_id, actual_cost): + _audit_log("budget_reservation_settle_failed", reservation_id=reservation_id) + payload["_budget_cleanup_pending"] = True + _write_job(payload) + return + payload.pop("_budget_reservation_id", None) + _write_job(payload) + + +def _release_budget_reservation(payload: dict[str, Any]) -> None: + reservation_id = str(payload.get("_budget_reservation_id") or "") + if not reservation_id: + return + from service.ai_budget_guard import get_ai_budget_guard + + if not get_ai_budget_guard().release(reservation_id): + _audit_log("budget_reservation_release_failed", reservation_id=reservation_id) + payload["_budget_cleanup_pending"] = True + _write_job(payload) + return + payload.pop("_budget_reservation_id", None) + _write_job(payload) + + +def _mark_budget_reservation_uncertain(payload: dict[str, Any]) -> None: + reservation_id = str(payload.get("_budget_reservation_id") or "") + if not reservation_id: + return + from service.ai_budget_guard import get_ai_budget_guard + + if not get_ai_budget_guard().mark_uncertain(reservation_id): + _audit_log("budget_reservation_uncertain_mark_failed", reservation_id=reservation_id) + payload["_budget_cleanup_pending"] = True + else: + payload["_budget_reservation_state"] = "pending_uncertain" + _write_job(payload) + + +def _finalize_budget_reservation(payload: dict[str, Any], actual_cost: float = 0.0) -> None: + state = str(payload.get("dispatch_state") or "") + if state == "pending_uncertain": + _mark_budget_reservation_uncertain(payload) + elif state == "dispatched" or (not state and bool(payload.get("dispatch_started"))): + _settle_budget_reservation(payload, actual_cost) + else: + _release_budget_reservation(payload) + + def _assert_job_access(job: dict[str, Any], claims: dict[str, Any]) -> None: repository = str(claims.get("repository") or "") if repository != str(job.get("repository") or ""): @@ -1108,14 +1178,27 @@ def _run_job(job_id: str, payload: dict[str, Any]) -> None: adapter = CodexAdapter() sandbox = _validate_sandbox(str(payload.get("sandbox") or "")) reasoning_effort = _resolve_codex_reasoning_effort(payload, str(payload.get("task") or TASK_EXECUTE)) + prompt = str(payload["prompt"]) + timeout = int(payload.get("timeout_seconds", 2700)) + # A process crash at this boundary cannot prove whether the provider + # accepted work. Preserve that uncertainty without settling it. + job["dispatch_state"] = "pending_uncertain" + job["updated_at"] = _now() + _write_job(job) result = adapter.execute( - prompt=str(payload["prompt"]), + prompt=prompt, sandbox=sandbox, model=str(payload.get("model") or "").strip() or None, reasoning_effort=reasoning_effort, - timeout=int(payload.get("timeout_seconds", 2700)), + timeout=timeout, ) job = _read_job(job_id) + job["dispatch_started"] = bool(getattr(result, "dispatch_started", False)) + job["dispatch_state"] = ( + "dispatched" if job["dispatch_started"] + else "pending_uncertain" if bool(getattr(result, "dispatch_uncertain", False)) + else "not_dispatched" + ) if result.success: job["status"] = "succeeded" job["output"] = result.output @@ -1147,6 +1230,7 @@ def _run_job(job_id: str, payload: dict[str, Any]) -> None: }, domain=str(job.get("domain") or ""), ) + _finalize_budget_reservation(job, 0.10) _audit_log("job_completed", job_id=job_id, status=job["status"], repository=job.get("repository"), task=job.get("task")) except Exception as exc: @@ -1172,6 +1256,7 @@ def _run_job(job_id: str, payload: dict[str, Any]) -> None: }, domain=str(job.get("domain") or ""), ) + _finalize_budget_reservation(job, 0.10) _audit_log("job_failed", job_id=job_id, error=type(exc).__name__, repository=job.get("repository")) @@ -1216,6 +1301,9 @@ def _submit_job(claims: dict[str, Any], payload: dict[str, Any]) -> dict[str, ob "mode": str(payload.get("mode") or MODE_REVIEW_ONLY), "timeout_seconds": int(payload.get("timeout_seconds", 2700)), "dedupe_key": dedupe_key, + "_budget_reservation_id": str(payload.get("_budget_reservation_id") or ""), + "dispatch_started": False, + "dispatch_state": "not_dispatched", } _write_job(job) _record_job_automation_run(job) @@ -1402,26 +1490,60 @@ def _handle_analyze(self, payload: dict[str, Any]) -> None: # Quota check quota = get_quota_manager() resolved_model = _resolve_analyze_model(req.model) - qr = quota.check(source_repo, resolved_model, req.prompt) + qr = quota.check( + source_repo, resolved_model, req.prompt, task_class="research", + estimated_output_tokens=req.max_tokens, + budget_guard=_budget_gate_enabled(), + ) if not qr["allowed"]: _json_response(self, HTTPStatus.TOO_MANY_REQUESTS, { - "status": "error", + "status": "error", "deferred_budget": True, "error": qr["reason"], "recommended_model": qr.get("recommended_model", ""), "remaining_usd": qr.get("remaining_usd", 0), + "budget_decision": qr.get("budget_decision", {}), }) return started = time.time() adapter = LlmAdapter() - result = adapter.complete( - model=resolved_model, system=req.system, user=req.prompt, - max_tokens=req.max_tokens, timeout=req.timeout_seconds, - ) + reservation_id = str(qr.get("budget_reservation_id") or "") + try: + result = adapter.complete( + model=resolved_model, system=req.system, user=req.prompt, + max_tokens=req.max_tokens, timeout=req.timeout_seconds, + ) + except Exception: + if reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().mark_uncertain(reservation_id) + raise + 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, + float(qr.get("cost_estimate_usd") or 0.0), + ) + elif reservation_id and bool(getattr(result, "dispatch_uncertain", False)): + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().mark_uncertain(reservation_id) + elif reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().release(reservation_id) + # A local quota-recording failure must not roll back provider spend. + 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 "") + except Exception as exc: # noqa: BLE001 - provider response is already settled. + quota.mark_recording_failed(source_repo) + _audit_log("analyze_quota_record_failed", error=type(exc).__name__) latency = time.time() - started # Record quota and health - quota.record(source_repo, resolved_model, req.prompt, result.output if result.success else "") get_health_monitor().record("/v1/ai/analyze", latency, result.success, result.error if not result.success else "") _audit_log("analyze_completed", model=result.model, provider=result.provider, @@ -1453,17 +1575,43 @@ def _handle_execute_async(self, claims: dict[str, Any], payload: dict[str, Any]) # Quota check quota = get_quota_manager() - qr = quota.check(quota_repo, "codex-cli", req.prompt, codex_account=True) + qr = quota.check( + quota_repo, "codex-cli", req.prompt, codex_account=True, + budget_guard=_budget_gate_enabled(), + task_class="auto_fix" if str(payload.get("task") or "").lower() in {"auto_fix", "autofix"} else "maintenance", + ) if not qr["allowed"]: _json_response(self, HTTPStatus.TOO_MANY_REQUESTS, { - "status": "error", "error": qr["reason"], + "status": "error", "deferred_budget": True, "error": qr["reason"], "remaining_usd": qr.get("remaining_usd", 0), + "budget_decision": qr.get("budget_decision", {}), }) return - quota.record_execute(quota_repo) + reservation_id = str(qr.get("budget_reservation_id") or "") + try: + quota.record_execute(quota_repo) + except Exception: + if reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().release(reservation_id) + raise payload.setdefault("task", TASK_EXECUTE) - job = _submit_job(claims, payload) + if reservation_id: + payload["_budget_reservation_id"] = reservation_id + try: + job = _submit_job(claims, payload) + except Exception: + if reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().release(reservation_id) + raise + if reservation_id and bool(job.get("deduped")): + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().release(reservation_id) get_health_monitor().record("/v1/ai/execute/jobs", time.time() - started, True) _json_response(self, HTTPStatus.ACCEPTED, job) @@ -1477,24 +1625,61 @@ def _handle_execute_sync(self, claims: dict[str, Any], payload: dict[str, Any]) _validate_source_repo_org(claims, source_repo) quota_repo = source_repo or str(claims.get("repository") or "unknown") quota = get_quota_manager() - qr = quota.check(quota_repo, "codex-cli", req.prompt, codex_account=True) + qr = quota.check( + quota_repo, "codex-cli", req.prompt, codex_account=True, + budget_guard=_budget_gate_enabled(), + task_class="auto_fix" if str(payload.get("task") or "").lower() in {"auto_fix", "autofix"} else "maintenance", + ) if not qr["allowed"]: _json_response(self, HTTPStatus.TOO_MANY_REQUESTS, { - "status": "error", "error": qr["reason"], + "status": "error", "deferred_budget": True, "error": qr["reason"], "remaining_usd": qr.get("remaining_usd", 0), + "budget_decision": qr.get("budget_decision", {}), }) return - quota.record_execute(quota_repo) adapter = CodexAdapter() - sandbox = _validate_sandbox(str(payload.get("sandbox") or "")) - reasoning_effort = _resolve_codex_reasoning_effort(payload, str(payload.get("task") or TASK_EXECUTE)) - result = adapter.execute( - prompt=req.prompt, - sandbox=sandbox, - model=req.model or None, - reasoning_effort=reasoning_effort, - timeout=req.timeout_seconds, - ) + reservation_id = str(qr.get("budget_reservation_id") or "") + try: + quota.record_execute(quota_repo) + except Exception: + if reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().release(reservation_id) + raise + adapter_invoked = False + try: + sandbox = _validate_sandbox(str(payload.get("sandbox") or "")) + reasoning_effort = _resolve_codex_reasoning_effort(payload, str(payload.get("task") or TASK_EXECUTE)) + adapter_invoked = True + result = adapter.execute( + prompt=req.prompt, + sandbox=sandbox, + model=req.model or None, + reasoning_effort=reasoning_effort, + timeout=req.timeout_seconds, + ) + except Exception: + if reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + if adapter_invoked: + get_ai_budget_guard().mark_uncertain(reservation_id) + else: + get_ai_budget_guard().release(reservation_id) + raise + 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) + elif reservation_id and bool(getattr(result, "dispatch_uncertain", False)): + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().mark_uncertain(reservation_id) + elif reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().release(reservation_id) get_health_monitor().record("/v1/ai/execute", time.time() - started, result.success, result.error if not result.success else "") if result.success: _json_response(self, HTTPStatus.OK, {"status": "ok", "output": result.output}) @@ -1528,22 +1713,190 @@ def _handle_review(self, claims: dict[str, Any], payload: dict[str, Any]) -> Non ] _audit_log("review_started", reviewers=req.reviewers, verifier=req.verifier, changed_paths_count=len(changed_paths)) - llm_results = llm.parallel_review( - reviewers=reviewer_tuples, - system="You are a careful quantitative strategy reviewer. Return JSON with verdict, confidence (0.0-1.0), and summary.", - user=req.prompt, - timeout=req.timeout_seconds, - ) - - # Step 2: optional Codex verification - codex_result = None + source_repository = str(payload.get("source_repository") or "") + if source_repository: + _validate_source_repo(source_repository) + if not _automation_operator_claims(claims): + _validate_source_repo_org(claims, source_repository) + review_repo = source_repository or str(claims.get("repository") or "unknown") + codex_quota: dict[str, Any] | None = None + codex_reservation_id = "" if req.verifier == "codex": - codex_result = codex.execute( - prompt=req.prompt, - sandbox=_validate_sandbox("read-only"), - reasoning_effort=_resolve_codex_reasoning_effort(payload, TASK_REVIEW), + codex_quota = get_quota_manager().check( + review_repo, + "codex-cli", + req.prompt, + codex_account=True, + budget_guard=_budget_gate_enabled(), + task_class="review", + ) + if not codex_quota.get("allowed"): + _json_response(self, HTTPStatus.TOO_MANY_REQUESTS, { + "status": "error", "deferred_budget": True, + "error": codex_quota.get("reason", "AI budget gate denied Codex review"), + "budget_decision": codex_quota.get("budget_decision", {}), + }) + return + codex_reservation_id = str(codex_quota.get("budget_reservation_id") or "") + review_reservations: list[tuple[str, float, int]] = [] + if _budget_gate_enabled(): + estimated_total = sum( + float( + get_quota_manager().check( + review_repo, + reviewer_model, + req.prompt, + estimated_output_tokens=DEFAULT_MAX_TOKENS, + ).get("cost_estimate_usd") or 0.0 + ) + for _reviewer_name, reviewer_model in reviewer_tuples + ) + if codex_quota is not None: + estimated_total += float(codex_quota.get("cost_estimate_usd") or 0.0) + if get_quota_manager().remaining_daily(review_repo) < estimated_total: + if codex_reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().release(codex_reservation_id) + _json_response(self, HTTPStatus.TOO_MANY_REQUESTS, { + "status": "error", "deferred_budget": True, + "error": "deferred_budget: combined reviewer daily budget is insufficient", + }) + return + for reviewer_index, (reviewer_name, reviewer_model) in enumerate(reviewer_tuples): + review_quota = get_quota_manager().check( + review_repo, + reviewer_model, + req.prompt, + estimated_output_tokens=DEFAULT_MAX_TOKENS, + task_class="review", + budget_guard=_budget_gate_enabled(), + ) + if not review_quota.get("allowed"): + if review_reservations: + from service.ai_budget_guard import get_ai_budget_guard + + for reservation_id, _cost, _index in review_reservations: + get_ai_budget_guard().release(reservation_id) + if codex_reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().release(codex_reservation_id) + _json_response(self, HTTPStatus.TOO_MANY_REQUESTS, { + "status": "error", "deferred_budget": True, + "error": review_quota.get("reason", "AI budget gate denied review"), + "reviewer": reviewer_name, + "budget_decision": review_quota.get("budget_decision", {}), + }) + return + if review_quota.get("budget_reservation_id"): + review_reservations.append( + ( + str(review_quota["budget_reservation_id"]), + float(review_quota.get("cost_estimate_usd") or 0.0), + reviewer_index, + ) + ) + try: + llm_results = llm.parallel_review( + reviewers=reviewer_tuples, + system="You are a careful quantitative strategy reviewer. Return JSON with verdict, confidence (0.0-1.0), and summary.", + user=req.prompt, timeout=req.timeout_seconds, ) + except Exception: + if review_reservations: + from service.ai_budget_guard import get_ai_budget_guard + + for reservation_id, _estimated_cost, _index in review_reservations: + get_ai_budget_guard().release(reservation_id) + if codex_reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().release(codex_reservation_id) + raise + # Keep the legacy per-repo quota ledger aligned with the provider + # reservation ledger for every reviewer invocation. + quota_record_failed = False + for review_result in llm_results: + if not review_result.success and not bool(getattr(review_result, "dispatch_started", False)): + continue + try: + get_quota_manager().record( + review_repo, + str(review_result.model or ""), + req.prompt, + review_result.output if review_result.success else "", + ) + except Exception as exc: # noqa: BLE001 - provider spend is already reserved/settled. + get_quota_manager().mark_recording_failed(review_repo) + quota_record_failed = True + _audit_log("review_quota_record_failed", error=type(exc).__name__) + if review_reservations: + from service.ai_budget_guard import get_ai_budget_guard + + 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) + elif matched_result is None or bool(getattr(matched_result, "dispatch_uncertain", False)): + get_ai_budget_guard().mark_uncertain(reservation_id) + else: + get_ai_budget_guard().release(reservation_id) + review_reservations = [] + if quota_record_failed and codex_reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().release(codex_reservation_id) + codex_reservation_id = "" + codex_quota = None + if quota_record_failed: + _json_response(self, HTTPStatus.SERVICE_UNAVAILABLE, { + "status": "error", + "deferred_budget": True, + "error": "quota_ledger_unavailable", + }) + return + # Step 2: optional Codex verification + codex_result = None + if req.verifier == "codex" and codex_quota is not None: + codex_adapter_invoked = False + try: + codex_sandbox = _validate_sandbox("read-only") + codex_reasoning_effort = _resolve_codex_reasoning_effort(payload, TASK_REVIEW) + codex_adapter_invoked = True + codex_result = codex.execute( + prompt=req.prompt, + sandbox=codex_sandbox, + reasoning_effort=codex_reasoning_effort, + timeout=req.timeout_seconds, + ) + except Exception: + if codex_reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + if codex_adapter_invoked: + get_ai_budget_guard().mark_uncertain(codex_reservation_id) + else: + get_ai_budget_guard().release(codex_reservation_id) + raise + if codex_reservation_id and bool(getattr(codex_result, "dispatch_started", False)): + from service.ai_budget_guard import get_ai_budget_guard + + try: + get_quota_manager().record_execute(review_repo) + except Exception as exc: # noqa: BLE001 - central reservation remains conservative. + get_quota_manager().mark_recording_failed(review_repo) + _audit_log("review_codex_quota_record_failed", error=type(exc).__name__) + get_ai_budget_guard().settle(codex_reservation_id, 0.10) + elif codex_reservation_id and bool(getattr(codex_result, "dispatch_uncertain", False)): + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().mark_uncertain(codex_reservation_id) + elif codex_reservation_id: + from service.ai_budget_guard import get_ai_budget_guard + + get_ai_budget_guard().release(codex_reservation_id) # Step 3: build per-reviewer results with extracted confidence results: list[dict[str, Any]] = [] diff --git a/service/anthropic_admin_usage.py b/service/anthropic_admin_usage.py index 0b4d7683..09fdc911 100644 --- a/service/anthropic_admin_usage.py +++ b/service/anthropic_admin_usage.py @@ -31,6 +31,22 @@ def _int_env(name: str, default: int, *, minimum: int = 1, maximum: int | None = return value +def _usage_window(end_time: int, billing_timezone: str = "UTC") -> tuple[int, int]: + configured = os.environ.get("CODEX_AUDIT_SERVICE_ANTHROPIC_USAGE_WINDOW_DAYS", "").strip() + if configured: + days = _int_env("CODEX_AUDIT_SERVICE_ANTHROPIC_USAGE_WINDOW_DAYS", 1, minimum=1, maximum=31) + return end_time - days * 86400, days + try: + from zoneinfo import ZoneInfo + + zone = ZoneInfo(billing_timezone) + except Exception: # noqa: BLE001 - invalid configuration must use the guard's UTC fallback. + zone = UTC + current = datetime.fromtimestamp(end_time, zone) + start = current.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + return int(start.timestamp()), (current.date() - start.date()).days + 1 + + def _iso_utc(timestamp: int) -> str: return datetime.fromtimestamp(timestamp, UTC).strftime("%Y-%m-%dT%H:%M:%SZ") @@ -131,14 +147,17 @@ def _sum_costs(payloads: list[dict[str, Any]]) -> dict[str, Any]: return {"total_cost": round(by_currency.get(currency, 0.0), 4), "currency": currency, "result_count": result_count} -def read_anthropic_admin_usage(now: int | None = None, timeout_seconds: float | None = None) -> dict[str, Any] | None: +def read_anthropic_admin_usage( + now: int | None = None, + timeout_seconds: float | None = None, + billing_timezone: str = "UTC", +) -> dict[str, Any] | None: """Return a sanitized Anthropic org usage/cost snapshot, or None when unavailable.""" admin_key = _admin_key() if not admin_key: return None - days = _int_env("CODEX_AUDIT_SERVICE_ANTHROPIC_USAGE_WINDOW_DAYS", 7, minimum=1, maximum=31) end_time = int(now if now is not None else time.time()) - start_time = end_time - days * 86400 + start_time, days = _usage_window(end_time, billing_timezone) timeout = timeout_seconds if timeout_seconds is not None else float( _int_env("CODEX_AUDIT_SERVICE_ANTHROPIC_ADMIN_TIMEOUT_SECONDS", 8, minimum=1, maximum=60) ) diff --git a/service/automation_decision.py b/service/automation_decision.py index 769798ad..62de844f 100644 --- a/service/automation_decision.py +++ b/service/automation_decision.py @@ -333,6 +333,7 @@ def decide_automation_execution( recent_runs: list[dict[str, Any]] | None = None, failure_history_complete: bool = True, policy: dict[str, Any] | None = None, + ai_budget_decision: dict[str, Any] | None = None, ) -> dict[str, Any]: """Produce a safe execution decision from health, quota, failures, and repo policy.""" repo_policy = repo_execution_policy(repo, policy) @@ -357,6 +358,7 @@ def decide_automation_execution( action = EXECUTION_RUN human_review_required = False defer = False + budget_deferred = False service = _normalize_status(service_health) quota = _normalize_quota_status(quota_status) @@ -412,6 +414,20 @@ def decide_automation_execution( human_review_required = True reasons.append("health unhealthy; forcing human review") + # Budget is a central preflight gate. A missing/failed decision is not an + # invitation to select another paid provider; it is deferred_budget. + if ai_budget_decision is not None: + budget_result = str(ai_budget_decision.get("decision") or "block").strip().lower() + if budget_result != "allow": + if action != EXECUTION_HUMAN_REVIEW: + action = EXECUTION_DEFER + defer = True + budget_deferred = True + effective_mode = MODE_REVIEW_ONLY + reasons.append( + "ai budget decision is " + (budget_result or "unavailable") + "; execution deferred_budget" + ) + if quota in {"low", "constrained"}: if action in {EXECUTION_HUMAN_REVIEW, EXECUTION_DEFER}: reasons.append(f"quota status is {quota}; execution already blocked") @@ -422,7 +438,7 @@ def decide_automation_execution( reasons.append(f"quota status is {quota}; deferring automation") elif quota_low_behavior == "defer": reasons.append(f"quota status is {quota}; automation already review_only") - else: + elif not budget_deferred: effective_model = low_cost_model or recommend_model(0.0) effective_provider = low_cost_provider or "auto" reasons.append(f"quota status is {quota}; recommending low-cost model") @@ -463,5 +479,7 @@ def decide_automation_execution( and not human_review_required ), "defer": defer, + "deferred_budget": budget_deferred, + "ai_budget_decision": ai_budget_decision, "reasons": reasons or ["execution allowed"], } diff --git a/service/openai_admin_usage.py b/service/openai_admin_usage.py index 8649c5f0..17b7e2a4 100644 --- a/service/openai_admin_usage.py +++ b/service/openai_admin_usage.py @@ -5,6 +5,7 @@ import json import os import time +from datetime import UTC, datetime from typing import Any from urllib.error import HTTPError, URLError from urllib.parse import urlencode, urlparse @@ -29,6 +30,22 @@ def _int_env(name: str, default: int, *, minimum: int = 1, maximum: int | None = return value +def _usage_window(end_time: int, billing_timezone: str = "UTC") -> tuple[int, int]: + configured = os.environ.get("CODEX_AUDIT_SERVICE_OPENAI_USAGE_WINDOW_DAYS", "").strip() + if configured: + days = _int_env("CODEX_AUDIT_SERVICE_OPENAI_USAGE_WINDOW_DAYS", 1, minimum=1, maximum=31) + return end_time - days * 86400, days + try: + from zoneinfo import ZoneInfo + + zone = ZoneInfo(billing_timezone) + except Exception: # noqa: BLE001 - invalid configuration must use the guard's UTC fallback. + zone = UTC + current = datetime.fromtimestamp(end_time, zone) + start = current.replace(day=1, hour=0, minute=0, second=0, microsecond=0) + return int(start.timestamp()), (current.date() - start.date()).days + 1 + + def _split_csv_env(name: str) -> list[str]: return [item.strip() for item in os.environ.get(name, "").split(",") if item.strip()] @@ -123,14 +140,17 @@ def _sum_costs(payloads: list[dict[str, Any]]) -> dict[str, Any]: } -def read_openai_admin_usage(now: int | None = None, timeout_seconds: float | None = None) -> dict[str, Any] | None: +def read_openai_admin_usage( + now: int | None = None, + timeout_seconds: float | None = None, + billing_timezone: str = "UTC", +) -> dict[str, Any] | None: """Return a sanitized OpenAI completions usage snapshot, or None when unavailable.""" admin_key = _admin_key() if not admin_key: return None - days = _int_env("CODEX_AUDIT_SERVICE_OPENAI_USAGE_WINDOW_DAYS", 7, minimum=1, maximum=31) end_time = int(now if now is not None else time.time()) - start_time = end_time - days * 86400 + start_time, days = _usage_window(end_time, billing_timezone) timeout = timeout_seconds if timeout_seconds is not None else float( _int_env("CODEX_AUDIT_SERVICE_OPENAI_ADMIN_TIMEOUT_SECONDS", 8, minimum=1, maximum=60) ) diff --git a/service/quota.py b/service/quota.py index b27bdfd8..32aa606a 100644 --- a/service/quota.py +++ b/service/quota.py @@ -219,6 +219,7 @@ def __init__(self): self._daily_budget = DEFAULT_DAILY_BUDGET_USD self._weekly_budget = DEFAULT_WEEKLY_BUDGET_USD self._repo_budgets: dict[str, dict[str, float]] = {} + self._record_failures: set[str] = set() self._codex_account_cache: dict[str, Any] | None = None self._codex_account_cache_ts = 0.0 self._codex_account_attempt_ts = 0.0 @@ -275,6 +276,9 @@ def _load_records(self) -> None: for repo, item in records.items() if isinstance(repo, str) and isinstance(item, dict) } + persisted_failures = raw.get("record_failures") if isinstance(raw, dict) else None + if isinstance(persisted_failures, list): + self._record_failures.update(str(repo) for repo in persisted_failures if isinstance(repo, str)) def _save_records_locked(self) -> None: path = self._store_path() @@ -282,7 +286,10 @@ def _save_records_locked(self) -> None: return path.parent.mkdir(mode=0o700, parents=True, exist_ok=True) payload = json.dumps( - {"records": {repo: record.to_dict() for repo, record in self._records.items()}}, + { + "records": {repo: record.to_dict() for repo, record in self._records.items()}, + "record_failures": sorted(self._record_failures), + }, ensure_ascii=False, sort_keys=True, ).encode("utf-8") @@ -322,6 +329,24 @@ def get_daily_budget(self, repo: str) -> float: def get_weekly_budget(self, repo: str) -> float: return self._repo_budgets.get(repo, {}).get("weekly", self._weekly_budget) + def mark_recording_failed(self, repo: str) -> None: + """Fail closed for a repo after its quota record cannot be persisted.""" + with self._lock: + self._record_failures.add(repo or "unknown") + try: + self._save_records_locked() + except Exception: # noqa: BLE001 - preserve in-memory fail-closed state. + pass + + def recover_recording_failure(self, repo: str) -> None: + """Explicit recovery after quota state reconciliation.""" + with self._lock: + self._record_failures.discard(repo or "unknown") + try: + self._save_records_locked() + except Exception: # noqa: BLE001 - keep the in-memory fail-closed state. + pass + @staticmethod def _api_budget_cost(record: QuotaRecord) -> float: """Return spend governed by the internal API-key USD budget. @@ -387,6 +412,8 @@ def check( estimated_output_tokens: int = 0, *, codex_account: bool = False, + task_class: str = "review", + budget_guard: bool = False, ) -> dict[str, Any]: """Check whether the selected provider is allowed to run. @@ -394,11 +421,110 @@ def check( execution handlers set ``codex_account`` after authenticating a request; only that internal signal may bypass the API-key budget. """ + with self._lock: + if repo in self._record_failures: + return { + "allowed": False, + "reason": "quota_ledger_unavailable", + "remaining_usd": 0.0, + "decision": "block", + "deferred_budget": True, + } tokens_input = estimate_tokens(prompt) - cost = estimate_cost(model, tokens_input, estimated_output_tokens) - if codex_account: + resolved_model = model + provider = "" + reservation = None + if budget_guard and not codex_account: + from service.adapters.llm_adapter import resolve_model + + provider, resolved_model = resolve_model(model) + cost = estimate_cost(resolved_model, tokens_input, estimated_output_tokens) + if budget_guard and not codex_account and self.remaining_daily(repo) < cost: + return { + "allowed": False, + "reason": f"deferred_budget: daily budget has ${self.remaining_daily(repo):.4f} remaining, ${cost:.4f} needed", + "remaining_usd": self.remaining_daily(repo), + "cost_estimate_usd": cost, + "decision": "defer", + "deferred_budget": True, + } + if budget_guard and not codex_account: + from service.ai_budget_guard import get_ai_budget_guard + + budget_guard_instance = get_ai_budget_guard() + snapshot = ( + self._anthropic_account_snapshot() + if provider == "anthropic" + else self._openai_account_snapshot() + ) + provider_scope = ( + os.environ.get("CODEX_AUDIT_SERVICE_ANTHROPIC_WORKSPACE_SCOPE", "") + if provider == "anthropic" + else os.environ.get("CODEX_AUDIT_SERVICE_OPENAI_PROJECT_SCOPE", "") + ).strip() or "default" + budget_decision = budget_guard_instance.preflight( + task_class=task_class, + provider=provider, + provider_scope=provider_scope, + repo=repo, + estimated_cost_usd=cost, + usage_snapshot=snapshot, + ) + if budget_decision.get("decision") != "allow": + return { + "allowed": False, + "reason": "; ".join(budget_decision.get("reason_codes") or ["AI budget gate denied API execution"]), + "budget_decision": budget_decision, + "remaining_usd": 0.0, + } + reservation = budget_guard_instance.reserve(budget_decision, cost) + if reservation is None: + return { + "allowed": False, + "reason": "deferred_budget: atomic reservation unavailable", + "budget_decision": {**budget_decision, "decision": "defer", "reason_codes": ["reservation_conflict"]}, + "remaining_usd": 0.0, + } + if codex_account and budget_guard: if model != "codex-cli": raise ValueError("codex_account quota checks require model=codex-cli") + from service.ai_budget_guard import get_ai_budget_guard + + budget_guard_instance = get_ai_budget_guard() + provider_scope = os.environ.get("CODEX_AUDIT_SERVICE_CODEX_ACCOUNT_SCOPE", "").strip() or "codex-account" + budget_decision = budget_guard_instance.preflight( + task_class=task_class, + provider="codex", + provider_scope=provider_scope, + repo=repo, + codex_snapshot=self._codex_account_snapshot(), + ) + if budget_decision.get("decision") != "allow": + return { + "allowed": False, + "reason": "; ".join(budget_decision.get("reason_codes") or ["AI budget gate denied Codex execution"]), + "quota_scope": "codex_account", + "budget_decision": budget_decision, + "remaining_usd": 0.0, + } + reservation = budget_guard_instance.reserve(budget_decision, 0.10) + if reservation is None: + return { + "allowed": False, + "reason": "deferred_budget: atomic Codex reservation unavailable", + "quota_scope": "codex_account", + "budget_decision": {**budget_decision, "decision": "defer", "reason_codes": ["reservation_conflict"]}, + "remaining_usd": 0.0, + } + return { + "allowed": True, + "cost_estimate_usd": cost, + "remaining_usd": self.remaining_daily(repo), + "quota_scope": "codex_account", + "budget_decision": budget_decision, + "budget_reservation_id": reservation.reservation_id, + } + if codex_account: return { "allowed": True, "cost_estimate_usd": cost, @@ -408,18 +534,19 @@ def check( remaining = self.remaining_daily(repo) if remaining < cost: - recommended = recommend_model(remaining) return { "allowed": False, - "reason": f"Daily budget exceeded: ${remaining:.4f} remaining, ${cost:.4f} needed", - "recommended_model": recommended, + "reason": f"deferred_budget: daily budget has ${remaining:.4f} remaining, ${cost:.4f} needed", "remaining_usd": remaining, "cost_estimate_usd": cost, + "decision": "defer", + "deferred_budget": True, } return { "allowed": True, "cost_estimate_usd": cost, "remaining_usd": remaining - cost, + "budget_reservation_id": reservation.reservation_id if reservation is not None else "", } def record(self, repo: str, model: str, prompt: str, output: str = "") -> None: @@ -447,6 +574,7 @@ def record(self, repo: str, model: str, prompt: str, output: str = "") -> None: record.total_cost_usd += cost self._records[repo] = record self._save_records_locked() + self._record_failures.discard(repo or "unknown") def record_execute(self, repo: str) -> None: """Record a Codex exec call with a nominal dashboard estimate only.""" @@ -461,6 +589,7 @@ def record_execute(self, repo: str) -> None: record.total_cost_usd += cost self._records[repo] = record self._save_records_locked() + self._record_failures.discard(repo or "unknown") def _codex_account_snapshot(self, timeout_seconds: float | None = None) -> dict[str, Any] | None: with self._codex_account_lock: @@ -494,7 +623,12 @@ def _openai_account_snapshot(self, timeout_seconds: float | None = None) -> dict if now - self._openai_account_attempt_ts < failure_ttl: return None self._openai_account_attempt_ts = now - snapshot = read_openai_admin_usage(timeout_seconds=timeout_seconds) + from service.ai_budget_guard import get_ai_budget_guard + + snapshot = read_openai_admin_usage( + timeout_seconds=timeout_seconds, + billing_timezone=get_ai_budget_guard().billing_timezone, + ) if snapshot: self._openai_account_cache = snapshot self._openai_account_cache_ts = now @@ -513,7 +647,12 @@ def _anthropic_account_snapshot(self, timeout_seconds: float | None = None) -> d if now - self._anthropic_account_attempt_ts < failure_ttl: return None self._anthropic_account_attempt_ts = now - snapshot = read_anthropic_admin_usage(timeout_seconds=timeout_seconds) + from service.ai_budget_guard import get_ai_budget_guard + + snapshot = read_anthropic_admin_usage( + timeout_seconds=timeout_seconds, + billing_timezone=get_ai_budget_guard().billing_timezone, + ) if snapshot: self._anthropic_account_cache = snapshot self._anthropic_account_cache_ts = now diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..192e383f --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,4 @@ +import os + + +os.environ.setdefault("CODEX_AUDIT_SERVICE_ALLOW_IN_MEMORY_LEDGER", "1") diff --git a/tests/test_ai_budget_guard.py b/tests/test_ai_budget_guard.py new file mode 100644 index 00000000..4b6011d4 --- /dev/null +++ b/tests/test_ai_budget_guard.py @@ -0,0 +1,254 @@ +from __future__ import annotations + +import time +from pathlib import Path +from tempfile import TemporaryDirectory + +from service.ai_budget_guard import AIBudgetGuard, DECISION_SCHEMA + + +def _codex_snapshot( + primary: int = 50, + secondary: int = 80, + observed_at: float | None = None, + reset_at: float | None = None, +) -> dict: + reset = time.time() + 3600 if reset_at is None else reset_at + return { + "updated_at": time.time() if observed_at is None else observed_at, + "rate_limits": { + "primary": {"remaining_percent": primary, "resets_at": reset}, + "secondary": {"remaining_percent": secondary, "resets_at": reset}, + }, + } + + +def test_unconfigured_api_budget_is_zero_and_deferred() -> None: + guard = AIBudgetGuard({"billing_timezone": "UTC"}) + decision = guard.preflight( + task_class="research", provider="openai", provider_scope="project", repo="o/r", estimated_cost_usd=1, + usage_snapshot={"updated_at": time.time(), "used_usd": 0}, + ) + assert decision["schema"] == DECISION_SCHEMA + assert decision["decision"] == "defer" + assert decision["hard_limit"] == 0 + assert "monthly_budget_not_configured" in decision["reason_codes"] + + +def test_api_hard_limit_is_min_of_user_and_provider_eighty_percent() -> None: + guard = AIBudgetGuard({"monthly_budgets": {"openai:project": { + "user_monthly_budget_usd": 100, "provider_project_limit_usd": 50, + }}}) + decision = guard.preflight( + task_class="review", provider="openai", provider_scope="project", repo="o/r", estimated_cost_usd=1, + usage_snapshot={"updated_at": time.time(), "used_usd": 0}, + ) + assert decision["decision"] == "allow" + assert decision["hard_limit"] == 40 + + +def test_repo_budget_includes_fresh_provider_usage_without_project_limit() -> None: + guard = AIBudgetGuard({"monthly_budgets": {"openai:default:repo-a:review": {"user_monthly_budget_usd": 10}}}) + decision = guard.preflight( + task_class="review", provider="openai", provider_scope="default", repo="repo-a", + estimated_cost_usd=1, usage_snapshot={"updated_at": time.time(), "used_usd": 10}, + ) + assert decision["decision"] == "defer" + + +def test_nested_budget_scope_resolves_repo_and_task_class() -> None: + guard = AIBudgetGuard({"monthly_budgets": {"openai": {"project": {"o/r": { + "research": {"user_monthly_budget_usd": 12}, + }}}}}) + decision = guard.preflight( + task_class="research", provider="openai", provider_scope="project", repo="o/r", estimated_cost_usd=1, + usage_snapshot={"updated_at": time.time(), "used_usd": 0}, + ) + assert decision["decision"] == "allow" + assert decision["hard_limit"] == 9.6 + + +def test_sqlite_ledger_shares_atomic_reservations_between_guard_instances(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setenv("CODEX_AUDIT_SERVICE_AI_BUDGET_LEDGER_PATH", str(tmp_path / "ledger.sqlite3")) + config = {"monthly_budgets": {"openai": {"user_monthly_budget_usd": 10}}} + first_guard = AIBudgetGuard(config) + second_guard = AIBudgetGuard(config) + snapshot = {"updated_at": time.time(), "used_usd": 0} + first = first_guard.preflight(task_class="review", provider="openai", estimated_cost_usd=6, usage_snapshot=snapshot) + second = second_guard.preflight(task_class="review", provider="openai", estimated_cost_usd=6, usage_snapshot=snapshot) + assert first_guard.reserve(first, 6) is not None + assert second_guard.reserve(second, 6) is None + + +def test_stale_usage_fails_closed_and_does_not_fallback() -> None: + guard = AIBudgetGuard({"monthly_budgets": {"openai": {"user_monthly_budget_usd": 10}}}) + decision = guard.preflight( + task_class="auto_fix", provider="openai", estimated_cost_usd=1, + usage_snapshot={"updated_at": time.time() - 100_000, "used_usd": 0}, + ) + assert decision["decision"] == "block" + assert decision["auto_fallback_allowed"] is False + + +def test_malformed_usage_and_rate_limit_values_fail_closed() -> None: + guard = AIBudgetGuard({"monthly_budgets": {"openai": {"user_monthly_budget_usd": 10}}}) + usage = guard.preflight( + task_class="review", provider="openai", estimated_cost_usd=1, + usage_snapshot={"updated_at": time.time(), "used_usd": "N/A"}, + ) + assert usage["decision"] == "block" + codex = guard.preflight( + task_class="review", provider="codex", + codex_snapshot={"updated_at": time.time(), "rate_limits": { + "primary": {"remaining_percent": "unknown"}, + "secondary": {"remaining_percent": 80}, + }}, + ) + assert codex["decision"] == "defer" + + +def test_codex_uses_tightest_window_and_keeps_reserve() -> None: + guard = AIBudgetGuard() + assert guard.preflight(task_class="research", provider="codex", codex_snapshot=_codex_snapshot(35, 90))["decision"] == "allow" + assert guard.preflight(task_class="research", provider="codex", codex_snapshot=_codex_snapshot(29, 90))["decision"] == "defer" + assert guard.preflight(task_class="incident", provider="codex", codex_snapshot=_codex_snapshot(11, 90))["decision"] == "defer" + assert guard.preflight(task_class="incident", provider="codex", codex_snapshot=_codex_snapshot(19, 90))["remaining_after_reservation"] == 0.09 + + +def test_codex_reservation_uses_live_headroom_shared_provider_scope() -> None: + guard = AIBudgetGuard() + snapshot = _codex_snapshot(21, 90) + first = guard.preflight(task_class="review", provider="codex", provider_scope="shared", repo="o/a", codex_snapshot=snapshot) + second = guard.preflight(task_class="review", provider="codex", provider_scope="shared", repo="o/b", codex_snapshot=snapshot) + assert guard.reserve(first, 0.10) is not None + assert guard.reserve(second, 0.10) is None + + +def test_codex_settled_usage_remains_reserved_until_reset() -> None: + guard = AIBudgetGuard() + snapshot = _codex_snapshot(50, 90, reset_at=time.time() + 3600) + for _ in range(3): + decision = guard.preflight(task_class="review", provider="codex", provider_scope="shared", codex_snapshot=snapshot) + reservation = guard.reserve(decision, 0.10) + assert reservation is not None + assert guard.settle(reservation, 0.10) + follow_up = guard.preflight(task_class="review", provider="codex", provider_scope="shared", codex_snapshot=snapshot) + assert follow_up["decision"] == "defer" + + +def test_missing_codex_snapshot_defers_research_and_auto_fix() -> None: + guard = AIBudgetGuard() + for task in ("research", "auto_fix"): + decision = guard.preflight(task_class=task, provider="codex") + assert decision["decision"] == "defer" + assert decision["reason_codes"] == ["codex_rate_limit_snapshot_unavailable"] + + +def test_reservations_are_atomic_and_released_or_settled() -> None: + guard = AIBudgetGuard({"monthly_budgets": {"openai": {"user_monthly_budget_usd": 10}}}) + kwargs = dict(task_class="review", provider="openai", estimated_cost_usd=6, + usage_snapshot={"updated_at": time.time(), "used_usd": 0}) + first = guard.preflight(**kwargs) + second = guard.preflight(**kwargs) + one = guard.reserve(first, 6) + two = guard.reserve(second, 6) + assert one is not None + assert two is None + assert guard.settle(one, 4) + assert guard.release(one) is False + + +def test_settled_spend_remains_reserved_until_usage_snapshot_catches_up() -> None: + guard = AIBudgetGuard({"monthly_budgets": {"openai": {"user_monthly_budget_usd": 10}}}) + snapshot = {"updated_at": time.time(), "used_usd": 0} + decision = guard.preflight(task_class="review", provider="openai", estimated_cost_usd=6, usage_snapshot=snapshot) + reservation = guard.reserve(decision, 6) + assert reservation is not None + assert guard.settle(reservation, 6) + follow_up = guard.preflight(task_class="review", provider="openai", estimated_cost_usd=3, usage_snapshot=snapshot) + assert follow_up["decision"] == "defer" + + +def test_settled_delta_is_not_cleared_by_historical_provider_usage() -> None: + guard = AIBudgetGuard({"monthly_budgets": {"openai": {"user_monthly_budget_usd": 200, "provider_project_limit_usd": 200}}}) + snapshot = {"updated_at": time.time(), "used_usd": 100} + decision = guard.preflight(task_class="review", provider="openai", estimated_cost_usd=50, usage_snapshot=snapshot) + reservation = guard.reserve(decision, 50) + assert reservation is not None + assert guard.settle(reservation, 10) + follow_up = guard.preflight(task_class="review", provider="openai", estimated_cost_usd=55, usage_snapshot=snapshot) + assert follow_up["decision"] == "defer" + + +def test_sqlite_ledger_shares_reservations_across_guard_instances() -> None: + with TemporaryDirectory() as tmp: + config = {"ledger_path": str(Path(tmp) / "budget.sqlite3"), "monthly_budgets": {"openai": {"user_monthly_budget_usd": 10}}} + snapshot = {"updated_at": time.time(), "used_usd": 0} + first_guard = AIBudgetGuard(config) + first = first_guard.preflight(task_class="review", provider="openai", estimated_cost_usd=6, usage_snapshot=snapshot) + reservation = first_guard.reserve(first, 6) + assert reservation is not None + + second_guard = AIBudgetGuard(config) + second = second_guard.preflight(task_class="review", provider="openai", estimated_cost_usd=6, usage_snapshot=snapshot) + assert second_guard.reserve(second, 6) is None + assert second_guard.settle(reservation.reservation_id, 6) + + +def test_uncertain_reservation_survives_restart_until_explicit_reconciliation() -> None: + with TemporaryDirectory() as tmp: + config = { + "ledger_path": str(Path(tmp) / "budget.sqlite3"), + "monthly_budgets": {"openai": {"user_monthly_budget_usd": 10}}, + } + snapshot = {"updated_at": time.time(), "used_usd": 0} + first_guard = AIBudgetGuard(config) + decision = first_guard.preflight(task_class="review", provider="openai", estimated_cost_usd=8, usage_snapshot=snapshot) + reservation = first_guard.reserve(decision, 8) + assert reservation is not None + assert first_guard.mark_uncertain(reservation) + + restarted_guard = AIBudgetGuard(config) + blocked = restarted_guard.preflight(task_class="review", provider="openai", estimated_cost_usd=0.1, usage_snapshot=snapshot) + assert blocked["decision"] == "defer" + assert not restarted_guard.reconcile_pending_uncertain(reservation.reservation_id, dispatched=None) + assert restarted_guard.reconcile_pending_uncertain(reservation.reservation_id, dispatched=False) + + +def test_codex_reset_only_clears_matching_account_scope() -> None: + now = time.time() + guard = AIBudgetGuard(clock=lambda: now) + for scope in ("account-a", "account-b"): + decision = guard.preflight( + task_class="review", provider="codex", provider_scope=scope, + codex_snapshot=_codex_snapshot(50, 80, observed_at=now, reset_at=now + 60), + ) + reservation = guard.reserve(decision, 0.10) + assert reservation is not None + assert guard.settle(reservation, 0.10) + + guard.preflight( + task_class="review", provider="codex", provider_scope="account-a", + codex_snapshot=_codex_snapshot(50, 80, observed_at=now, reset_at=now), + ) + account_a = guard.preflight( + task_class="review", provider="codex", provider_scope="account-a", + codex_snapshot=_codex_snapshot(50, 80, observed_at=now, reset_at=now), + ) + account_b = guard.preflight( + task_class="review", provider="codex", provider_scope="account-b", + codex_snapshot=_codex_snapshot(50, 80, observed_at=now, reset_at=now + 60), + ) + assert account_a["reserved_usage"] == 0 + assert account_b["reserved_usage"] == 0.1 + + +def test_month_period_is_explicit_and_fallback_requires_human_approval() -> None: + guard = AIBudgetGuard({"billing_timezone": "UTC", "monthly_budgets": {"openai": {"user_monthly_budget_usd": 10}}}) + decision = guard.preflight( + task_class="maintenance", provider="openai", estimated_cost_usd=1, + usage_snapshot={"updated_at": time.time(), "organization_costs": {"total_cost": 0}}, + ) + assert len(decision["period"]) == 7 + assert decision["billing_timezone"] == "UTC" + assert decision["auto_fallback_allowed"] is False diff --git a/tests/test_ai_gateway_service_job_recovery.py b/tests/test_ai_gateway_service_job_recovery.py index e418359a..ea792ef8 100644 --- a/tests/test_ai_gateway_service_job_recovery.py +++ b/tests/test_ai_gateway_service_job_recovery.py @@ -39,6 +39,32 @@ def test_restart_marks_orphaned_active_jobs_failed(self) -> None: record_automation_run.assert_has_calls([call(queued), call(running)], any_order=True) self.assertEqual(record_automation_run.call_count, 2) + def test_restart_keeps_ambiguous_dispatch_reservation_pending(self) -> None: + now = time.time() + job = { + "job_id": "d" * 24, + "status": "running", + "created_at": now, + "updated_at": now, + "_budget_reservation_id": "reservation-1", + "dispatch_state": "pending_uncertain", + } + with tempfile.TemporaryDirectory() as tmp: + with patch.dict(gateway.os.environ, {"CODEX_AUDIT_SERVICE_JOB_DIR": tmp}, clear=False): + gateway._write_job(job) + with ( + patch.object(gateway, "_record_job_automation_run"), + patch.object(gateway, "_audit_log"), + patch.object(gateway, "_mark_budget_reservation_uncertain") as mark_uncertain, + patch.object(gateway, "_settle_budget_reservation") as settle, + patch.object(gateway, "_release_budget_reservation") as release, + ): + self.assertEqual(gateway._recover_orphaned_jobs(), 1) + + mark_uncertain.assert_called_once() + settle.assert_not_called() + release.assert_not_called() + if __name__ == "__main__": unittest.main() diff --git a/tests/test_anthropic_admin_usage.py b/tests/test_anthropic_admin_usage.py index 6e572b4a..d617f8ab 100644 --- a/tests/test_anthropic_admin_usage.py +++ b/tests/test_anthropic_admin_usage.py @@ -5,10 +5,11 @@ import json import os import unittest +from datetime import UTC, datetime from urllib.parse import parse_qs, urlparse from unittest.mock import patch -from service.anthropic_admin_usage import read_anthropic_admin_usage +from service.anthropic_admin_usage import _usage_window, read_anthropic_admin_usage class _FakeResponse: @@ -26,6 +27,13 @@ def read(self) -> bytes: class TestAnthropicAdminUsage(unittest.TestCase): + def test_default_window_uses_configured_billing_timezone(self) -> None: + end_time = int(datetime(2026, 7, 31, 18, tzinfo=UTC).timestamp()) + start_time, days = _usage_window(end_time, "Asia/Shanghai") + + self.assertEqual(start_time, end_time - 2 * 3600) + self.assertEqual(days, 1) + def test_disabled_without_admin_key(self) -> None: with patch.dict(os.environ, {}, clear=True): self.assertIsNone(read_anthropic_admin_usage(now=1783139561, timeout_seconds=1)) diff --git a/tests/test_automation_decision.py b/tests/test_automation_decision.py index c4d7ef98..6c5d78bb 100644 --- a/tests/test_automation_decision.py +++ b/tests/test_automation_decision.py @@ -157,6 +157,24 @@ def test_low_quota_recommends_low_cost_model(self) -> None: self.assertEqual(result["effective_provider"], "openai") self.assertEqual(result["effective_model"], "gpt-5.4-mini") + def test_budget_defer_never_selects_paid_fallback(self) -> None: + result = decide_automation_execution( + repo="org/repo", + requested_provider="codex", + requested_model="codex-cli", + quota_status="ok", + control_action=CONTROL_CONTINUE, + ai_budget_decision={ + "schema": "ai_budget_decision.v1", + "decision": "defer", + "reason_codes": ["codex_rate_limit_below_task_threshold"], + }, + ) + self.assertEqual(result["action"], EXECUTION_DEFER) + self.assertTrue(result["deferred_budget"]) + self.assertNotEqual(result["effective_provider"], "openai") + self.assertTrue(any("deferred_budget" in reason for reason in result["reasons"])) + def test_low_quota_overrides_requested_expensive_model(self) -> None: result = decide_automation_execution( repo="QuantStrategyLab/AIAuditBridge", diff --git a/tests/test_codex_adapter.py b/tests/test_codex_adapter.py new file mode 100644 index 00000000..740d783d --- /dev/null +++ b/tests/test_codex_adapter.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import subprocess +import unittest +from unittest.mock import Mock, patch + +from service.adapters.codex_adapter import CodexAdapter + + +class CodexAdapterDispatchTests(unittest.TestCase): + def test_timeout_is_pending_uncertain_not_confirmed_dispatch(self) -> None: + with ( + patch("service.adapters.codex_adapter.shutil.which", return_value="/usr/bin/codex"), + patch("service.adapters.codex_adapter.subprocess.run", side_effect=subprocess.TimeoutExpired("codex", 1)), + ): + result = CodexAdapter().execute(prompt="review", timeout=1) + + self.assertFalse(result.dispatch_started) + self.assertTrue(result.dispatch_uncertain) + + def test_nonzero_exit_is_pending_uncertain_not_confirmed_dispatch(self) -> None: + completed = Mock(returncode=1, stdout="upstream request interrupted", stderr="") + with ( + patch("service.adapters.codex_adapter.shutil.which", return_value="/usr/bin/codex"), + patch("service.adapters.codex_adapter.subprocess.run", return_value=completed), + ): + result = CodexAdapter().execute(prompt="review") + + self.assertFalse(result.dispatch_started) + self.assertTrue(result.dispatch_uncertain) + + def test_known_local_nonzero_exit_is_not_dispatched(self) -> None: + completed = Mock(returncode=2, stdout="", stderr="unknown option --bad-flag") + with ( + patch("service.adapters.codex_adapter.shutil.which", return_value="/usr/bin/codex"), + patch("service.adapters.codex_adapter.subprocess.run", return_value=completed), + ): + result = CodexAdapter().execute(prompt="review") + + self.assertFalse(result.dispatch_started) + self.assertFalse(result.dispatch_uncertain) + + def test_prelaunch_command_failure_is_not_dispatched(self) -> None: + with patch( + "service.adapters.codex_adapter._codex_command", + side_effect=RuntimeError("codex CLI was not found on the service host"), + ): + result = CodexAdapter().execute(prompt="review") + + self.assertFalse(result.dispatch_started) + self.assertFalse(result.dispatch_uncertain) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_llm_adapter.py b/tests/test_llm_adapter.py index 91549cf9..58af41c8 100644 --- a/tests/test_llm_adapter.py +++ b/tests/test_llm_adapter.py @@ -17,6 +17,7 @@ def test_complete_returns_empty_output_on_provider_failure(self) -> None: self.assertFalse(result.success) self.assertEqual(result.output, "") self.assertEqual(result.error, "provider unavailable") + self.assertTrue(result.dispatch_uncertain) def test_parallel_review_returns_empty_output_on_worker_failure(self) -> None: with patch.object(LlmAdapter, "complete", side_effect=RuntimeError("worker failed")): @@ -29,6 +30,17 @@ def test_parallel_review_returns_empty_output_on_worker_failure(self) -> None: self.assertFalse(results[0].success) self.assertEqual(results[0].output, "") self.assertEqual(results[0].error, "worker failed") + self.assertFalse(results[0].dispatch_uncertain) + + def test_provider_rejection_is_not_ambiguous_dispatch(self) -> None: + with patch( + "service.adapters.llm_adapter._openai_completion", + side_effect=LlmAdapterError("OpenAI HTTP 401: invalid API key"), + ): + result = LlmAdapter().complete(model="gpt-5.4-mini", user="review") + + self.assertFalse(result.dispatch_started) + self.assertFalse(result.dispatch_uncertain) if __name__ == "__main__": diff --git a/tests/test_openai_admin_usage.py b/tests/test_openai_admin_usage.py index d4924012..95123d83 100644 --- a/tests/test_openai_admin_usage.py +++ b/tests/test_openai_admin_usage.py @@ -5,9 +5,10 @@ import json import os import unittest +from datetime import UTC, datetime from unittest.mock import patch -from service.openai_admin_usage import read_openai_admin_usage +from service.openai_admin_usage import _usage_window, read_openai_admin_usage class _FakeResponse: @@ -25,6 +26,13 @@ def read(self) -> bytes: class TestOpenAIAdminUsage(unittest.TestCase): + def test_default_window_uses_configured_billing_timezone(self) -> None: + end_time = int(datetime(2026, 7, 31, 18, tzinfo=UTC).timestamp()) + start_time, days = _usage_window(end_time, "Asia/Shanghai") + + self.assertEqual(start_time, end_time - 2 * 3600) + self.assertEqual(days, 1) + def test_disabled_without_admin_key(self) -> None: with patch.dict(os.environ, {}, clear=True): self.assertIsNone(read_openai_admin_usage(now=1783139561, timeout_seconds=1)) diff --git a/tests/test_quota.py b/tests/test_quota.py index 84bb2bc9..2ec7ddb4 100644 --- a/tests/test_quota.py +++ b/tests/test_quota.py @@ -162,6 +162,21 @@ def test_legacy_record_starts_a_fresh_weekly_budget_window(self) -> None: class TestQuotaManager(unittest.TestCase): + def test_recording_failure_requires_explicit_recovery(self) -> None: + manager = QuotaManager() + manager.mark_recording_failed("repo/recovery") + self.assertFalse(manager.check("repo/recovery", "gpt-5.4-mini")["allowed"]) + manager.recover_recording_failure("repo/recovery") + self.assertTrue(manager.check("repo/recovery", "gpt-5.4-mini")["allowed"]) + + def test_recording_failure_persists_across_restart(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + store = str(Path(tmp) / "quota.json") + with patch.dict(os.environ, {"CODEX_AUDIT_SERVICE_QUOTA_STORE": store}, clear=False): + manager = QuotaManager() + manager.mark_recording_failed("repo/persisted") + restarted = QuotaManager() + self.assertFalse(restarted.check("repo/persisted", "gpt-5.4-mini")["allowed"]) """QuotaManager budget tracking and enforcement.""" def setUp(self) -> None: @@ -195,6 +210,14 @@ def test_only_trusted_codex_account_checks_ignore_api_budget(self) -> None: self.manager._records["test/repo"] = QuotaRecord( repo="test/repo", total_cost_usd=0.05, api_key_cost_usd=0.05 ) + self.manager._codex_account_cache = { + "updated_at": time.time(), + "rate_limits": { + "primary": {"remaining_percent": 80}, + "secondary": {"remaining_percent": 80}, + }, + } + self.manager._codex_account_cache_ts = time.time() untrusted_result = self.manager.check("test/repo", "codex-cli", "review this pull request") trusted_result = self.manager.check( @@ -279,6 +302,15 @@ def test_status_summary_can_include_live_openai_account_snapshot(self) -> None: status = self.manager.status() self.assertEqual(status["summary"]["openai_account"], snapshot) + def test_openai_snapshot_uses_budget_guard_billing_timezone(self) -> None: + guard = type("Guard", (), {"billing_timezone": "Asia/Shanghai"})() + with ( + patch("service.quota.read_openai_admin_usage", return_value=None) as read_snapshot, + patch("service.ai_budget_guard.get_ai_budget_guard", return_value=guard), + ): + self.manager._openai_account_snapshot(timeout_seconds=1) + read_snapshot.assert_called_once_with(timeout_seconds=1, billing_timezone="Asia/Shanghai") + def test_status_summary_can_include_live_anthropic_account_snapshot(self) -> None: snapshot = {"source": "anthropic_admin_api", "status": "available", "costs": {"total_cost": 1.23}} with patch("service.quota.read_anthropic_admin_usage", return_value=snapshot): @@ -298,7 +330,7 @@ def test_anthropic_account_failures_are_negative_cached(self) -> None: self.assertEqual(read_snapshot.call_count, 1) def test_account_snapshot_reads_use_shared_status_timeout(self) -> None: - def slow_snapshot(timeout_seconds: float | None = None) -> dict[str, object]: + def slow_snapshot(timeout_seconds: float | None = None, billing_timezone: str = "UTC") -> dict[str, object]: time.sleep(timeout_seconds or 0.25) return {"source": "slow", "status": "available"} @@ -321,7 +353,7 @@ def slow_snapshot(timeout_seconds: float | None = None) -> dict[str, object]: def test_openai_account_snapshot_refresh_is_single_flight(self) -> None: snapshot = {"source": "openai_admin_api", "status": "available"} - def slow_snapshot(timeout_seconds: float | None = None) -> dict[str, object]: + def slow_snapshot(timeout_seconds: float | None = None, billing_timezone: str = "UTC") -> dict[str, object]: time.sleep(0.05) return snapshot diff --git a/tests/test_run_codex_pr_review.py b/tests/test_run_codex_pr_review.py index 768439ac..8b591458 100644 --- a/tests/test_run_codex_pr_review.py +++ b/tests/test_run_codex_pr_review.py @@ -534,6 +534,7 @@ def test_direct_api_runs_when_service_url_is_unset(self) -> None: { "OPENAI_API_KEY": "test-key", "CODEX_PR_REVIEW_API_FALLBACK_ENABLED": "true", + "CODEX_PR_REVIEW_DIRECT_API_PRIMARY_ENABLED": "true", }, clear=True, ), @@ -1387,7 +1388,14 @@ def test_oidc_repo_not_allowed_counts_as_unconfigured_backend(self) -> None: def test_service_exec_failure_falls_back_to_direct_api(self) -> None: with ( - patch.dict(os.environ, {"CODEX_AUDIT_SERVICE_URL": "https://service.example"}, clear=True), + patch.dict( + os.environ, + { + "CODEX_AUDIT_SERVICE_URL": "https://service.example", + "CODEX_PR_REVIEW_API_FALLBACK_ENABLED": "true", + }, + clear=True, + ), patch( "scripts.run_codex_pr_review.run_codex_service_review", side_effect=ReviewError("Codex service job failed: codex exec failed (rc=1): boom"),