diff --git a/app/budget_pricing.py b/app/budget_pricing.py new file mode 100644 index 00000000..d026fe0f --- /dev/null +++ b/app/budget_pricing.py @@ -0,0 +1,213 @@ +"""Pure pricing helpers for budgeted requests. + +A budgeted key is charged for what the upstream delivered. Most of the time the +provider reports token usage and the charge is that measurement. When it does +not, the charge is estimated from the characters that actually reached the +client. Everything here is a pure function over plain dicts and strings, so it +has no database, no FastAPI, and no route state. +""" + +from __future__ import annotations + +import json + +# Crude character-to-token divisor, used only to price a delivery the provider +# did not measure. +CHARS_PER_TOKEN = 4 + +_PROMPT_USAGE_KEYS = ("prompt_tokens", "input_tokens") +_COMPLETION_USAGE_KEYS = ("completion_tokens", "output_tokens") +_COUNTABLE_USAGE_KEYS = _PROMPT_USAGE_KEYS + _COMPLETION_USAGE_KEYS + +# Stands in for a delivered part with no character count (an image, an audio +# clip). One character: enough to be a delivery, floored to one token. +_OPAQUE_PART_MARKER = "?" + + +def _part_texts(part) -> list[str]: + """The character-bearing strings one content part carries. + + A part is a bare string, a `{"text": ...}` dict, an Anthropic + `{"type": "tool_use", "name": ..., "input": ...}` block, or some other + non-text part (an image, an audio clip, a b64 payload). + + A tool_use block is billed as its name plus its input serialized the same + way the translators serialize it (`json.dumps(input or {})`), so an object + input is counted in full, not dropped. A non-text part has no honest + character count, but it was delivered, so it counts as one character: that + keeps it from settling at zero while `chars_to_tokens` floors it to a token. + """ + if isinstance(part, str): + return [part] + if not isinstance(part, dict): + return [] + if part.get("type") == "tool_use": + texts = [part["name"]] if isinstance(part.get("name"), str) else [] + raw_input = part.get("input") + if isinstance(raw_input, str): + texts.append(raw_input) + else: + texts.append(json.dumps(raw_input or {}, separators=(", ", ": "))) + return texts + if isinstance(part.get("text"), str): + return [part["text"]] + if part.get("type") == "tool_result": + # A tool result's billable text is its `content`: a string, or a list of + # parts in the same shape a message's content takes. Count that text, + # not a single marker character for the whole tool output. + return _content_texts(part.get("content")) + if part.get("type") in (None, "text"): + return [] + # A non-text part (image, audio, ...) has no honest character count, but it + # was sent or delivered and is billed. Count it as one character, so it + # floors to a token and never prices at zero. + return [_OPAQUE_PART_MARKER] + + +def _content_texts(content) -> list[str]: + """Text strings of a content value: a bare string, or a list of parts.""" + if isinstance(content, str): + return [content] + if isinstance(content, list): + texts: list[str] = [] + for part in content: + texts.extend(_part_texts(part)) + return texts + return [] + + +def text_chars(content) -> int: + """Character count of message content, for a str or a list of content parts. + + Counts every character of the text it is given, whitespace included. The + prompt is billed for all of it. This is deliberately not `blocking_delivery_chars`' + rule: that function drops whitespace-only values from a completion, which is + right for deciding whether something was delivered, but the streaming path + also prices each delta through here, and a whitespace-only delta is real + output that must not be dropped. + """ + if isinstance(content, str): + return len(content) + if isinstance(content, list): + return sum(len(text) for part in content for text in _part_texts(part)) + return 0 + + +def message_chars(content, tool_calls=None) -> int: + """Character count of one prompt message: its content plus its tool calls. + + An assistant turn that only called tools has `content` of None and its + billable text in `tool_calls`; pricing content alone would charge it zero. + """ + return text_chars(content) + len(tool_call_text(tool_calls)) + + +def _positive_number(value) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) and value > 0 + + +def countable_usage(usage, *, delivered: bool = False) -> bool: + """Whether a usage dict carries a measured token total. + + A truthy usage dict is not proof of measurement. Some upstreams report + `{"total_tokens": 123}` and nothing else, and LiteLLM emits a frame of + `{"prompt_tokens": 0, "completion_tokens": 0}` when the upstream reported + nothing at all. Settling on either records a zero cost for a delivered + completion, so only a countable key holding a positive number counts as a + measurement. + + When `delivered` is True the frame must also carry a positive completion-side + count. A prompt-only frame such as `{"prompt_tokens": 100, + "completion_tokens": 0}` would otherwise settle a delivered completion at + zero. A delivered completion that really cost tokens always reports them. + + `bool` is rejected explicitly: `True` is an `int` in Python, and a flag that + leaked into a token field is not a measurement. + """ + if not isinstance(usage, dict) or not usage: + return False + if delivered: + return any(_positive_number(usage.get(key)) for key in _COMPLETION_USAGE_KEYS) + return any(_positive_number(usage.get(key)) for key in _COUNTABLE_USAGE_KEYS) + + +def chars_to_tokens(chars: int) -> int: + """Characters to tokens. Zero characters is zero tokens. + + A non-zero count floors at one, so a prompt that was really sent is never + recorded as empty just because it is shorter than one token's worth of text. + """ + if chars <= 0: + return 0 + return max(1, chars // CHARS_PER_TOKEN) + + +def estimate_usage(prompt_chars: int, completion_chars: int) -> dict: + """Token counts priced from characters, for a delivery never measured.""" + return { + "prompt_tokens": chars_to_tokens(prompt_chars), + "completion_tokens": chars_to_tokens(completion_chars), + } + + +def tool_call_text(tool_calls) -> str: + """Text a message's tool calls carry: function names and string arguments. + + A tool call is delivered content the upstream bills for. Arguments sent as + an object rather than a string have no honest character count, so they + contribute nothing. + """ + if not isinstance(tool_calls, list): + return "" + parts: list[str] = [] + for call in tool_calls: + if not isinstance(call, dict): + continue + fn = call.get("function") + if not isinstance(fn, dict): + continue + for key in ("name", "arguments"): + value = fn.get(key) + if isinstance(value, str): + parts.append(value) + return "".join(parts) + + +def blocking_delivery_chars(response: dict) -> tuple[bool, int]: + """Whether a blocking completion delivered content, and how many characters. + + Whitespace-only values are not a delivery, and they add nothing to the + count either: the flag and the count answer the same question and must agree. + """ + choices = response.get("choices") + if not isinstance(choices, list): + return False, 0 + has_content = False + chars = 0 + for choice in choices: + if not isinstance(choice, dict): + continue + candidates: list[str] = [] + text = choice.get("text") + if isinstance(text, str): + candidates.append(text) + message = choice.get("message") + if isinstance(message, dict): + content = message.get("content") + if isinstance(content, str): + candidates.append(content) + elif isinstance(content, list): + for part in content: + candidates.extend(_part_texts(part)) + # A refusal is the completion itself when content is None: it + # reached the client, so it is billed output. + refusal = message.get("refusal") + if isinstance(refusal, str): + candidates.append(refusal) + candidates.append(tool_call_text(message.get("tool_calls"))) + for value in candidates: + if not value.strip(): + continue + chars += len(value) + has_content = True + return has_content, chars diff --git a/app/main.py b/app/main.py index 8fd75800..12f8aa03 100644 --- a/app/main.py +++ b/app/main.py @@ -66,8 +66,23 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: engine = get_engine(settings.database_url) + # budget_parks and budget_folded are created only by + # ensure_budget_columns, through the race-tolerant _apply_ddl path. Every + # worker runs this at boot, so an unguarded CREATE here would let two + # workers race on the same table and crash the loser's lifespan on a + # duplicate pg_class/pg_type entry. + from packages.db.models.budget_folded import BudgetFolded + from packages.db.models.budget_park import BudgetPark + + deferred = {BudgetPark.__table__, BudgetFolded.__table__} + boot_tables = [ + table for table in Base.metadata.sorted_tables + if table not in deferred + ] async with engine.begin() as conn: - await conn.run_sync(Base.metadata.create_all) + await conn.run_sync( + lambda sync: Base.metadata.create_all(sync, tables=boot_tables) + ) # create_all only makes missing tables, never alters existing ones. Bring # existing deployments (SQLite volume, Postgres) up to date with columns added diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 27649968..9a999463 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -1,7 +1,7 @@ """Per-key lifetime spend tracking schema and accounting primitives for ``ApiKey.budget_limit_cents``. -This module provides the accounting schema foundation and atomic charge primitives -(part 1 of the 4-part budget subsystem; request-path enforcement is wired in #161). +This module holds the atomic charge primitives, the pre-check that reads them, +and the park ledger for settlements that could not be recorded. The cap is a hard lifetime limit on the key's total spend, in microcents (1 cent = 10_000 microcents; 1 USD = 1_000_000 microcents, matching chat.py's @@ -31,16 +31,538 @@ from __future__ import annotations -from sqlalchemy import select, update +import asyncio +import uuid + +from sqlalchemy import delete, select, update +from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from packages.db.models.api_key import ApiKey +from packages.db.models.budget_folded import BudgetFolded +from packages.db.models.budget_park import BudgetPark +from packages.db.models.request_log import RequestLog # Defined in `packages.db.units` so the boot repair that clamps a counter to the # same cap scales by the identical number; re-exported here because this module # is the documented home of the budget accounting primitives. from packages.db.units import MICROCENTS_PER_CENT as MICROCENTS_PER_CENT +# A settlement that gives up after every retry leaves a delivered cost with no +# record anywhere: the log row and the charge are one transaction, so both roll +# back and the counter never moves. The obligation is parked — one row per +# settlement, keyed by its `trace_id` — and keeps counting against the cap until +# a budget pre-check folds it into `spent_microcents`. +# +# The park is a database table, not process memory, because the deployment +# stops its machine whenever it goes idle: an in-memory obligation is lost on +# the next cold start, which reopens the cap for exactly the key the failure +# was about to protect. `_unsettled` below is the hold for when even that write +# cannot be made. Preserve each settlement's trace_id so an already-charged +# settlement cannot cause another settlement to be discarded. +# +# One entry is retained per settlement trace: an already-charged settlement +# must not clear the hold of another trace, so each trace_id gets its own +# obligation rather than being merged into a single per-key total. That leaves +# a trace per settled trace retained in process memory while the database stays +# down — bounded by the in-flight settlement rate, not capped. A cap that +# merges traces reintroduces the race this structure exists to close. +_unsettled: dict[str, dict[str, int]] = {} + + +def _hold(api_key_id: str, trace_id: str, microcents: int) -> None: + """Keep each unsettled obligation separately, keyed by its trace_id.""" + holds = _unsettled.setdefault(api_key_id, {}) + # A trace identifies one settlement. Repeated attempts for that trace + # must not add the same obligation twice. + holds[trace_id] = max(holds.get(trace_id, 0), microcents) + + +class _FoldConflict(Exception): + """A concurrent worker folded the same park first; the loser retries later.""" + + +async def _park_is_durable( + *, trace_id: str, api_key_id: str, expected_microcents: int +) -> bool | None: + """Whether the durable row covers the entire expected obligation. + + ``None`` means the probe could not read the database at all — a different + answer from ``False`` (the row is not there). Callers deciding whether to + hold the obligation in memory must never treat an unreadable probe as + absence: a hold beside its own durable row lets another worker fold the + row and this worker re-file the same obligation later, billing it twice. + They retry until this probe answers (`_decide_park_durability`). + """ + from packages.db import session as session_mod + + factory = session_mod._session_factory + if factory is None: + # No database is configured, so no row can be durable here; the write + # path does not even attempt one, which makes "absent" a real answer. + return False + try: + async with factory() as s: + amount = await s.scalar( + select(BudgetPark.microcents).where( + BudgetPark.trace_id == trace_id, + BudgetPark.api_key_id == api_key_id, + ) + ) + return ( + amount is not None + and int(amount) >= expected_microcents + ) + except Exception: + return None + + +# How long to wait between durability-probe retries while the database is +# unreadable. The probe is cheap and the caller is already unwinding a failure, +# so the delay only keeps an outage from spinning the retry loop. +_PARK_PROBE_RETRY_SECONDS = 0.1 + + +async def _decide_park_durability( + *, trace_id: str, api_key_id: str, expected_microcents: int +) -> bool: + """Answer whether the row is durable, retrying a probe that cannot read. + + The answer decides whether the obligation is held in memory, and it cannot + be guessed: holding a durable obligation double-charges it (the row and the + hold), while dropping an absent one loses the spend. So an unreadable probe + (``None``) is retried until the database can answer. + """ + while True: + durable = await _park_is_durable( + trace_id=trace_id, + api_key_id=api_key_id, + expected_microcents=expected_microcents, + ) + if durable is not None: + return durable + await asyncio.sleep(_PARK_PROBE_RETRY_SECONDS) + + +async def _insert_park(*, trace_id: str, api_key_id: str, microcents: int) -> bool: + """Persist one parked obligation. Returns True when it is durable. + + Idempotent on `trace_id`: a commit that applied but whose ack was lost + retries into the same primary key instead of recording the obligation a + second time. On general failure, `_decide_park_durability` probes whether + the write landed anyway, retrying until the database can answer. Returns + False only when the obligation is confirmed not durable, leaving the caller + to hold the amount in memory. A cancellation propagates with the memory + copy still held. + """ + from packages.db import session as session_mod + + factory = session_mod._session_factory + if factory is None: + return False + try: + async with factory() as s: + try: + s.add( + BudgetPark( + trace_id=trace_id, api_key_id=api_key_id, microcents=microcents, + ) + ) + await s.commit() + except IntegrityError: + # A row for this `trace_id` is already there — our own retried + # write after an ack loss, or a repeated give-up for the same + # settlement. The row is the authoritative remaining debt for + # the trace: the only other writer that lowers it is the fold's + # trim, and the difference a trim removes moved into + # `spent_microcents` in that transaction. Raising the row back + # to `microcents` would add that already-billed difference a + # second time; accept the row exactly as it is. + await s.rollback() + return True + except asyncio.CancelledError: + raise + except Exception: + return await _decide_park_durability( + trace_id=trace_id, + api_key_id=api_key_id, + expected_microcents=microcents, + ) + + +async def record_unsettled_spend( + *, trace_id: str, api_key_id: str, microcents: int +) -> None: + """Keep a settlement that gave up counting against the key's cap. + + The caller is unwinding a failure, so this never loses the amount: a park + that cannot be written durably is held in memory under its `trace_id` for + the next pre-check to re-file, and a cancellation holds it before + propagating rather than taking the obligation with it. + """ + if microcents <= 0 or not trace_id or not api_key_id: + return + key = str(api_key_id) + try: + if await _insert_park( + trace_id=trace_id, api_key_id=key, microcents=microcents + ): + return + except asyncio.CancelledError: + # The cancellation can land after the park commit already applied. Hold + # only if the row is not durable: a durable row with a memory hold on + # top lets another worker fold the row and this worker re-file the same + # obligation later, billing it twice. + # + # The probe must reach a decision even if the task is cancelled again + # while it runs. Shielding it alone is not enough: a second cancel + # interrupts the await, and the amount would be neither held nor + # confirmed. So the loop keeps waiting on the probe task until it + # finishes, and the cancellation is re-raised afterwards. + probe = asyncio.ensure_future( + _decide_park_durability( + trace_id=str(trace_id), + api_key_id=key, + expected_microcents=microcents, + ) + ) + while not probe.done(): + try: + await asyncio.shield(probe) + except asyncio.CancelledError: + # Another cancel while waiting: keep waiting for the decision. + continue + durable = not probe.cancelled() and probe.exception() is None and bool(probe.result()) + if not durable: + _hold(key, str(trace_id), microcents) + raise + _hold(key, str(trace_id), microcents) + + +async def pending_parked_spend(api_key_id: str) -> int | None: + """The outstanding park for a key: durable rows plus whatever is memory-only. + + ``None`` means the durable ledger could not be read, which is a different + answer from ``0``: the table is the only record of a park written by another + worker, or by this one before it stopped, so folding a failed read into the + total reopens the cap for exactly the key the park exists to hold shut. The + memory half still counts on the way to a real durable read failing, because + that half is what an outage is expected to lose and what recovers after it. + """ + from packages.db import session as session_mod + + key = str(api_key_id) + factory = session_mod._session_factory + held = _unsettled.get(key, {}) + if factory is None: + return sum(held.values()) + try: + async with factory() as s: + rows = ( + await s.execute( + select(BudgetPark.trace_id, BudgetPark.microcents).where( + BudgetPark.api_key_id == key + ) + ) + ).all() + # A tombstoned trace was already folded into the counter by some + # worker: neither its rows nor its holds are outstanding anymore. + # A stale re-filed row for one still sits in the table until a + # fold deletes it unbilled. + folded = set( + ( + await s.execute( + select(BudgetFolded.trace_id).where( + BudgetFolded.api_key_id == key + ) + ) + ).scalars().all() + ) + except Exception: + return None + # Subtract only the part of each hold already represented durably. + # A matching trace alone does not prove the entire amount was persisted. + durable_amounts = { + trace_id: int(microcents) for trace_id, microcents in rows + } + stored = sum( + microcents + for trace_id, microcents in durable_amounts.items() + if trace_id not in folded + ) + + held_amount = sum( + max(0, amount - durable_amounts.get(trace_id, 0)) + for trace_id, amount in held.items() + if trace_id not in folded + ) + return stored + held_amount + + +async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: + """Fold a key's parked obligations into its recorded spend. Returns what moved. + + The park exists because a charge could not be recorded; leaving it parked + forever would mean a key at its cap is rejected by an amount that never + settles and never clears, so every pre-check tries to move it. Park rows are + read oldest-`created_at` first and the remaining allowance is applied to them + in that order; a park larger than the allowance bills what fits and is + rewritten to its remainder, rather than staying parked whole. That keeps the + invariant the fold exists to hold: either the queue is empty, or the counter + sits exactly on the cap. Without it a key can be refused at a lifetime spend + below its limit with a row that nothing will ever shrink, which is the state + this function is supposed to drain. The remainder is still a real debt — the + over-claim is the fail-closed policy — so it stays visible and keeps + `is_exhausted` blocking; it is never written off, and it folds for free the + moment the cap is raised. + + Two parks carrying the same `created_at` fall through to the `trace_id` + tiebreak, and `trace_id` is a uuid4 — so for obligations stamped within the + same clock tick the order is arbitrary rather than oldest-first. The + accounting does not depend on which row wins: the parked total is conserved + either way, the counter still reaches `min(cap, spent + debt)`, and only + which row is trimmed differs. + + The charge and the row writes share one transaction with compare-and-swap + guards on each: two workers folding the same park cannot double-bill it, + because the loser's UPDATE or DELETE matches nothing and its next request + folds what the winner left. + + Every fully-billed trace also leaves a tombstone (`budget_folded`) in the + same transaction. A memory hold that outlives another worker's fold — the + hold is process-local, the ledger is shared — is dropped at re-file time + instead of being re-inserted, and a row re-filed in the race window is + deleted unbilled here. Either way the obligation is counted exactly once. + """ + from packages.db import session as session_mod + + key = str(api_key_id) + factory = session_mod._session_factory + if factory is None: + return 0 + # Re-file each obligation under its own trace. A single merged row would + # let the fold mistake unrelated spend for an already-charged settlement. + # Bounded to the holds present on entry: an amount that grows during the + # write stays for the next pre-check instead of spinning here. + holds = _unsettled.get(key) + if holds: + try: + async with factory() as _s: + folded_traces = set( + ( + await _s.execute( + select(BudgetFolded.trace_id).where( + BudgetFolded.api_key_id == key + ) + ) + ).scalars().all() + ) + except Exception: + # Tombstones unreadable: leave every hold for a later pre-check + # rather than risk re-filing an obligation another worker folded. + folded_traces = None + for trace_id, amount in ( + list(holds.items()) if folded_traces is not None else [] + ): + if trace_id in folded_traces: + # Another worker already folded this obligation into the + # counter: drop the stale hold instead of re-inserting it. + current = _unsettled.get(key) + if current is not None and current.get(trace_id) == amount: + current.pop(trace_id) + if not current: + _unsettled.pop(key, None) + continue + # A cancellation here propagates with the hold still in place; a later + # pre-check re-files it under the same trace_id. + if not await _insert_park( + trace_id=trace_id, api_key_id=key, microcents=amount + ): + break + current = _unsettled.get(key) + if current is not None and current.get(trace_id) == amount: + current.pop(trace_id) + if not current: + _unsettled.pop(key, None) + # If the amount grew during the write, leave it in place for the + # next pre-check. + move = 0 + settling: list[tuple[str, int]] = [] + billed_amounts: dict[str, int] = {} + trim: tuple[str, int, int] | None = None + try: + async with factory() as s: + async with s.begin(): + spent = ( + await s.execute( + select(ApiKey.spent_microcents).where(ApiKey.id == key) + ) + ).scalar_one_or_none() + if spent is None: + return 0 + spent = int(spent) + rows = ( + await s.execute( + select(BudgetPark.trace_id, BudgetPark.microcents) + .where(BudgetPark.api_key_id == key) + # Oldest debt first. `created_at` is stamped + # Python-side with sub-second resolution, but two + # workers computing the same fold still have to agree on + # which row is the partial one, so `trace_id` breaks any + # tie rather than letting the order depend on a race. + .order_by(BudgetPark.created_at, BudgetPark.trace_id) + ) + ).all() + # A log row and its budget charge are committed atomically. + # If a commit acknowledgement was lost and the durability + # probe also failed, the matching park is only a fallback + # record; the request-log row proves the charge already landed. + logged = set( + ( + await s.scalars( + select(RequestLog.trace_id).where( + RequestLog.trace_id.in_( + [trace_id for trace_id, _amount in rows] + ) + ) + ) + ).all() + ) if rows else set() + # Traces some worker already folded: a row for one is a stale + # re-file that landed after the tombstone. It was already + # billed once — delete it without billing, or the same + # obligation is charged twice. + folded = set( + ( + await s.scalars( + select(BudgetFolded.trace_id).where( + BudgetFolded.api_key_id == key + ) + ) + ).all() + ) + stale = [ + (trace_id, int(amount)) + for trace_id, amount in rows + if trace_id in folded + ] + for trace_id, _amount in stale: + cleared = await s.execute( + delete(BudgetPark).where( + BudgetPark.trace_id == trace_id, + BudgetPark.api_key_id == key, + ) + ) + if cleared.rowcount != 1: + raise _FoldConflict + rows = [row for row in rows if row[0] not in folded] + already_charged = [ + (trace_id, int(amount)) + for trace_id, amount in rows + if trace_id in logged + ] + if already_charged: + billed_amounts.update({ + trace_id: int(amount) + for trace_id, amount in rows + if trace_id in logged + }) + for trace_id, amount in already_charged: + cleared = await s.execute( + delete(BudgetPark).where( + BudgetPark.trace_id == trace_id, + BudgetPark.api_key_id == key, + BudgetPark.microcents == amount, + ) + ) + if cleared.rowcount != 1: + raise _FoldConflict + s.add(BudgetFolded( + trace_id=trace_id, + api_key_id=key, + microcents=amount, + )) + rows = [row for row in rows if row[0] not in logged] + room = cap_microcents - spent + for trace_id, microcents in rows: + microcents = int(microcents) + if microcents <= room: + room -= microcents + move += microcents + settling.append((trace_id, microcents)) + billed_amounts[trace_id] = microcents + continue + if room > 0: + trim = (trace_id, microcents, microcents - room) + move += room + # Record the full original row amount so reconciliation + # credits the coexisting memory hold in full. The row + # below is only trimmed to the remainder, so the hold + # counts the trimmed remainder: original hold minus what + # the fold actually moved. + billed_amounts[trace_id] = microcents + break + # Even with no spend to move, already-charged rows may have + # been deleted and their corresponding holds need reconciliation. + if move <= 0 and not already_charged: + return 0 + else: + charged = await s.execute( + update(ApiKey) + .where(ApiKey.id == key, ApiKey.spent_microcents == spent) + .values(spent_microcents=spent + move) + ) + if charged.rowcount != 1: + raise _FoldConflict + if settling: + for trace_id, amount in settling: + cleared = await s.execute( + delete(BudgetPark).where( + BudgetPark.trace_id == trace_id, + BudgetPark.api_key_id == key, + BudgetPark.microcents == amount, + ) + ) + if cleared.rowcount != 1: + raise _FoldConflict + s.add(BudgetFolded( + trace_id=trace_id, + api_key_id=key, + microcents=amount, + )) + if trim is not None: + trace_id, whole, remainder = trim + trimmed = await s.execute( + update(BudgetPark) + .where( + BudgetPark.trace_id == trace_id, + BudgetPark.api_key_id == key, + BudgetPark.microcents == whole, + ) + .values(microcents=remainder) + ) + if trimmed.rowcount != 1: + raise _FoldConflict + except _FoldConflict: + return 0 + # Reconcile only the amount actually covered by committed rows. A hold + # for another trace is never cleared just because one row was billed. + holds = _unsettled.get(key) + if holds: + for trace_id, covered in billed_amounts.items(): + held_amount = holds.get(trace_id) + if held_amount is None: + continue + holds.pop(trace_id) + remainder = held_amount - covered + if remainder > 0: + # The original park was deleted. Give the remaining debt a + # fresh identity so re-filing cannot collide with that row. + holds[str(uuid.uuid4())] = remainder + if not holds: + _unsettled.pop(key, None) + return move + async def read_spent(db: AsyncSession, api_key_id: str) -> int: """Return the key's currently-recorded lifetime spend in microcents.""" @@ -50,6 +572,52 @@ async def read_spent(db: AsyncSession, api_key_id: str) -> int: return int(spent or 0) +async def budget_precheck(db: AsyncSession, api_key_id: str, cap_microcents: int) -> int: + """The key's spend as of this pre-check, in microcents. + + One read behind both halves of the caller's decision: whether to reject the + request, and what allowance is left for a cost that is not known yet. It + folds a parked obligation into the counter before answering, so a write + outage is neither a window of free requests nor a park that can never clear, + and the result adds whatever is still pending rather than reading only the + counter: when the fold could not commit, the obligation still has to block + dispatch. A park ledger that cannot be read at all is answered as the cap + rather than as no debt. + """ + key = str(api_key_id) + pending = await pending_parked_spend(key) + if pending is None: + # Unknown debt is answered as full debt, the way the durability probe + # that wrote the park reads an unreadable database as "not settled". + # The alternative is a key whose cap is held shut only by parked rows + # dispatching freely while the ledger is down. + return cap_microcents + # The park ledger is read before the counter, and the counter on a fresh + # snapshot: a fold both moves `spent_microcents` and empties the park + # queue, so a counter read first can pair a pre-fold spend with a zero + # pending — understating the spend by exactly the folded amount with no + # evidence left to trigger the re-read. Parks-first keeps the evidence: + # a fold racing these reads leaves `pending > 0`, which takes the + # re-read below. + try: + await db.rollback() + except Exception: + pass + spent = await read_spent(db, key) + if pending: + try: + await settle_parked_spend(key, cap_microcents) + except Exception: + # The fold runs on sessions of its own, so there is nothing to + # roll back here — and the re-read below still counts the park. + pass + spent = await read_spent(db, key) + pending = await pending_parked_spend(key) + if pending is None: + return cap_microcents + return spent + pending + + async def is_exhausted(db: AsyncSession, api_key_id: str, cap_microcents: int) -> bool: """Fast pre-check: has the key already reached its lifetime cap? @@ -57,8 +625,7 @@ async def is_exhausted(db: AsyncSession, api_key_id: str, cap_microcents: int) - ``MICROCENTS_PER_CENT``, not the column itself — passing the raw cents value asks whether the key has spent a ten-thousandth of its budget. """ - spent = await read_spent(db, api_key_id) - return spent >= cap_microcents + return await budget_precheck(db, api_key_id, cap_microcents) >= cap_microcents async def charge_budget( @@ -84,9 +651,9 @@ async def charge_budget( Both UPDATEs run with ``synchronize_session=False``: the session default evaluates the SET in Python against any identity-map copy of the row and - marks it dirty, so a caller that loaded the key before this charge (the - documented #161 flow: validate_api_key loaded it, same session commits the - charge) would flush that stale computed value over the DB's atomic result — + marks it dirty, so a caller that loaded the key before this charge (same + session, loaded by auth, then charged) would flush that stale computed value + over the DB's atomic result — silently dropping a concurrent charge. The DB row stays the only source of truth; read the counter back with ``read_spent`` if you need it. """ diff --git a/packages/db/migrate.py b/packages/db/migrate.py index 34c5e041..7afe6449 100644 --- a/packages/db/migrate.py +++ b/packages/db/migrate.py @@ -14,26 +14,34 @@ from __future__ import annotations +from collections.abc import Callable +from typing import Any + from sqlalchemy import BigInteger, inspect, text from sqlalchemy.exc import DBAPIError +from packages.db.models.budget_folded import BudgetFolded +from packages.db.models.budget_park import BudgetPark from packages.db.units import MICROCENTS_PER_CENT def _already_applied(err: DBAPIError) -> bool: """Whether a DDL failure means someone else applied the change first.""" msg = str(err).lower() + if "already exists" in msg or "duplicate column" in msg: + return True + # Two concurrent CREATE TABLE are serialised by the catalog rather than by + # the wording of a complaint, so the loser gets a unique violation on + # pg_class/pg_type (`*_relname_nsp_index`, `*_typname_nsp_index`) instead of + # "already exists". Same meaning: the object is there now. return ( - "already exists" in msg - or "duplicate column" in msg - # Concurrent CREATE INDEX on Postgres can lose the race at the catalog - # insert rather than the IF NOT EXISTS probe, surfacing as a verror on - # pg_class's unique index instead of the usual "already exists". - or "pg_class_relname_nsp_index" in msg + "duplicate key value" in msg + and "nsp_index" in msg + and ("relname" in msg or "typname" in msg) ) -async def _apply_ddl(conn, statement: str) -> None: +async def _apply_ddl(conn, statement: str | Callable[..., Any]) -> None: """Run one startup DDL statement, tolerating a boot that raced us to it. Every worker runs this in its lifespan, so the first boot after an upgrade @@ -43,10 +51,16 @@ async def _apply_ddl(conn, statement: str) -> None: from here, not a reason to keep the worker from booting. The failure is caught inside a SAVEPOINT because on Postgres an error would otherwise abort the whole transaction and take the rest of the startup with it. + + `statement` is raw SQL, or a callable handed to `run_sync` for DDL that only + the dialect's own generator can emit (a `Table.create`). """ try: async with conn.begin_nested(): - await conn.execute(text(statement)) + if callable(statement): + await conn.run_sync(statement) + else: + await conn.execute(text(statement)) except DBAPIError as err: if not _already_applied(err): raise @@ -61,16 +75,45 @@ async def ensure_budget_columns(engine) -> None: Postgres (the column is scaled into microcents for every comparison against spend, and an int4 ceiling is about 214,748 dollars of lifetime budget) and creates the `ix_requests_log_api_key_spend` index that create_all only builds - on fresh databases. Each step costs nothing on a database that needs none of - it — there the repair is a single indexed UPDATE that matches no row. + on fresh databases, and creates `budget_parks` for deployments that predate + the durable-recovery release — a lost settlement needs somewhere every + worker, and every reboot, can see it. Each step costs nothing on a database + that needs none of it — there the repair is a single indexed UPDATE that + matches no row. """ async with engine.begin() as conn: + tables = set( + await conn.run_sync(lambda sync: inspect(sync).get_table_names()) + ) cols = { c["name"]: c["type"] for c in await conn.run_sync(lambda sync: inspect(sync).get_columns("api_keys")) } is_postgres = engine.dialect.name == "postgresql" + if BudgetPark.__tablename__ not in tables: + # The only creator of this table at boot: app/main.py's create_all + # skips it, so fresh databases and upgrades both come through here. + # `checkfirst` re-reads the catalog, and + # that read cannot see another boot's uncommitted CREATE — so two + # workers both get here and one loses anyway. It goes through + # `_apply_ddl` for the same reason the ALTER does: losing that race + # has to count as having won, and the savepoint is what keeps the + # error from aborting the transaction the rest of startup runs in. + await _apply_ddl( + conn, + lambda sync: BudgetPark.__table__.create(sync, checkfirst=True), + ) + + if BudgetFolded.__tablename__ not in tables: + # Same guarded path as budget_parks: the fold's tombstones must + # exist before any pre-check can rely on them, and two workers + # racing the first boot after the upgrade must both survive. + await _apply_ddl( + conn, + lambda sync: BudgetFolded.__table__.create(sync, checkfirst=True), + ) + # The model declares ix_requests_log_api_key_spend (api_key_id, # is_deleted); create_all only builds it on fresh databases, so an # upgraded deployment would drift. Built before the seed below, which is @@ -110,8 +153,8 @@ async def ensure_budget_columns(engine) -> None: # Not a one-shot `= 0` backfill but a monotonic, cap-clamped repair run # every boot: the counter rises to max(current, log total) and is never # written past budget_limit_cents scaled to microcents or below its - # current value. Until #161 wires the per-request charge, request logs - # are the only recorded spend, so a counter seeded at an earlier boot + # current value. Request logs are the only spend record the counter can + # be rebuilt from, so a counter seeded at an earlier boot # goes stale as traffic flows and only re-aggregating closes the gap. # Afterwards the clamp keeps this a no-op on correct state: a key whose # final charge was clamped deliberately under-reports its log total, and diff --git a/packages/db/models/__init__.py b/packages/db/models/__init__.py index ca728ffa..57900b39 100644 --- a/packages/db/models/__init__.py +++ b/packages/db/models/__init__.py @@ -2,6 +2,8 @@ from packages.db.models.api_key import ApiKey from packages.db.models.base import Base, SoftDeleteMixin, TimestampMixin, UUIDMixin +from packages.db.models.budget_folded import BudgetFolded +from packages.db.models.budget_park import BudgetPark from packages.db.models.provider_key import ProviderKey from packages.db.models.quality_score_override import QualityScoreOverride from packages.db.models.quality_score_snapshot import QualityScoreSnapshot @@ -15,6 +17,8 @@ "TimestampMixin", "UUIDMixin", "ApiKey", + "BudgetFolded", + "BudgetPark", "ProviderKey", "QualityScoreOverride", "QualityScoreSnapshot", diff --git a/packages/db/models/budget_folded.py b/packages/db/models/budget_folded.py new file mode 100644 index 00000000..e47fca3f --- /dev/null +++ b/packages/db/models/budget_folded.py @@ -0,0 +1,36 @@ +"""Traces whose parked obligation has been folded into the spend counter. + +A memory hold is process-local but the park ledger is shared: worker A can +hold trace T while worker B folds T's durable row and deletes it. Without a +durable record of that fold, A's next pre-check re-files the hold as a fresh +row and the same obligation is billed a second time. A tombstone row, written +in the same transaction that billed the park row, makes "already folded" +durable: re-filing a tombstoned trace drops the hold instead of re-inserting +it, and pending spend stops counting it. + +One row per fully-folded trace. A trimmed row keeps its remainder under the +same trace and is not tombstoned — the remainder is still owed. +""" + +from datetime import datetime, timezone + +from sqlalchemy import BigInteger, DateTime, String, func +from sqlalchemy.orm import Mapped, mapped_column + +from packages.db.models.base import Base, TimestampMixin, UUIDMixin + + +class BudgetFolded(Base, UUIDMixin, TimestampMixin): + __tablename__ = "budget_folded" + + trace_id: Mapped[str] = mapped_column(String(36), nullable=False, unique=True) + api_key_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True) + microcents: Mapped[int] = mapped_column(BigInteger, nullable=False) + # Same resolution concern as BudgetPark.created_at: tombstones are only + # ever read by trace_id, but keep the stamp precise anyway. + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(timezone.utc), + server_default=func.now(), + nullable=False, + ) diff --git a/packages/db/models/budget_park.py b/packages/db/models/budget_park.py new file mode 100644 index 00000000..0b9dafba --- /dev/null +++ b/packages/db/models/budget_park.py @@ -0,0 +1,45 @@ +"""Unsettled budget obligations — delivered spend the ledger never recorded. + +A settlement that gives up after every retry leaves a delivered cost with no +record anywhere: the log row and the charge are one transaction, so both roll +back and `spent_microcents` never moves. The obligation is parked here, one row +per settlement, so it outlives the process that lost it and is visible to every +worker behind the same database. + +Rows are keyed by the settlement's `trace_id`, which makes every park write +idempotent: a commit that applied but whose ack was lost retries into the same +primary key instead of recording the obligation twice, and a write that fails +outright is checked for having landed anyway before the process falls back to +holding it in memory. A fold that bills a row either deletes it or shrinks it to +what the cap could not absorb, in the same transaction that moves +`spent_microcents`, and it drops the writer's memory hold for any row it fully +billed. +""" + +from datetime import datetime, timezone + +from sqlalchemy import BigInteger, DateTime, String, func +from sqlalchemy.orm import Mapped, mapped_column + +from packages.db.models.base import Base, TimestampMixin, UUIDMixin + + +class BudgetPark(Base, UUIDMixin, TimestampMixin): + __tablename__ = "budget_parks" + + trace_id: Mapped[str] = mapped_column(String(36), nullable=False, unique=True) + api_key_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True) + microcents: Mapped[int] = mapped_column(BigInteger, nullable=False) + # Overrides TimestampMixin's column, whose `server_default=func.now()` is + # CURRENT_TIMESTAMP: one second wide on SQLite, where a recovered outage + # re-files a whole batch of obligations in a single pre-check and every row + # in it ties. The fold bills oldest debt first, and two workers have to + # agree on which row the partial one is, so the stamp needs enough + # resolution to settle that on its own instead of falling through to the + # `trace_id` tiebreak — which is a uuid4, and so picks at random. + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(timezone.utc), + server_default=func.now(), + nullable=False, + ) diff --git a/tests/unit/test_budget_migration.py b/tests/unit/test_budget_migration.py index e1ec8ffd..8700ab50 100644 --- a/tests/unit/test_budget_migration.py +++ b/tests/unit/test_budget_migration.py @@ -60,6 +60,10 @@ async def _legacy_deploy_engine(tmp_sqlite_url): async with engine.begin() as conn: await conn.execute(text("DROP INDEX IF EXISTS ix_requests_log_api_key_spend")) await conn.execute(text("ALTER TABLE api_keys DROP COLUMN spent_microcents")) + # The durable park table is newer still: the last release had no + # budget_parks, so the upgrade has to create it. + await conn.execute(text("DROP TABLE IF EXISTS budget_parks")) + await conn.execute(text("DROP TABLE IF EXISTS budget_folded")) await engine.dispose() return build_engine(tmp_sqlite_url) @@ -239,10 +243,9 @@ def __getattr__(self, name): async def test_repair_raises_a_counter_staled_by_traffic(tmp_sqlite_url): """Traffic logged after one boot's seed is counted by the next. - Until #161 wires the per-request charge, the counter only moves at boot: - requests keep landing in requests_log while a seeded key's spent_microcents - sits frozen, so a counter from an earlier boot is stale by everything that - flowed since. A seed gated on `= 0` can never see that traffic (the counter + The counter is rebuilt at boot from request logs: requests keep landing in + requests_log while a seeded key's spent_microcents sits frozen, so a counter + from an earlier boot is stale by everything that flowed since. A seed gated on `= 0` can never see that traffic (the counter is already nonzero), and the undercount is permanent; the every-boot monotonic repair re-aggregates and closes the window at each restart. """ @@ -268,7 +271,7 @@ async def test_repair_raises_a_counter_staled_by_traffic(tmp_sqlite_url): async def test_repair_clamps_to_the_budget_cap(tmp_sqlite_url): """Re-aggregating never writes spend past the lifetime cap. - charge_budget (once #161 wires it) deliberately clamps a breaching request's + charge_budget deliberately clamps a breaching request's counter to the cap while the log row records the true, larger cost — so `SUM(logs) > cap` is correct steady state, and an unclamped re-seed would both overshoot the cap and undo that clamp on every boot. w1's cap is 100 @@ -357,3 +360,44 @@ def _record(conn, cursor, statement, parameters, context, executemany): finally: await engine.dispose() + + +async def test_budget_parks_is_created_by_the_guarded_path_on_upgrade(tmp_sqlite_url): + """An upgrade must create budget_parks through ensure_budget_columns. + + The boot-time create_all skips the table, so the guarded branch is the only + creator. Creating it unguarded there would let racing workers crash the + loser on a duplicate catalog entry. + """ + from packages.db.models.base import Base + from packages.db.models.budget_folded import BudgetFolded + from packages.db.models.budget_park import BudgetPark + + engine = await _legacy_deploy_engine(tmp_sqlite_url) + try: + assert "budget_parks" not in await _table_names(engine) + assert "budget_folded" not in await _table_names(engine) + + # Mirror app.main: create_all with the ledger tables excluded. + deferred = {BudgetPark.__table__, BudgetFolded.__table__} + boot_tables = [t for t in Base.metadata.sorted_tables if t not in deferred] + async with engine.begin() as conn: + await conn.run_sync(lambda sync: Base.metadata.create_all(sync, tables=boot_tables)) + assert "budget_parks" not in await _table_names(engine) + assert "budget_folded" not in await _table_names(engine) + + await ensure_budget_columns(engine) + assert "budget_parks" in await _table_names(engine) + assert "budget_folded" in await _table_names(engine) + + # A worker that races the same boot must not fail on the existing tables. + await ensure_budget_columns(engine) + assert "budget_parks" in await _table_names(engine) + assert "budget_folded" in await _table_names(engine) + finally: + await engine.dispose() + + +async def _table_names(engine) -> set[str]: + async with engine.connect() as conn: + return set(await conn.run_sync(lambda sync: sa_inspect(sync).get_table_names())) diff --git a/tests/unit/test_budget_park.py b/tests/unit/test_budget_park.py new file mode 100644 index 00000000..6871515a --- /dev/null +++ b/tests/unit/test_budget_park.py @@ -0,0 +1,566 @@ +"""Unit tests for the budget park ledger in packages.auth.spend. + +A settlement that gives up leaves a delivered cost parked against its key. These +tests pin the properties the cap depends on: a park is recorded once per +trace_id, a parked amount keeps counting against the cap, a fold bills the +oldest debt first and never double-bills, and a row already present in the +request log is cleared instead of billed twice. +""" + +import asyncio +from datetime import datetime, timedelta, timezone + +import pytest +from sqlalchemy import delete, select, update +from sqlalchemy.exc import OperationalError +from sqlalchemy.ext.asyncio import async_sessionmaker + +from packages.auth import spend +from packages.auth.spend import ( + budget_precheck, + pending_parked_spend, + read_spent, + record_unsettled_spend, + settle_parked_spend, +) +from packages.db import session as session_mod +from packages.db.models.api_key import ApiKey +from packages.db.models.budget_park import BudgetPark +from packages.db.models.request_log import RequestLog + + +@pytest.fixture +async def factory(tmp_sqlite_url): + """Session factory on a fresh schema, installed as the app's global factory.""" + from packages.db.engine import build_engine + from packages.db.models.base import Base + + engine = build_engine(tmp_sqlite_url) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + maker = async_sessionmaker(engine, expire_on_commit=False) + previous = session_mod._session_factory + session_mod._session_factory = maker + spend._unsettled.clear() + try: + yield maker + finally: + session_mod._session_factory = previous + spend._unsettled.clear() + await engine.dispose() + + +async def _make_key(maker, *, key_id: str, spent: int = 0) -> str: + async with maker() as s: + s.add( + ApiKey( + id=key_id, + workspace_id="default", + name=key_id, + key_hash=f"h-{key_id}", + key_prefix=f"p-{key_id}", + spent_microcents=spent, + ) + ) + await s.commit() + return key_id + + +async def _spent(maker, key_id: str) -> int: + async with maker() as s: + return await read_spent(s, key_id) + + +async def _park_rows(maker, key_id: str) -> list[tuple[str, int]]: + async with maker() as s: + rows = ( + await s.execute( + select(BudgetPark.trace_id, BudgetPark.microcents) + .where(BudgetPark.api_key_id == key_id) + .order_by(BudgetPark.trace_id) + ) + ).all() + return [(trace, int(micro)) for trace, micro in rows] + + +async def test_record_unsettled_spend_persists_one_row_per_trace(factory): + await _make_key(factory, key_id="k1") + await record_unsettled_spend(trace_id="t1", api_key_id="k1", microcents=500) + # A retried write of the same settlement must not park it twice. + await record_unsettled_spend(trace_id="t1", api_key_id="k1", microcents=500) + assert await _park_rows(factory, "k1") == [("t1", 500)] + + +async def test_record_ignores_zero_and_missing_identifiers(factory): + await _make_key(factory, key_id="k1") + await record_unsettled_spend(trace_id="t0", api_key_id="k1", microcents=0) + await record_unsettled_spend(trace_id="", api_key_id="k1", microcents=5) + await record_unsettled_spend(trace_id="t2", api_key_id="", microcents=5) + assert await _park_rows(factory, "k1") == [] + assert spend._unsettled == {} + + +async def test_pending_parked_spend_sums_durable_rows(factory): + await _make_key(factory, key_id="k1") + await record_unsettled_spend(trace_id="a", api_key_id="k1", microcents=300) + await record_unsettled_spend(trace_id="b", api_key_id="k1", microcents=200) + assert await pending_parked_spend("k1") == 500 + + +async def test_memory_hold_counts_when_the_database_is_unavailable(monkeypatch): + """With no session factory the only record is the memory hold; it must count.""" + previous = session_mod._session_factory + session_mod._session_factory = None + spend._unsettled.clear() + try: + await record_unsettled_spend(trace_id="m1", api_key_id="k9", microcents=400) + assert await pending_parked_spend("k9") == 400 + finally: + session_mod._session_factory = previous + spend._unsettled.clear() + + +async def test_memory_hold_is_bounded_by_key_not_by_failed_settlements(monkeypatch): + """One entry per key, however many settlements fail during an outage. + + Entries only leave the hold when a fold commits, and no fold can commit + while the database is down, so a per-settlement entry would grow for as + long as the outage lasted. + """ + previous = session_mod._session_factory + session_mod._session_factory = None + spend._unsettled.clear() + try: + for i in range(1000): + await record_unsettled_spend( + trace_id=f"t{i}", api_key_id="k9", microcents=1 + ) + assert len(spend._unsettled) == 1 + # Bounded, but nothing is lost: the full amount still counts. + assert await pending_parked_spend("k9") == 1000 + # And it stays attributable to one key. + assert await pending_parked_spend("other-key") == 0 + finally: + session_mod._session_factory = previous + spend._unsettled.clear() + + +async def test_merging_into_a_hold_raises_a_colliding_park_row(factory): + """A hold that merged more settlements must not lose the difference. + + The row for the hold's trace can already exist when the fold re-files it — + the original write landed and its ack was lost. Re-filing then collides, and + treating that collision as success would drop everything merged in since. + """ + await _make_key(factory, key_id="k1", spent=0) + async with factory() as s: + s.add(BudgetPark(trace_id="t", api_key_id="k1", microcents=100)) + await s.commit() + + # The writer believed all three failed and merged them into one hold. + spend._hold("k1", "t", 100) + spend._hold("k1", "t2", 50) + spend._hold("k1", "t3", 25) + + moved = await settle_parked_spend("k1", cap_microcents=10_000) + + assert moved == 175 + assert await _spent(factory, "k1") == 175 + assert await _park_rows(factory, "k1") == [] + assert spend._unsettled == {} + + +async def test_fold_bills_the_oldest_park_first_and_trims_the_partial(factory): + await _make_key(factory, key_id="k1", spent=0) + base = datetime(2026, 1, 1, tzinfo=timezone.utc) + async with factory() as s: + s.add_all( + [ + BudgetPark(trace_id="new", api_key_id="k1", microcents=600, + created_at=base + timedelta(seconds=2)), + BudgetPark(trace_id="old", api_key_id="k1", microcents=500, + created_at=base), + ] + ) + await s.commit() + + # Cap of 800: "old" (500) fits whole, "new" (600) only has 300 of room. + moved = await settle_parked_spend("k1", cap_microcents=800) + + assert moved == 800 + assert await _spent(factory, "k1") == 800 + # "old" is fully billed and gone; "new" is trimmed to its 300 remainder. + assert await _park_rows(factory, "k1") == [("new", 300)] + + +async def test_fold_clears_a_park_whose_charge_is_already_logged(factory): + """A park whose request-log row landed was already billed; folding it again doubles the charge.""" + await _make_key(factory, key_id="k1", spent=100) + async with factory() as s: + s.add(BudgetPark(trace_id="done", api_key_id="k1", microcents=250)) + s.add( + RequestLog( + workspace_id="default", + api_key_id="k1", + trace_id="done", + cost_microcents=250, + model_requested="m", + model_resolved="m", + provider="openai", + routing_strategy="passthrough", + input_tokens=10, + output_tokens=20, + latency_ms=5, + status_code=200, + ) + ) + await s.commit() + + moved = await settle_parked_spend("k1", cap_microcents=10_000) + + assert moved == 0 + assert await _spent(factory, "k1") == 100 + assert await _park_rows(factory, "k1") == [] + + +async def test_fold_does_nothing_when_the_cap_has_no_room(factory): + await _make_key(factory, key_id="k1", spent=1_000) + await record_unsettled_spend(trace_id="t", api_key_id="k1", microcents=50) + assert await settle_parked_spend("k1", cap_microcents=1_000) == 0 + assert await _spent(factory, "k1") == 1_000 + assert await _park_rows(factory, "k1") == [("t", 50)] + + +async def test_precheck_counts_pending_park_against_the_cap(factory): + """A parked obligation blocks dispatch even when the fold cannot commit.""" + await _make_key(factory, key_id="k1", spent=900) + await record_unsettled_spend(trace_id="t", api_key_id="k1", microcents=200) + async with factory() as s: + assert await budget_precheck(s, "k1", cap_microcents=1_000) >= 1_000 + + +async def test_precheck_folds_pending_park_into_the_counter(factory): + await _make_key(factory, key_id="k1", spent=0) + await record_unsettled_spend(trace_id="t", api_key_id="k1", microcents=300) + async with factory() as s: + total = await budget_precheck(s, "k1", cap_microcents=10_000) + assert total == 300 + assert await _spent(factory, "k1") == 300 + assert await _park_rows(factory, "k1") == [] + + +async def test_precheck_treats_an_unreadable_ledger_as_full_debt(factory, monkeypatch): + await _make_key(factory, key_id="k1") + + async def _unreadable(_key): + return None + + monkeypatch.setattr(spend, "pending_parked_spend", _unreadable) + async with factory() as s: + assert await budget_precheck(s, "k1", cap_microcents=777) == 777 + + +async def test_cancel_after_a_landed_park_commit_leaves_no_hold(factory, monkeypatch): + """A commit that landed before the cancellation is durable: no memory hold. + + A hold alongside the durable row lets another worker fold the row and this + worker re-file the same obligation later, billing it twice. + """ + await _make_key(factory, key_id="k", spent=0) + real_insert = spend._insert_park + + async def insert_then_cancel(**kw): + assert await real_insert(**kw) + raise asyncio.CancelledError() + + monkeypatch.setattr(spend, "_insert_park", insert_then_cancel) + with pytest.raises(asyncio.CancelledError): + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=600) + monkeypatch.setattr(spend, "_insert_park", real_insert) + + assert spend._unsettled.get("k") in (None, {}) + await settle_parked_spend("k", cap_microcents=10_000) + await settle_parked_spend("k", cap_microcents=10_000) + assert await _spent(factory, "k") == 600 + + +async def test_cancel_before_any_commit_still_holds_the_amount(factory, monkeypatch): + """If the park never landed, a cancellation must still hold the obligation.""" + await _make_key(factory, key_id="k", spent=0) + + async def insert_cancelled(**kw): + raise asyncio.CancelledError() + + monkeypatch.setattr(spend, "_insert_park", insert_cancelled) + with pytest.raises(asyncio.CancelledError): + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=600) + assert spend._unsettled.get("k") == {"t": 600} + + +async def test_second_cancel_during_the_durability_probe_still_decides_the_hold( + factory, monkeypatch +): + """A cancel can arrive again while the cancel-path probe is in flight. + + The probe must still reach its decision, so a durable row is not held and an + absent row is held, rather than the cancellation abandoning the decision. + """ + await _make_key(factory, key_id="k", spent=0) + + async def insert_cancelled(**kw): + raise asyncio.CancelledError() + + slow_probe_started = asyncio.Event() + + real_probe = spend._park_is_durable + + async def slow_probe(**kw): + slow_probe_started.set() + await asyncio.sleep(0.05) + return await real_probe(**kw) + + monkeypatch.setattr(spend, "_insert_park", insert_cancelled) + monkeypatch.setattr(spend, "_park_is_durable", slow_probe) + + async def run(): + with pytest.raises(asyncio.CancelledError): + await spend.record_unsettled_spend(trace_id="t", api_key_id="k", microcents=600) + + task = asyncio.create_task(run()) + await slow_probe_started.wait() + task.cancel() # second cancel while the probe is in flight + await task + # No durable row exists, so the probe decided "not durable" and the amount is held. + assert spend._unsettled.get("k") == {"t": 600} + + +class _CommitFailingSession: + """A session whose `commit` fails, simulating a write outage. + + ``land=True`` commits for real and then raises — an acknowledgement loss: + the row is durable but the writer cannot know it. Otherwise the commit + raises before anything lands. + """ + + def __init__(self, session, *, land: bool): + self._session = session + self._land = land + + def __getattr__(self, name): + return getattr(self._session, name) + + async def __aenter__(self): + await self._session.__aenter__() + return self + + async def __aexit__(self, *exc): + return await self._session.__aexit__(*exc) + + async def commit(self): + if self._land: + await self._session.commit() + raise OperationalError("COMMIT", None, Exception("commit failed")) + + +def _commit_failing_factory(maker, *, land: bool): + """A session factory whose sessions fail their commit as described above.""" + + def factory(): + return _CommitFailingSession(maker(), land=land) + + return factory + + +async def test_ack_lost_write_is_not_held_when_the_probe_has_to_retry( + factory, monkeypatch +): + """A park that landed must never gain a hold on a failed probe. + + Regression for the P1 where the general-failure path probed once and read + an unreadable probe as "row absent": the obligation was held on top of its + own durable row, another worker folded the row, and this worker's hold + re-filed the same obligation and billed it a second time. + """ + await _make_key(factory, key_id="k", spent=0) + real_probe = spend._park_is_durable + probes = {"n": 0} + + async def unreadable_once(**kw): + probes["n"] += 1 + if probes["n"] == 1: + return None # the database is unreachable at probe time + return await real_probe(**kw) + + monkeypatch.setattr(spend, "_park_is_durable", unreadable_once) + monkeypatch.setattr(spend, "_PARK_PROBE_RETRY_SECONDS", 0) + previous = session_mod._session_factory + session_mod._session_factory = _commit_failing_factory(factory, land=True) + try: + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=600) + finally: + session_mod._session_factory = previous + + # The probe was retried instead of being read as "absent", so no hold. + assert probes["n"] == 2 + assert spend._unsettled.get("k") in (None, {}) + assert await _park_rows(factory, "k") == [("t", 600)] + + # Another worker folds the durable row while this worker holds nothing. + async with factory() as s: + await s.execute( + update(ApiKey).where(ApiKey.id == "k").values(spent_microcents=600) + ) + await s.execute(delete(BudgetPark).where(BudgetPark.trace_id == "t")) + await s.commit() + + # Nothing is left to re-file, so the obligation is billed exactly once. + assert await settle_parked_spend("k", cap_microcents=10_000) == 0 + assert await _spent(factory, "k") == 600 + assert await _park_rows(factory, "k") == [] + + +async def test_unlanded_write_with_a_retried_probe_still_holds_the_amount( + factory, monkeypatch +): + """A park that never landed is held — but only once the probe says absent.""" + await _make_key(factory, key_id="k", spent=0) + real_probe = spend._park_is_durable + probes = {"n": 0} + + async def unreadable_once(**kw): + probes["n"] += 1 + if probes["n"] == 1: + return None + return await real_probe(**kw) + + monkeypatch.setattr(spend, "_park_is_durable", unreadable_once) + monkeypatch.setattr(spend, "_PARK_PROBE_RETRY_SECONDS", 0) + previous = session_mod._session_factory + session_mod._session_factory = _commit_failing_factory(factory, land=False) + try: + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=600) + finally: + session_mod._session_factory = previous + + assert probes["n"] == 2 + assert await _park_rows(factory, "k") == [] + assert spend._unsettled.get("k") == {"t": 600} + + +async def test_second_give_up_after_a_partial_fold_does_not_raise_the_row(factory): + """A repeated park write must not resurrect debt the fold already billed. + + Regression for the P1 where the collision path raised the row back to the + incoming full amount: after the fold billed part of a parked obligation + and trimmed the row to the remainder, a repeated give-up for the same + settlement raised the row to the full amount again, double-counting the + share already moved into `spent_microcents`. + """ + await _make_key(factory, key_id="k", spent=0) + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=200) + + # Room for 140 of the 200: the fold bills 140 and trims the row to 60. + assert await settle_parked_spend("k", cap_microcents=140) == 140 + assert await _spent(factory, "k") == 140 + assert await _park_rows(factory, "k") == [("t", 60)] + + # The settlement is retried and gives up again for the same trace. + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=200) + + assert spend._unsettled.get("k") in (None, {}) + assert await _park_rows(factory, "k") == [("t", 60)] + + # The row is the remaining debt: 140 + 60 = the 200 obligation, once. + assert await settle_parked_spend("k", cap_microcents=10_000) == 60 + assert await _spent(factory, "k") == 200 + assert await _park_rows(factory, "k") == [] + + +async def test_folded_trace_is_never_rebilled_across_workers(factory): + """Worker A holds T while worker B folds T's row: no double charge. + + The commit landed but the ack was lost, so A's hold coexists with its + durable row. B's fold bills the row and leaves a tombstone; A's next + pre-check must drop the stale hold instead of re-filing it. + """ + from packages.db.models.budget_folded import BudgetFolded + + await _make_key(factory, key_id="k", spent=0) + async with factory() as s: + s.add(BudgetPark(trace_id="T", api_key_id="k", microcents=100)) + await s.commit() + # A's view: the hold survived beside its durable row. + spend._hold("k", "T", 100) + # B's fold (B never saw the hold): bills 100, deletes the row. + holds = spend._unsettled.pop("k") + assert await settle_parked_spend("k", cap_microcents=10_000) == 100 + spend._unsettled["k"] = holds + # A's pre-check: the tombstone drops the hold, nothing is re-filed. + assert await settle_parked_spend("k", cap_microcents=10_000) == 0 + assert await _spent(factory, "k") == 100 + assert await _park_rows(factory, "k") == [] + assert spend._unsettled.get("k") in (None, {}) + async with factory() as s: + tombstoned = ( + await s.execute( + select(BudgetFolded.trace_id).where(BudgetFolded.api_key_id == "k") + ) + ).scalars().all() + assert list(tombstoned) == ["T"] + + +async def test_tombstoned_hold_is_not_counted_in_pending(factory): + """A hold whose trace was already folded is spent, not outstanding.""" + from packages.db.models.budget_folded import BudgetFolded + + await _make_key(factory, key_id="k", spent=100) + async with factory() as s: + s.add(BudgetFolded(trace_id="T", api_key_id="k", microcents=100)) + await s.commit() + spend._hold("k", "T", 100) + assert await pending_parked_spend("k") == 0 + + +async def test_stale_refiled_row_is_deleted_unbilled(factory): + """A row re-filed after its trace was tombstoned is removed, not billed. + + The race window: A re-inserts T after B's fold tombstoned it but before + A's fold reads. The fold must delete the stale row without moving the + counter a second time. + """ + from packages.db.models.budget_folded import BudgetFolded + + await _make_key(factory, key_id="k", spent=100) + async with factory() as s: + s.add(BudgetFolded(trace_id="T", api_key_id="k", microcents=100)) + s.add(BudgetPark(trace_id="T", api_key_id="k", microcents=100)) + await s.commit() + assert await settle_parked_spend("k", cap_microcents=10_000) == 0 + assert await _spent(factory, "k") == 100 + assert await _park_rows(factory, "k") == [] + + +async def test_repeated_give_up_after_a_full_fold_is_not_rebilled(factory): + """The unique trace key cannot be freed by the fold: re-parking is not re-billing. + + Regression for the P1 where a full bill deleted the park row and left the + `trace_id` free: a repeated give-up for the same settlement re-inserted the + full amount, and the next fold — finding no request-log row for a charge + that never landed — billed the same obligation a second time. + """ + await _make_key(factory, key_id="k", spent=0) + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=100) + assert await settle_parked_spend("k", cap_microcents=10_000) == 100 + assert await _spent(factory, "k") == 100 + assert await _park_rows(factory, "k") == [] + + # The settlement is retried and gives up again for the same trace: the + # write lands, because the fold freed the unique key when it deleted the + # row. It must not read as outstanding debt. + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=100) + assert await pending_parked_spend("k") == 0 + + # The next fold deletes the stale re-file unbilled: one obligation, one bill. + assert await settle_parked_spend("k", cap_microcents=10_000) == 0 + assert await _spent(factory, "k") == 100 + assert await _park_rows(factory, "k") == [] diff --git a/tests/unit/test_budget_pricing.py b/tests/unit/test_budget_pricing.py new file mode 100644 index 00000000..fa8ed684 --- /dev/null +++ b/tests/unit/test_budget_pricing.py @@ -0,0 +1,279 @@ +"""Unit tests for app.budget_pricing: estimating what an unmeasured delivery cost.""" + +import pytest + +from app.budget_pricing import ( + blocking_delivery_chars, + chars_to_tokens, + countable_usage, + estimate_usage, + message_chars, + text_chars, + tool_call_text, +) + + +@pytest.mark.parametrize( + ("usage", "expected"), + [ + ({"prompt_tokens": 5}, True), + ({"completion_tokens": 1}, True), + ({"input_tokens": 5}, True), + ({"output_tokens": 5}, True), + # A zero on one side is still a measurement when the other side counts. + ({"prompt_tokens": 0, "completion_tokens": 7}, True), + ({"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, False), + ({"total_tokens": 123}, False), + ({"prompt_tokens": None}, False), + ({"prompt_tokens": "12"}, False), + ({"prompt_tokens": True}, False), + ({"prompt_tokens": -5}, False), + ({}, False), + (None, False), + ("not a dict", False), + ], +) +def test_countable_usage_requires_a_positive_measurement(usage, expected): + assert countable_usage(usage) is expected + + +def test_countable_usage_rejects_a_zero_placeholder_frame(): + """The frame LiteLLM emits when the upstream reported nothing. + + A delivered completion settled on this records a zero cost, which is the + bypass the guard exists to prevent. + """ + assert countable_usage({"prompt_tokens": 0, "completion_tokens": 0}) is False + assert ( + countable_usage( + {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + ) + is False + ) + + +@pytest.mark.parametrize( + ("chars", "expected"), + [(-5, 0), (0, 0), (1, 1), (3, 1), (4, 1), (8, 2), (9, 2)], +) +def test_chars_to_tokens_floors_nonzero_at_one(chars, expected): + assert chars_to_tokens(chars) == expected + + +def test_zero_characters_is_zero_tokens_not_one(): + assert estimate_usage(0, 0) == {"prompt_tokens": 0, "completion_tokens": 0} + + +def test_estimate_usage_prices_both_sides_from_characters(): + assert estimate_usage(400, 40) == {"prompt_tokens": 100, "completion_tokens": 10} + + +def test_text_chars_across_str_and_content_parts(): + assert text_chars("hello") == 5 + # A non-text part is billed, so it counts as one character, not zero. + assert text_chars([{"type": "text", "text": "ab"}, {"type": "image_url"}]) == 3 + assert text_chars(None) == 0 + + +def test_text_chars_counts_bare_string_parts(): + """Anthropic-style content may be a list of plain strings, not just dicts. + + Mirrors `blocking_delivery_chars`, which accepts a bare `str` part: a + prompt and a completion carrying the same bare-string content must price + the same characters rather than the prompt pricing at zero. + """ + assert text_chars(["hello", "world"]) == 10 + assert text_chars([{"text": "ab"}, "cd", {"no": "text"}, 42]) == 4 + + +def test_text_chars_keeps_whitespace_so_streamed_deltas_are_not_undercounted(): + """Whitespace-only deltas are real streamed output and must be priced. + + The delivery path skips whitespace-only values when deciding whether a + completion delivered anything; `text_chars` must not, or streamed spaces + between words would be dropped from the count. + """ + assert text_chars(" ") == 1 + assert text_chars(" ") == 3 + assert text_chars([" ", "x"]) == 3 + + +def test_tool_call_text_counts_names_and_string_arguments(): + calls = [ + {"function": {"name": "search", "arguments": '{"q":"x"}'}}, + {"function": {"name": "fetch", "arguments": {"url": "obj-not-counted"}}}, + ] + assert tool_call_text(calls) == 'search{"q":"x"}fetch' + + +def test_tool_call_text_ignores_malformed_entries(): + assert tool_call_text(None) == "" + assert tool_call_text(["oops", {"function": "nope"}]) == "" + + +def test_blocking_delivery_counts_content_and_tool_calls(): + response = { + "choices": [ + { + "message": { + "content": "hi", + "tool_calls": [{"function": {"name": "f", "arguments": "{}"}}], + } + } + ] + } + assert blocking_delivery_chars(response) == (True, len("hi") + len("f{}")) + + +def test_whitespace_only_is_not_a_delivery_and_adds_no_characters(): + response = {"choices": [{"message": {"content": " \n"}}]} + assert blocking_delivery_chars(response) == (False, 0) + + +def test_blocking_delivery_reads_legacy_text_and_part_lists(): + response = { + "choices": [ + {"text": "abc"}, + {"message": {"content": [{"text": "de"}, "f"]}}, + ] + } + assert blocking_delivery_chars(response) == (True, 6) + + +def test_blocking_delivery_empty_or_malformed_response(): + assert blocking_delivery_chars({}) == (False, 0) + assert blocking_delivery_chars({"choices": "nope"}) == (False, 0) + + +def test_message_chars_counts_tool_calls_on_an_assistant_turn(): + """An assistant turn with content None and only tool_calls must not price at zero.""" + calls = [{"function": {"name": "search", "arguments": '{"q":"x"}'}}] + assert message_chars(None, calls) == len('search{"q":"x"}') + assert message_chars("hi", calls) == len("hi") + len('search{"q":"x"}') + assert message_chars("hi") == 2 + + +def test_prompt_and_completion_price_the_same_tool_call_characters(): + """A tool-call turn in the prompt and the same turn delivered must agree.""" + calls = [{"function": {"name": "search", "arguments": '{"q":"x"}'}}] + prompt_side = message_chars(None, calls) + _, delivered_side = blocking_delivery_chars( + {"choices": [{"message": {"content": None, "tool_calls": calls}}]} + ) + assert prompt_side == delivered_side > 0 + + +def test_prompt_only_frame_is_not_countable_when_content_was_delivered(): + """A prompt-measured frame with a zero completion would settle a delivered completion at zero.""" + frame = {"prompt_tokens": 100, "completion_tokens": 0} + assert countable_usage(frame) is True + assert countable_usage(frame, delivered=True) is False + + +def test_delivered_frame_with_positive_completion_is_countable(): + assert countable_usage({"prompt_tokens": 100, "completion_tokens": 7}, delivered=True) is True + assert countable_usage({"output_tokens": 7}, delivered=True) is True + + +def test_anthropic_tool_use_part_is_counted_as_delivered_content(): + response = { + "choices": [ + { + "message": { + "content": [ + {"type": "tool_use", "name": "search", "input": '{"q":"x"}'}, + ] + } + } + ] + } + assert blocking_delivery_chars(response) == (True, len("search") + len('{"q":"x"}')) + + +def test_tool_use_part_with_object_input_counts_name_and_serialized_input(): + """Anthropic tool input is an object on the wire. The translators bill it as + json.dumps(input), so the prompt must count that same serialization.""" + part = {"type": "tool_use", "name": "get_weather", "input": {"city": "SF"}} + assert text_chars([part]) == len("get_weather") + len('{"city": "SF"}') + + +def test_tool_use_object_input_matches_the_translator_serialization(): + from app.protocols.anthropic import _translate_assistant_message + + blocks = [{"type": "tool_use", "id": "t1", "name": "get_weather", "input": {"city": "SF"}}] + translated = _translate_assistant_message(blocks) + billed = len(translated["tool_calls"][0]["function"]["name"]) + len( + translated["tool_calls"][0]["function"]["arguments"] + ) + assert text_chars([blocks[0]]) == billed + + +def test_tool_use_with_missing_input_counts_empty_object(): + part = {"type": "tool_use", "name": "ping"} + assert text_chars([part]) == len("ping") + len("{}") + + +def test_tool_use_only_completion_is_a_delivery_not_zero(): + response = {"choices": [{"message": {"content": [{"type": "tool_use", "name": "f"}]}}]} + # The tool name plus the empty object a missing input serializes to. + assert blocking_delivery_chars(response) == (True, len("f") + len("{}")) + + +def test_refusal_with_no_content_is_a_delivery(): + """A moderation refusal reaches the client as output and must not settle at zero.""" + response = {"choices": [{"message": {"content": None, "refusal": "I can't help with that."}}]} + assert blocking_delivery_chars(response) == (True, len("I can't help with that.")) + + +def test_non_text_delivered_part_counts_as_one_character_not_zero(): + """An image or audio part was delivered and billed; it has no character + count, so it counts as one character, which floors to one token.""" + response = {"choices": [{"message": {"content": [{"type": "image_url", "image_url": {"url": "x"}}]}}]} + delivered, chars = blocking_delivery_chars(response) + assert delivered is True + assert chars == 1 + assert chars_to_tokens(chars) == 1 + + +def test_text_and_tool_use_parts_are_not_treated_as_opaque(): + parts = [{"type": "text", "text": "hi"}, {"type": "tool_use", "name": "f", "input": {}}] + assert text_chars(parts) == len("hi") + len("f") + len("{}") + + +def test_image_only_prompt_is_not_priced_at_zero(): + """A vision prompt with no text still sent billed input. It must floor to a token.""" + content = [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}] + assert chars_to_tokens(text_chars(content)) == 1 + + +def test_tool_result_string_content_counts_its_real_length(): + """A tool output is billed text. It must not price as one marker character.""" + part = {"type": "tool_result", "tool_use_id": "t1", "content": "x" * 400} + assert text_chars([part]) == 400 + + +def test_tool_result_list_content_counts_its_text_parts(): + part = { + "type": "tool_result", + "tool_use_id": "t1", + "content": [{"type": "text", "text": "abc"}, {"type": "text", "text": "de"}], + } + assert text_chars([part]) == len("abc") + len("de") + + +def test_tool_result_with_an_image_inside_still_floors_the_image(): + part = { + "type": "tool_result", + "tool_use_id": "t1", + "content": [{"type": "image", "source": {"data": "AAAA"}}], + } + assert text_chars([part]) == 1 + + +def test_tool_result_without_content_counts_nothing(): + assert text_chars([{"type": "tool_result", "tool_use_id": "t1"}]) == 0 + + +def test_image_part_still_floors_to_one_character_on_the_prompt_side(): + """An image is billed; it has no honest character count, so it floors to one.""" + assert text_chars([{"type": "image_url", "image_url": {"url": "data:x"}}]) == 1 diff --git a/tests/unit/test_budget_spend.py b/tests/unit/test_budget_spend.py index ffb1a0af..23edd2b9 100644 --- a/tests/unit/test_budget_spend.py +++ b/tests/unit/test_budget_spend.py @@ -4,6 +4,7 @@ import pytest +from packages.auth import spend as spend_mod from packages.auth.spend import ( MICROCENTS_PER_CENT, charge_budget, @@ -91,8 +92,8 @@ async def test_stale_identity_map_cannot_clobber_a_concurrent_charge(tmp_sqlite_ """A session that loaded the key before a concurrent charge must not flush its stale value over the DB's atomic result. - #161's documented usage runs charge_budget on the same session that - validate_api_key already used to load the ApiKey. Without + The request path runs charge_budget on the same session that auth already + used to load the ApiKey. Without synchronize_session=False the ORM's 'auto' sync evaluates the SET in Python against that stale identity-map copy, marks it dirty, and the commit flushes it as a plain unguarded UPDATE — silently dropping the other @@ -138,3 +139,174 @@ async def test_stale_identity_map_cannot_clobber_a_concurrent_charge(tmp_sqlite_ def test_microcent_conversion_constant(): assert MICROCENTS_PER_CENT == 10_000 + + +@pytest.fixture +async def park_factory(tmp_sqlite_url): + """Session factory on a fresh schema, installed as the app's global factory. + + The park regression tests below need the module-global session factory + (the fold and the probes do not take a session), while the charge tests + above use an explicit session. Installed and cleared like the fixture in + test_budget_park.py. + """ + from sqlalchemy.ext.asyncio import async_sessionmaker + + from packages.db import session as session_mod + from packages.db.engine import build_engine + from packages.db.models.base import Base + + engine = build_engine(tmp_sqlite_url) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + maker = async_sessionmaker(engine, expire_on_commit=False) + previous = session_mod._session_factory + session_mod._session_factory = maker + spend_mod._unsettled.clear() + try: + yield maker + finally: + session_mod._session_factory = previous + spend_mod._unsettled.clear() + await engine.dispose() + + +async def _make_park_key(maker, *, key_id: str, spent: int = 0) -> str: + from packages.db.models.api_key import ApiKey + + async with maker() as s: + s.add( + ApiKey( + id=key_id, + workspace_id="default", + name=key_id, + key_hash=f"h-{key_id}", + key_prefix=f"p-{key_id}", + spent_microcents=spent, + ) + ) + await s.commit() + return key_id + + +async def test_durability_probe_compares_amounts_not_existence(park_factory): + """A row holding less than the obligation is not durable for it. + + Regression for the P1 where `_park_is_durable` answered True on mere row + existence: a 100-microcent row must not vouch for a 150-microcent hold. + """ + from packages.db.models.budget_park import BudgetPark + + await _make_park_key(park_factory, key_id="k") + async with park_factory() as s: + s.add(BudgetPark(trace_id="t1", api_key_id="k", microcents=100)) + await s.commit() + + assert ( + await spend_mod._park_is_durable( + trace_id="t1", api_key_id="k", expected_microcents=150 + ) + is False + ) + assert ( + await spend_mod._park_is_durable( + trace_id="t1", api_key_id="k", expected_microcents=100 + ) + is True + ) + + +async def test_pending_counts_hold_excess_over_durable_row(park_factory): + """A 100-microcent row with a 150 hold still counts 150 toward the cap. + + Regression for the P1 where pending zeroed the whole hold when its trace + had any row: the uncovered 50 must keep counting. + """ + from packages.db.models.budget_park import BudgetPark + + await _make_park_key(park_factory, key_id="k") + async with park_factory() as s: + s.add(BudgetPark(trace_id="t1", api_key_id="k", microcents=100)) + await s.commit() + spend_mod._hold("k", "t1", 150) + + assert await spend_mod.pending_parked_spend("k") == 150 + + +async def test_fold_keeps_hold_excess_when_refile_cannot_land(park_factory, monkeypatch): + """Row 100, hold 150, re-file fails: the fold bills 100 and keeps 50 held. + + Regression for the P1's concrete scenario (write outage: the row stays at + 100 while the hold is 150). The excess must survive under a fresh identity, + never be dropped. + """ + from packages.db.models.budget_park import BudgetPark + + await _make_park_key(park_factory, key_id="k", spent=0) + async with park_factory() as s: + s.add(BudgetPark(trace_id="t1", api_key_id="k", microcents=100)) + await s.commit() + spend_mod._hold("k", "t1", 150) + + async def _failing_insert(**kw): + return False + + monkeypatch.setattr(spend_mod, "_insert_park", _failing_insert) + moved = await spend_mod.settle_parked_spend("k", cap_microcents=10_000) + + assert moved == 100 + async with park_factory() as s: + assert await spend_mod.read_spent(s, "k") == 100 + assert sum(spend_mod._unsettled.get("k", {}).values()) == 50 + + +async def test_full_bill_delete_is_amount_guarded(park_factory): + """A delete naming a stale amount matches nothing; the raised row survives. + + Regression for the P1 where the full-bill DELETE matched by trace_id + alone: a concurrent re-file raising the row 100 -> 175 must make the stale + fold conflict instead of deleting 175 while billing 100. + """ + from sqlalchemy import delete, select + + from packages.db.models.budget_park import BudgetPark + + await _make_park_key(park_factory, key_id="k", spent=0) + async with park_factory() as s: + s.add(BudgetPark(trace_id="t", api_key_id="k", microcents=100)) + await s.commit() + + # The concurrent re-file lands first and raises the row. + async with park_factory() as s: + from sqlalchemy import update + + await s.execute( + update(BudgetPark) + .where(BudgetPark.trace_id == "t", BudgetPark.api_key_id == "k") + .values(microcents=175) + ) + await s.commit() + + # The stale fold's guarded delete (amount 100) must match nothing. + async with park_factory() as s: + stale = await s.execute( + delete(BudgetPark).where( + BudgetPark.trace_id == "t", + BudgetPark.api_key_id == "k", + BudgetPark.microcents == 100, + ) + ) + assert stale.rowcount == 0 + await s.rollback() + + # And a fresh fold bills the raised amount in full — nothing is lost. + moved = await spend_mod.settle_parked_spend("k", cap_microcents=10_000) + assert moved == 175 + async with park_factory() as s: + assert await spend_mod.read_spent(s, "k") == 175 + remaining = ( + await s.execute( + select(BudgetPark.trace_id).where(BudgetPark.api_key_id == "k") + ) + ).all() + assert remaining == []