Enforce per-project cost budgets with a hard-stop guardrail

[lucastononro]

Summary

Cost was fully observable after the fact (services/usage.py, live CostBadge) but nothing capped it — an orchestrator fanning out to specialists could silently rack up spend. This adds a per-project USD budget that the runner checks against accumulated UsageEvent rows and hard-stops the agent when exceeded, landing the session in a clean budget_exceeded terminal state (not failed) with a clear message telling the user how to raise or clear the cap.

The budget is project-scoped on purpose: sub-agents spawned via delegate-task share the session/project, so the whole fan-out is capped as one unit (spend includes both LLM tokens and sandbox compute).

Changes

Backend

• models.py / db.py: new nullable projects.budget_usd column (NULL = uncapped) + idempotent dev-DB migration; surfaced in Project.to_dict().
• schemas.py / routers/projects.py: budget_usd on ProjectCreate/ProjectUpdate (0 ≤ x ≤ 1M). PATCH distinguishes "omitted" from "explicit null clears the cap" via model_fields_set.
• services/budget.py (new): BudgetStatus (budget vs. project-wide spend summed over usage_events.cost_usd), get_budget_status() (by project or session), and check_budget() which raises BudgetExceededError.
• services/agent/runner.py:
  • Pre-run check — a project already over its cap never starts an agent (no provider call at all).
  • Per-usage-event check — after every recorded LLM usage event, the runner re-checks; crossing the cap mid-run unwinds the provider loop immediately (both the Claude/MCP path and the runner-managed tool loop go through the same _record_usage).
  • run_agent catches BudgetExceededError and publishes a budget_exceeded event (message + spent/cap amounts) plus a state_change to the clean budget_exceeded terminal state.
• routers/usage.py: /sessions/{id}/usage and /projects/{id}/usage now return a budget block (budget_usd, spent_usd, remaining_usd, exceeded) for the UI.

Frontend (minimal)

• Project Settings modal: new "Budget — spend cap (USD)" field (empty = no limit); saved through the existing PATCH /projects/{id} path in Sidebar.
• CostBadge: shows project spend vs. cap with a progress bar; flips to a red "over budget" state when exceeded.
• page.tsx (kept tiny — 5 small edits): budgetInfo state hydrated from session usage, a budget_exceeded SSE case (error message + stop running indicator), budget_exceeded added to the terminal-state list, prop pass-through to CostBadge. ⚠️ Overlaps with Break up the 4155-line page.tsx monolith #97 (page.tsx split) — these edits are small and localized (SSE switch + usage hydration block) but will need a trivial port if Break up the 4155-line page.tsx monolith #97 lands first.

Tests — backend/tests/test_budget.py (8 new):

• Budget math (exceeded / remaining / uncapped / under-budget / orphan session).
• API: PATCH set → untouched-by-other-PATCH → explicit-null clear → negative rejected; usage endpoints expose the budget block.
• Hard-stop: a session whose project has fake UsageEvents over the cap never drives the provider and persists budget_exceeded (and is not marked failed); a run that crosses the cap mid-flight is halted at the next usage event (events after the tripping usage event are never consumed).

Test plan

• Existing behavior: no budget set (budget_usd NULL) → check_budget is a no-op, runs behave exactly as today, cost still displays in CostBadge. Full backend suite: 304 passed, 8 skipped (was 296 before; +8 new). ruff check + ruff format --check clean.
• New behavior: set a low budget in Project Settings → agent halts with "Budget limit reached: this project has spent $X of its $Y cap…" as soon as spend crosses it (or refuses to start if already over); badge shows the red "over budget" state. Covered by the two runner hard-stop tests; frontend npm ci + next lint + tsc --noEmit + prettier --check + next build all green.

Follow-ups (out of scope)

• Token-count ceiling in addition to USD (the check point in services/budget.py makes this a small extension).
• Per-session budgets layered under the project cap.
• Pause-and-ask via the clarifications mechanism (resumable prompt instead of terminal state) — terminal state chosen here because a paused agent still holds a provider subprocess open indefinitely.
• Sandbox usage events trigger the halt only at the next LLM call (the check hooks LLM usage recording); a very long single sandbox run isn't interrupted mid-execution.

Closes #107

🤖 Generated with Claude Code

https://claude.ai/code/session_01Wp66uSq9saSpRq9Bbmjxj4

Greptile Summary

