From fc90745b15eb4f5edc98533d0915fe42c2c72560 Mon Sep 17 00:00:00 2001 From: Arik Levinsky Date: Mon, 20 Jul 2026 11:32:49 -0600 Subject: [PATCH] fix(chargeback): fail explicitly on a budget-less rollup root; bound read-route inputs (#97, #107) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #97: spend_rollup sums a node's OWN budget ledger (INNER join), which already aggregates everything beneath it (ADR 0033) — but only when the node carries a budget. A budget-less root accrues no ledger rows of its own, so the aggregate came back empty and /v1/spend silently reported zero spend while /v1/budgets listed the subtree's real committed spend. It now raises BudgetNotFound (403, already documented on the read routes) instead of a misleading [] — a cheap keyed existence check distinguishes "no budget at this node" from "budget exists, zero spend" (still []). Rolling the whole subtree up instead would double-count spend shared across ancestor nodes (the very thing the single-node design avoids), so the explicit precondition is the safe correct fix. #107 (read-route bounds): scope_id and the group_by label key were taken from unbounded query values straight into a WHERE / the rollup grouping; both are now bounded (MAX_STR_LEN, MAX_LABEL_KEY_LEN) with a 422, via a new _parse_group_by_query helper. _LabelKey gains min_length=1 so an empty-string label key is rejected on reserve/meter, matching the other command-path bounds. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01BjJ4QScUKbrMzF1jujL5Bg --- .../adapters/postgres/chargeback_repo.py | 23 ++++++++++- src/tollgate/api/routes/chargeback.py | 38 +++++++++++++++---- src/tollgate/api/schemas.py | 2 +- .../integration/test_chargeback_spend_repo.py | 12 ++++-- tests/unit/test_api_schemas.py | 13 +++++++ tests/unit/test_chargeback_route.py | 35 ++++++++++++++++- 6 files changed, 108 insertions(+), 15 deletions(-) diff --git a/src/tollgate/adapters/postgres/chargeback_repo.py b/src/tollgate/adapters/postgres/chargeback_repo.py index bb08a24..1900716 100644 --- a/src/tollgate/adapters/postgres/chargeback_repo.py +++ b/src/tollgate/adapters/postgres/chargeback_repo.py @@ -8,7 +8,9 @@ ``resolve_scope_ancestry`` returns a node's server-derived ancestry map for the authorization check (section 5.0). ``spend_rollup`` sums one node's own budget ledger rows for a period, grouped by provider, model, or a label key -- joined on ``budget`` for that single node so a -subtree union never double-counts shared spend (section 2). ``PostgresChargebackReader`` hands +subtree union never double-counts shared spend (section 2); a node that carries no budget of its +own has no such ledger rows, so it raises ``BudgetNotFound`` rather than silently reporting zero +spend (#97). ``PostgresChargebackReader`` hands out a repository bound to a fresh read-only connection per read. Explicit async SQLAlchemy Core, no ORM; satisfies the ports structurally without importing ``application``. """ @@ -35,6 +37,7 @@ user_principal, ) from tollgate.domain.chargeback import BudgetState, GroupBy, GroupByKind, SpendGroup +from tollgate.domain.errors import BudgetNotFound from tollgate.domain.ids import BudgetId from tollgate.domain.invariants import Balance from tollgate.domain.scopes import ScopeKind, scope_rank @@ -173,6 +176,17 @@ async def resolve_scope_ancestry( ).first() return None if row is None else {ScopeKind.ORG: row.org_id, ScopeKind.PROJECT: scope_id} + async def _node_has_own_budget(self, scope_kind: ScopeKind, scope_id: str) -> bool: + """Whether ``(scope_kind, scope_id)`` carries a budget of its own (a cheap keyed lookup).""" + row = ( + await self._conn.execute( + select(budget.c.budget_id) + .where(budget.c.scope_kind == scope_kind.value, budget.c.scope_id == scope_id) + .limit(1) + ) + ).first() + return row is not None + async def spend_rollup( self, scope_kind: ScopeKind, @@ -180,6 +194,13 @@ async def spend_rollup( period_start: datetime, group_by: GroupBy, ) -> Sequence[SpendGroup]: + # A single-node rollup sums only this node's own budget ledger, which already aggregates + # everything beneath it (ADR 0033) — but only if the node carries a budget. A budget-less + # node accrues no ledger rows of its own, so the aggregate would come back empty and + # silently report zero spend even while /v1/budgets shows the subtree's real spend. Fail + # explicitly so the two chargeback reads never disagree (#97). + if not await self._node_has_own_budget(scope_kind, scope_id): + raise BudgetNotFound(f"no budget at {scope_kind.value}:{scope_id} to roll up") led, bud, res = ledger.c, budget.c, reservation.c spend = func.sum(led.delta_committed_micro + led.delta_overage_micro) node = and_( diff --git a/src/tollgate/api/routes/chargeback.py b/src/tollgate/api/routes/chargeback.py index 1d6f0d9..cd7127b 100644 --- a/src/tollgate/api/routes/chargeback.py +++ b/src/tollgate/api/routes/chargeback.py @@ -18,6 +18,8 @@ from tollgate.api.dependencies import RequestAuth from tollgate.api.schemas import ( + MAX_LABEL_KEY_LEN, + MAX_STR_LEN, BudgetAlertState, BudgetStateResponse, BudgetStatesResponse, @@ -29,6 +31,7 @@ from tollgate.domain.chargeback import ( BudgetState, BudgetStatesView, + GroupBy, SpendRollup, crossed_thresholds, parse_group_by, @@ -52,11 +55,15 @@ def _parse_scope(scope: str | None) -> ScopeRef | None: - """Parse ``:`` into a :class:`ScopeRef`; ``None`` passes through. 422 on malformed.""" + """Parse ``:`` into a :class:`ScopeRef`; ``None`` passes through. 422 on malformed. + + ``scope_id`` is bounded to ``MAX_STR_LEN`` to match the command-path identifier limits — the + read path took an otherwise unbounded query value straight into a ``WHERE`` (#107). + """ if scope is None: return None kind, separator, scope_id = scope.partition(":") - if not separator or kind not in _SCOPE_KINDS or not scope_id: + if not separator or kind not in _SCOPE_KINDS or not scope_id or len(scope_id) > MAX_STR_LEN: raise HTTPException( status_code=422, detail="scope must be ':' where kind is one of org, team, user, project", @@ -64,6 +71,26 @@ def _parse_scope(scope: str | None) -> ScopeRef | None: return ScopeRef(scope_kind=ScopeKind(kind), scope_id=scope_id) +def _parse_group_by_query(raw: str) -> GroupBy: + """Parse the ``group_by`` query into a :class:`GroupBy`; 422 on malformed or an oversized key. + + The ``label:`` form otherwise carried an unbounded key straight into the rollup grouping; + bound it to ``MAX_LABEL_KEY_LEN`` like the command-path label keys (#107). + """ + parsed = parse_group_by(raw) + if parsed is None: + raise HTTPException( + status_code=422, + detail="group_by must be 'provider', 'model', or 'label:'", + ) + if parsed.label_key is not None and len(parsed.label_key) > MAX_LABEL_KEY_LEN: + raise HTTPException( + status_code=422, + detail=f"group_by label key must be at most {MAX_LABEL_KEY_LEN} characters", + ) + return parsed + + def _node_response(state: BudgetState) -> BudgetStateResponse: crossed = set(crossed_thresholds(state)) return BudgetStateResponse( @@ -117,12 +144,7 @@ async def spend( period_start: datetime | None = None, ) -> SpendRollupResponse: """Realized spend for a scope node, grouped by a dimension, for one period (section 2, 5.0).""" - parsed = parse_group_by(group_by) - if parsed is None: - raise HTTPException( - status_code=422, - detail="group_by must be 'provider', 'model', or 'label:'", - ) + parsed = _parse_group_by_query(group_by) handler: ChargebackHandler = request.app.state.chargeback_handler rollup = await handler.spend_rollup( auth, group_by=parsed, scope=_parse_scope(scope), period_start=period_start diff --git a/src/tollgate/api/schemas.py b/src/tollgate/api/schemas.py index 6f24d67..e554b52 100644 --- a/src/tollgate/api/schemas.py +++ b/src/tollgate/api/schemas.py @@ -30,7 +30,7 @@ MAX_LABEL_KEY_LEN: Final = 128 MAX_LABEL_VALUE_LEN: Final = 256 -_LabelKey = Annotated[str, StringConstraints(max_length=MAX_LABEL_KEY_LEN)] +_LabelKey = Annotated[str, StringConstraints(min_length=1, max_length=MAX_LABEL_KEY_LEN)] _LabelValue = Annotated[str, StringConstraints(max_length=MAX_LABEL_VALUE_LEN)] diff --git a/tests/integration/test_chargeback_spend_repo.py b/tests/integration/test_chargeback_spend_repo.py index 1099b60..3aa98ea 100644 --- a/tests/integration/test_chargeback_spend_repo.py +++ b/tests/integration/test_chargeback_spend_repo.py @@ -4,6 +4,7 @@ from datetime import UTC, datetime +import pytest from sqlalchemy.ext.asyncio import AsyncConnection from tollgate.adapters.postgres.chargeback_repo import PostgresChargebackRepository @@ -17,6 +18,7 @@ user_principal, ) from tollgate.domain.chargeback import GroupBy, GroupByKind +from tollgate.domain.errors import BudgetNotFound from tollgate.domain.scopes import ScopeKind _PERIOD = datetime(2026, 7, 1, tzinfo=UTC) @@ -352,8 +354,12 @@ async def test_other_period_and_other_node_excluded(db_conn: AsyncConnection) -> assert [(g.group, g.spend_micro) for g in rows] == [("anthropic", 200)] -async def test_node_without_a_budget_is_empty(db_conn: AsyncConnection) -> None: +async def test_node_without_a_budget_raises_budget_not_found(db_conn: AsyncConnection) -> None: + # A single-node rollup is only defined for a node that carries its own budget: its ledger + # already aggregates everything beneath it (ADR 0033). A budget-less node accrues no ledger + # rows of its own, so returning [] would silently report zero spend while /v1/budgets shows the + # subtree's real spend. Fail explicitly instead so the two reads never disagree (#97). await _tree(db_conn) # team t1 has no budget repo = PostgresChargebackRepository(db_conn) - rows = await repo.spend_rollup(ScopeKind.TEAM, "t1", _PERIOD, GroupBy(GroupByKind.PROVIDER)) - assert rows == [] + with pytest.raises(BudgetNotFound): + await repo.spend_rollup(ScopeKind.TEAM, "t1", _PERIOD, GroupBy(GroupByKind.PROVIDER)) diff --git a/tests/unit/test_api_schemas.py b/tests/unit/test_api_schemas.py index 5fcdbf6..847fc4e 100644 --- a/tests/unit/test_api_schemas.py +++ b/tests/unit/test_api_schemas.py @@ -144,6 +144,19 @@ def test_reserve_request_rejects_oversized_label_value() -> None: ) +def test_reserve_request_rejects_empty_label_key() -> None: + # An empty-string label key is a degenerate grouping dimension; reject it on the command path + # to match the other bounds (#107). + with pytest.raises(ValidationError): + ReserveRequest( + provider="anthropic", + model="claude", + input_bound_tokens=1, + max_output_tokens=1, + labels={"": "prod"}, + ) + + def test_reserve_request_rejects_empty_project_id() -> None: with pytest.raises(ValidationError): ReserveRequest( diff --git a/tests/unit/test_chargeback_route.py b/tests/unit/test_chargeback_route.py index 65ccb7a..f24bdb6 100644 --- a/tests/unit/test_chargeback_route.py +++ b/tests/unit/test_chargeback_route.py @@ -5,10 +5,16 @@ import pytest from fastapi import HTTPException -from tollgate.api.routes.chargeback import _node_response, _parse_scope, _spend_response +from tollgate.api.routes.chargeback import ( + _node_response, + _parse_group_by_query, + _parse_scope, + _spend_response, +) +from tollgate.api.schemas import MAX_LABEL_KEY_LEN, MAX_STR_LEN from tollgate.app import build_app from tollgate.config.settings import Settings -from tollgate.domain.chargeback import BudgetState, SpendGroup, SpendRollup +from tollgate.domain.chargeback import BudgetState, GroupBy, GroupByKind, SpendGroup, SpendRollup from tollgate.domain.ids import BudgetId from tollgate.domain.invariants import Balance from tollgate.domain.scopes import ScopeKind, ScopeRef @@ -30,6 +36,31 @@ def test_parse_scope_rejects_malformed_values_as_422(bad: str) -> None: assert excinfo.value.status_code == 422 +def test_parse_scope_rejects_an_oversized_scope_id_as_422() -> None: + # scope_id is otherwise unbounded on the read path; cap it like the command-path bounds (#107). + with pytest.raises(HTTPException) as excinfo: + _parse_scope("team:" + "x" * (MAX_STR_LEN + 1)) + assert excinfo.value.status_code == 422 + + +def test_parse_group_by_query_reads_the_dimensions() -> None: + assert _parse_group_by_query("provider") == GroupBy(GroupByKind.PROVIDER) + assert _parse_group_by_query("label:team") == GroupBy(GroupByKind.LABEL, "team") + + +@pytest.mark.parametrize("bad", ["", "unknown", "label:"]) +def test_parse_group_by_query_rejects_malformed_values_as_422(bad: str) -> None: + with pytest.raises(HTTPException) as excinfo: + _parse_group_by_query(bad) + assert excinfo.value.status_code == 422 + + +def test_parse_group_by_query_rejects_an_oversized_label_key_as_422() -> None: + with pytest.raises(HTTPException) as excinfo: + _parse_group_by_query("label:" + "k" * (MAX_LABEL_KEY_LEN + 1)) + assert excinfo.value.status_code == 422 + + def test_node_response_maps_amounts_utilization_and_alert_flags() -> None: state = BudgetState( budget_id=BudgetId("b-user"),