This PR introduces a per-project USD budget guardrail that hard-stops agents when accumulated spend (usage_events.cost_usd, LLM + sandbox compute) crosses a configurable cap. The backend adds a nullable budget_usd column on projects, a new services/budget.py module with BudgetStatus / check_budget, and hook points in the agent runner (pre-run and post-usage-event). The frontend adds a spend-cap field in Project Settings and updates CostBadge with a progress bar and an "over budget" state.

• services/budget.py resolves a session → project via the legacy experiment back-pointer and sums all UsageEvent.cost_usd rows; check_budget raises BudgetExceededError when spend meets or exceeds the cap.
• Runner integration: _check_budget_failopen (new wrapper) is called twice — once before the provider loop starts and once after every recorded usage event. BudgetExceededError is caught in run_agent and lands the session in a clean budget_exceeded terminal state with a descriptive user-facing message.
• Eight new tests cover budget math, API surface (PATCH set/clear/negative rejected, usage endpoint budget block), and two runner hard-stop scenarios (before-start and mid-run).

Confidence Score: 4/5

Safe to merge with one fix: a $0 budget silently blocks all agents immediately, before any spend occurs.

The only new defect introduced is the >= comparison in BudgetStatus.exceeded combined with the schema allowing budget_usd=0. A user who enters 0 in the settings field gets all agents hard-stopped with a "spent $0.00 of $0.00" message before any work happens. Everything else — the fail-open wrapper, the pre-run and mid-run check hooks, the clean terminal state, the frontend SSE handling — is implemented correctly and covered by tests.

backend/services/budget.py (the exceeded property) and backend/schemas.py (the ge=0.0 constraint) should be reviewed together for the $0-budget edge case.

Important Files Changed

Filename	Overview
backend/services/budget.py	New budget service: two separate SELECT queries, no explicit DB transaction (already flagged in previous thread); exceeded uses >= which means a budget_usd=0 project is immediately exceeded before any spend.
backend/services/agent/runner.py	Budget check correctly integrated via _check_budget_failopen wrapper; pre-run and per-usage-event hooks both propagate only BudgetExceededError, swallowing other exceptions (fail-open). BudgetExceededError handled cleanly before the generic Exception catch.
backend/schemas.py	Schema allows budget_usd=0.0 via ge=0.0; combined with the >= comparison in BudgetStatus.exceeded, a $0 budget is immediately exceeded even with $0 spend.
backend/tests/test_budget.py	Eight tests cover math, API surface, and runner hard-stop (pre-start and mid-run). Fail-open behavior is also verified with a dedicated test.
frontend/src/components/ProjectSettingsModal.tsx	Budget field added; empty field = null (no cap), numeric value parsed via parseFloat. Negative input silently sends null (already flagged in previous thread).

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant UI as Frontend
    participant API as projects / usage routers
    participant Runner as agent runner
    participant Budget as services/budget.py
    participant DB as SQLite (usage_events, projects)

    UI->>API: "PATCH /projects/{id} {budget_usd: 5.0}"
    API->>DB: "UPDATE projects SET budget_usd=5.0"
    API-->>UI: "200 {budget_usd: 5.0}"

    UI->>Runner: run_agent(session_id)
    Runner->>Budget: _check_budget_failopen(session_id) [pre-run]
    Budget->>DB: "SELECT budget_usd FROM projects WHERE id=pid"
    Budget->>DB: "SELECT SUM(cost_usd) FROM usage_events WHERE project_id=pid"
    DB-->>Budget: "budget=5.0, spent=2.0"
    Budget-->>Runner: OK (not exceeded)

    Runner->>Runner: _drive_provider() — LLM call
    Runner->>DB: record_llm_usage (new UsageEvent)
    Runner->>Budget: _check_budget_failopen(session_id) [post-usage]
    Budget->>DB: SELECT budget_usd / SUM(cost_usd)
    DB-->>Budget: "budget=5.0, spent=6.1"
    Budget-->>Runner: raises BudgetExceededError

    Runner->>UI: "SSE budget_exceeded {error, spent_usd, budget_usd}"
    Runner->>UI: "SSE state_change {state: budget_exceeded}"

    UI->>API: "GET /sessions/{id}/usage"
    API->>Budget: get_budget_status(session_id)
    Budget->>DB: SELECT budget_usd / SUM(cost_usd)
    API-->>UI: "{summary, budget: {exceeded: true, spent_usd: 6.1, budget_usd: 5.0}}"

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant UI as Frontend
    participant API as projects / usage routers
    participant Runner as agent runner
    participant Budget as services/budget.py
    participant DB as SQLite (usage_events, projects)

    UI->>API: "PATCH /projects/{id} {budget_usd: 5.0}"
    API->>DB: "UPDATE projects SET budget_usd=5.0"
    API-->>UI: "200 {budget_usd: 5.0}"

    UI->>Runner: run_agent(session_id)
    Runner->>Budget: _check_budget_failopen(session_id) [pre-run]
    Budget->>DB: "SELECT budget_usd FROM projects WHERE id=pid"
    Budget->>DB: "SELECT SUM(cost_usd) FROM usage_events WHERE project_id=pid"
    DB-->>Budget: "budget=5.0, spent=2.0"
    Budget-->>Runner: OK (not exceeded)

    Runner->>Runner: _drive_provider() — LLM call
    Runner->>DB: record_llm_usage (new UsageEvent)
    Runner->>Budget: _check_budget_failopen(session_id) [post-usage]
    Budget->>DB: SELECT budget_usd / SUM(cost_usd)
    DB-->>Budget: "budget=5.0, spent=6.1"
    Budget-->>Runner: raises BudgetExceededError

    Runner->>UI: "SSE budget_exceeded {error, spent_usd, budget_usd}"
    Runner->>UI: "SSE state_change {state: budget_exceeded}"

    UI->>API: "GET /sessions/{id}/usage"
    API->>Budget: get_budget_status(session_id)
    Budget->>DB: SELECT budget_usd / SUM(cost_usd)
    API-->>UI: "{summary, budget: {exceeded: true, spent_usd: 6.1, budget_usd: 5.0}}"

Loading

Comments Outside Diff (1)

1. frontend/src/components/ProjectSettingsModal.tsx, line 881-888 (link)

   Negative budget input silently clears the cap

   Number.isFinite(parsed) && parsed >= 0 correctly rejects negative numbers, but the fallback is null (no cap), not an error. A user who accidentally types -5 would see no validation message — the cap would simply be cleared. The backend validator (ge=0.0) will also reject negatives, but only via a 422 response that is swallowed by the catch block in handleSaveSandboxConfig. Consider either showing an inline validation message when the parsed value is negative, or clamping with Math.max(0, parsed) and a visible hint.

   

_{Reviews (2): Last reviewed commit: "fix(review): address Greptile findings o..." | Re-trigger Greptile}

[greptile-apps]

Two separate queries, no atomicity guarantee

get_budget_status runs two independent SELECTs — one for projects.budget_usd and one for SUM(usage_events.cost_usd) — within the same async with async_session() block but not inside a single database transaction. A concurrent agent that records a usage event between these two queries could make the spend sum appear slightly lower than reality, potentially allowing one more LLM round before the cap fires.

Suggested change

	async with async_session() as db:
	pid = project_id
	if pid is None and session_id:
	pid = await _project_id_for_session(db, session_id)
	if not pid:
	return None
	row = await db.execute(select(Project.budget_usd).where(Project.id == pid))
	budget = row.scalar_one_or_none()
	row = await db.execute(
	select(func.coalesce(func.sum(UsageEvent.cost_usd), 0.0)).where(
	UsageEvent.project_id == pid
	)
	)
	spent = float(row.scalar_one() or 0.0)
	async with async_session() as db:
	pid = project_id
	if pid is None and session_id:
	pid = await _project_id_for_session(db, session_id)
	if not pid:
	return None
	# Read both values inside an explicit transaction to minimise the
	# window where a concurrent usage event could be recorded between
	# the two queries and make the spend appear lower than reality.
	await db.execute(text("BEGIN"))
	row = await db.execute(select(Project.budget_usd).where(Project.id == pid))
	budget = row.scalar_one_or_none()
	row = await db.execute(
	select(func.coalesce(func.sum(UsageEvent.cost_usd), 0.0)).where(
	UsageEvent.project_id == pid
	)
	)
	spent = float(row.scalar_one() or 0.0)

[greptile-apps]

Too many files changed for review. (101 files found, 100 file limit)

Bypass the limit by tagging @greptile-apps to review.

[lucastononro]

Merged into integration branch staging-v0.0.5 (see the STAGING-v0.0.5.md ledger on that branch for merge order, Greptile follow-ups and test results). These changes will land on main via the staging merge — closing to clear the queue.