diff --git a/app/budget_pricing.py b/app/budget_pricing.py new file mode 100644 index 00000000..30c95b8d --- /dev/null +++ b/app/budget_pricing.py @@ -0,0 +1,220 @@ +"""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 (image, audio) 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. + + The one-character floor is deliberately approximate, not an exact cost: a + provider can bill an image for far more than one token, but its + tokenization is unknowable here, and counting base64 payload bytes would + overcharge by orders of magnitude more. The floor's job is only to keep a + delivered multimodal completion from settling at zero. + """ + 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/app/routes/chat.py b/app/routes/chat.py index ab9200d6..bc607437 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -12,11 +12,13 @@ import time import uuid from collections.abc import AsyncGenerator, AsyncIterable, Callable +from typing import Awaitable import anyio import structlog from fastapi import APIRouter, Depends, HTTPException, Request from fastapi.responses import JSONResponse, StreamingResponse +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app import prompt_cache, router_cache @@ -26,12 +28,27 @@ choose_auto_model, required_capabilities, ) +from app.budget_pricing import ( + blocking_delivery_chars, + countable_usage, + estimate_usage, + message_chars, + text_chars, + tool_call_text, +) from app.config import get_settings from app.deps import get_db, get_key_context from app.protocols.sse import AdapterError from app.quality_scores import resolve_model_metrics from app.schemas import ChatCompletionRequest +from packages.auth.spend import ( + MICROCENTS_PER_CENT, + budget_precheck, + charge_budget, + record_unsettled_spend, +) from packages.auth.types import KeyContext +from packages.db import session as session_mod from packages.db.models.request_log import RequestLog from packages.litellm_adapter.catalog import CATALOG, CATALOG_BY_ID from packages.litellm_adapter.hosted_fallback import ( @@ -58,6 +75,196 @@ # retries hold the (already [DONE]) stream open. Tests shrink this. _LOG_COMMIT_BACKOFF_S: tuple[float, ...] = (0.1, 0.4) +# Endings a settlement can have. Handlers record which one happened; the +# charge is decided from it in exactly one place (see `_unmeasured_charge`), +# so no handler can price its own ending differently from its siblings. +_STREAM_IN_FLIGHT = "in_flight" +_STREAM_COMPLETED = "completed" +_STREAM_CLIENT_DISCONNECT = "client_disconnect" +_STREAM_UPSTREAM_ERROR = "upstream_error" +_STREAM_ADAPTER_ERROR = "adapter_error" + + +def _cost_of_usage( + usage: dict, + *, + model_id: str | None, + fallback_model: str | None = None, +) -> int: + """Catalog cost of a token count, for a delivery the provider never priced.""" + return _compute_cost_microcents( + litellm_cost_usd=None, + model_id=model_id, + fallback_model=fallback_model, + input_tokens=usage.get("prompt_tokens", 0) or 0, + output_tokens=usage.get("completion_tokens", 0) or 0, + ) + + +def _unmeasured_charge( + *, + delivered: bool, + ending: str, + prompt_chars: int, + completion_chars: int, + model_id: str | None, + fallback_model: str | None = None, +) -> tuple[int, dict | None]: + """Charge for a request the provider never measured, from stream facts. + + The single decision point for every unmeasured ending, so the streaming and + blocking paths cannot drift apart. The rules, in order: + + - Nothing delivered: a provider or adapter failure that returned nothing + cost nothing measurable and settles at zero; a client that hung up + before the first byte still caused the prompt to go upstream, so it pays + the prompt estimate. + - Content delivered on a known-free model: charged nothing. + - Content delivered otherwise: priced from the delivery estimate, which is + the most we can honestly know about a delivery nobody measured. This is + the same price whether the stream ended normally or faulted — the ending + does not change what the client received, so it does not change the bill. + + Every delivered byte is counted (`completion_chars`), so pricing per + character already stops a client from opting out of measurement and + streaming for free. Charging the whole remaining allowance on top of that + would not fail closed more safely; it would let one request from a provider + that omits token counts consume a key's entire lifetime budget. + + A disconnect *after* a completed stream is not a separate ending: `ending` + is already `_STREAM_COMPLETED` by then, so where the client chose to hang up + cannot change the bill. + + Returns the charge together with the token counts it was derived from, so + the request-log row explains its own `cost_microcents` instead of recording + a cost with zero tokens behind it. The counts are `None` for the branch + that charges nothing, where no honest token estimate exists. + """ + # A settlement that ran without an ending recorded can only mean the + # consumer left before the provider finished. Priced as the hangup it is. + effective_ending = ( + _STREAM_CLIENT_DISCONNECT + if ending == _STREAM_IN_FLIGHT + else ending + ) + if not delivered: + if effective_ending == _STREAM_CLIENT_DISCONNECT: + estimate = estimate_usage(prompt_chars, 0) + return ( + max(1, _cost_of_usage( + estimate, + model_id=model_id, + fallback_model=fallback_model, + )), + estimate, + ) + return 0, None + # A catalog entry priced 0.0 in and 0.0 out is a known-free model: the + # completion costs nothing, regardless of how it ended. The ending does not + # change what the upstream agreed to bill for that delivery. + m = _lookup_priced_model(model_id) or _lookup_priced_model(fallback_model) + if ( + m is not None + and not m.input_cost_per_token + and not m.output_cost_per_token + ): + return 0, None + estimate = estimate_usage(prompt_chars, completion_chars) + # The charge is floored at one microcent, the smallest amount the counter can + # represent: a short answer estimates below that and would truncate to zero, + # which is exactly what the cap must not allow. + return ( + max(1, _cost_of_usage( + estimate, model_id=model_id, fallback_model=fallback_model, + )), + estimate, + ) + + +async def _settlement_is_durable(trace_id: str) -> bool: + """Whether the request-log row for this trace_id is committed. + + The row and the budget charge are one transaction, so the row is proof the + spend is already counted. It reads on its own session because the failed + attempt's session is closed or rolled back by the time the give-up runs, and + answers False when it cannot read at all: during a real outage the park is + the only thing keeping the cap honest, so an unreadable DB must look + not-durable rather than durable. + """ + factory = session_mod._session_factory + if factory is None: + return False + # Shielded: this runs while a give-up is already unwinding a cancellation, + # and the probe must still answer so the charge is parked, not dropped. + with anyio.CancelScope(shield=True): + try: + async with factory() as s: + return ( + await s.scalar( + select(RequestLog.id).where(RequestLog.trace_id == trace_id) + ) + ) is not None + except Exception: + # Broad catch is deliberate: any DB error (connectivity, auth, + # schema) means we cannot verify durability, so fail closed. + # + # This is a choice between two wrongs, made knowingly. When the + # commit landed but its ack was lost AND this read also fails, we + # cannot tell "already charged" from "never written", so we park and + # the key is billed twice. The alternative — trusting an unreadable + # database to have recorded the charge — reopens the cap for exactly + # the key whose settlement we are unwinding, which is the failure + # the park exists to prevent. Over-charging a capped key is the + # smaller harm; under-charging it is the one that lets spend through + # uncounted. + return False + + +async def _give_up_settlement( + kc: KeyContext, + trace_id: str, + amount: int, + attempts: int, + error: BaseException, + persisted: Callable[[], Awaitable[bool]], +) -> None: + """Last resort for a settlement that is not durable and will not be retried. + + The log row dies with the charge (one transaction), so nothing anywhere + remembers this cost. Park it — one idempotent row keyed by this + settlement's `trace_id` — against the key's cap, rather than leaving the + cap open for whoever reads the warning (see `packages.auth.spend`). + + `persisted()` is asked first, because the last attempt has no retry left to + run the trace_id check: a commit that applied but whose ack was lost (or a + cancellation that landed after it) leaves the charge already in + `spent_microcents`, and parking on top of it would bill one delivery twice. + """ + logger.warning("request_log_commit_failed", error=str(error), attempts=attempts) + if not await persisted(): + await record_unsettled_spend( + trace_id=trace_id, api_key_id=str(kc.key_id), microcents=amount, + ) + + +def _has_known_price( + *, + litellm_cost_usd: float | None, + model_id: str | None, + fallback_model: str | None = None, +) -> bool: + """True when a delivered completion's cost is measurable, even if it is 0. + + Mirrors `_compute_cost_microcents`' two tiers without computing: an + authoritative LiteLLM cost, or any catalog entry. A 0.0/0.0 entry is a + known-free model, not an unknown cost. Tokens with neither are + unknown. + """ + if litellm_cost_usd is not None and litellm_cost_usd > 0: + return True + m = _lookup_priced_model(model_id) or _lookup_priced_model(fallback_model) + return m is not None + def _chunk_to_dict(chunk) -> dict: """Normalize a litellm chunk (Pydantic model or dict) into a plain dict. @@ -209,8 +416,13 @@ async def _build_log_row( meta = response.get("_orca_meta", {}) or {} usage = response.get("usage", {}) or {} resolved = actual_resolved or response.get("model") or requested_model - input_t = usage.get("prompt_tokens", 0) or 0 - output_t = usage.get("completion_tokens", 0) or 0 + # Providers report usage under different keys: OpenAI-style + # prompt_tokens/completion_tokens, Anthropic-style input_tokens/ + # output_tokens (what /v1/messages forwards). Read both so a + # differently-keyed frame does not normalize to zero tokens — the + # token count is what the fail-closed budget gates key on. + input_t = usage.get("prompt_tokens", 0) or usage.get("input_tokens", 0) or 0 + output_t = usage.get("completion_tokens", 0) or usage.get("output_tokens", 0) or 0 return RequestLog( workspace_id=str(kc.workspace_id), api_key_id=str(kc.key_id), @@ -482,9 +694,46 @@ async def execute_chat( resolved_model = candidates[0] body.model = candidates[0] # mutate for downstream completion call + # `budget_limit_cents` is a hard lifetime cap. The check runs only after the + # request has passed every pre-dispatch validation (model allowlist, provider + # deployability), so a request we reject before touching an upstream never + # consumes budget. The real cost is only known once the upstream response or + # stream completes, so it is recorded atomically in `_settle_budget`: the + # `UPDATE spent = spent + actual WHERE spent + actual <= cap` guard makes + # this safe under concurrency and never lets the counter pass the cap. + # + # (cap_microcents, spent_at_precheck) for a capped key, None otherwise. The + # settle path and the charge estimate both read it, so the number a request + # is rejected on and the number it is charged against are the same. + budget: tuple[int, int] | None = None + if kc.budget_limit_cents is not None: + cap = kc.budget_limit_cents * MICROCENTS_PER_CENT + spent = await budget_precheck(db, str(kc.key_id), cap) + if spent >= cap: + raise HTTPException( + status_code=429, + detail="API key budget exhausted (lifetime cap reached).", + ) + budget = (cap, spent) + started_perf = time.perf_counter() completion_kwargs = body.model_dump(exclude_none=True) + async def _settle_budget(session, actual_microcents: int, *, commit: bool = True) -> None: + """Record `actual_microcents` of spend against the cap, if any. + + No-op when the key has no budget cap. When `commit` is False the UPDATE is + executed but not committed, so the caller commits it in the same + transaction as the request-log write — making the row and the charge one + atomic unit. Idempotency across retries comes from the row's trace_id + (a persisted trace_id proves the charge also landed), not from a + process-local flag. + """ + if budget is None: + return + cap, _spent = budget + await charge_budget(session, str(kc.key_id), cap, actual_microcents, commit=commit) + # Build the LiteLLM fallbacks argument from the auto candidate list. # Format: [{primary_model_name: [fallback_1, fallback_2, ...]}] # LiteLLM's async_function_with_fallbacks reads this from kwargs and @@ -553,6 +802,7 @@ async def execute_chat( log.cost_microcents = 0 db.add(log) try: + await _settle_budget(db, 0, commit=False) await db.commit() except Exception as commit_err: logger.warning("request_log_commit_failed", error=str(commit_err)) @@ -574,13 +824,13 @@ async def execute_chat( # the client decide what to do. if body.stream: # Auto-inject `stream_options.include_usage=True` if the client - # didn't set it. Without this, OpenAI/LiteLLM streaming responses - # omit the `usage` field entirely — chunks have no token counts, - # so our log row gets input=0, output=0 and the cost calculation - # rounds to zero. Almost no client knows to opt-in to this flag, - # which would silently zero out streaming spend in the dashboard. - # Honor an explicit `include_usage=False` from the client if they - # really want to disable it (e.g. wire-format compatibility tests). + # didn't set it, so streaming responses carry token counts and we bill + # the measured cost. A budgeted key needs it: without a usage frame the + # settlement falls back to pricing the delivery, which under-counts. + # A client that explicitly asked for `include_usage: False` still gets + # it — overriding that would hand every budgeted client an extra + # usage-only frame they asked not to receive. The cost of honouring it + # is bounded and visible: the stream settles from what it delivered. existing_so = completion_kwargs.get("stream_options") or {} if "include_usage" not in existing_so: completion_kwargs["stream_options"] = {**existing_so, "include_usage": True} @@ -606,6 +856,7 @@ async def _log_pre_stream_failure(status: int, err_type: str | None) -> None: ) db.add(log) try: + await _settle_budget(db, 0, commit=False) await db.commit() except Exception as commit_err: # Roll back so the request-scoped session is not left in a @@ -649,12 +900,21 @@ async def sse() -> AsyncGenerator[str, None]: agg_provider = "unknown" agg_fallback = False agg_latency = 0 + # Characters of assistant text handed to the client — the only + # measure of what a stream delivered when the provider never + # reported usage (see `app.budget_pricing`). + agg_output_chars = 0 # The first chunk's `model` field tells us what LiteLLM actually # served (could be a cascaded fallback, not the resolved primary). agg_model: str | None = None status_code = 200 error_type: str | None = None log_written = False + # How the stream ended. Handlers record the ending and nothing + # else — the charge is decided from these facts in one place + # (`_unmeasured_charge`), so a handler cannot price its own ending + # differently from a sibling. + stream_ending = _STREAM_IN_FLIGHT async def _finalize() -> None: """Write the request log row exactly once. @@ -712,9 +972,6 @@ async def _finalize() -> None: # _build_log_row fall back to the measured wall-clock # latency instead of persisting a constant 0. synthetic["_orca_meta"]["latency_ms"] = agg_latency - from sqlalchemy import select - - from packages.db import session as session_mod # The row's VALUES are computed exactly once — latency is # measured here, before any commit attempt, so retry @@ -752,27 +1009,107 @@ async def _already_persisted(s) -> bool: select(RequestLog.id).where(RequestLog.trace_id == row_values["trace_id"]) )) is not None + async def _durable() -> bool: + return await _settlement_is_durable(row_values["trace_id"]) + + def _settlement_amount() -> int: + """Budget charge for this request, in microcents. + + A usage frame carrying countable token totals is the + billing signal: the cost is known, and charging more would + over-bill a quantity the row already accounts for. A frame + without those keys is not a measurement at all. + + When nothing was measured, the charge comes from + `_unmeasured_charge` and the stream's recorded ending — + never from which exception handler happened to run. The + rules that matters here: + + - we asked for usage and the provider sent none: priced + from the delivery, the same as every other unmeasured + ending — the client received those characters, so they + are what gets charged; + - the client declined usage frames up front: priced from + the delivery, since we were told not to ask; + - a hangup or a fault on either side: priced from what + reached the client; + - nothing delivered: charged nothing, unless a client hung + up, in which case the prompt it caused was still billed. + + The amount is written back into the row, not just into the + counter: a charge only the counter saw would leave a key + exhausted by an amount nothing in its own request history + accounts for. + """ + actual = row_values.get("cost_microcents") or 0 + if ( + budget is not None + and countable_usage( + agg_usage, delivered=agg_output_chars > 0 + ) + and not actual + and not _has_known_price( + litellm_cost_usd=(agg_usage or {}).get("cost_usd"), + model_id=row_values.get("model_resolved"), + fallback_model=row_values.get("model_requested"), + ) + ): + # Measured but unpriceable, same as the blocking twin: + # tokens with no price any tier can honour (a custom + # upstream cannot cost, or the model is absent from the + # catalog). Fail closed — the counter must not stand + # still while the upstream still bills us. + actual = 1 + row_values["cost_microcents"] = actual + if budget is not None and not countable_usage( + agg_usage, delivered=agg_output_chars > 0 + ): + charge, estimate = _unmeasured_charge( + delivered=agg_output_chars > 0, + ending=stream_ending, + prompt_chars=sum( + message_chars(m.content, m.tool_calls) + for m in body.messages + ), + completion_chars=agg_output_chars, + model_id=row_values.get("model_resolved"), + fallback_model=row_values.get("model_requested"), + ) + actual = max(actual, charge) + row_values["cost_microcents"] = actual + if estimate is not None: + # The row explains its own charge: an estimated + # delivery records the estimated tokens, so a + # non-zero cost never sits on a zero-token row. + row_values["input_tokens"] = estimate["prompt_tokens"] + row_values["output_tokens"] = estimate["completion_tokens"] + return actual + async def _commit_row(*, retry: bool) -> None: - """INSERT + COMMIT the row on a session of its own. - - Only a failing `commit()` propagates; a failure while - closing the session AFTER the commit returned is - swallowed — the row is already in. A retry is - idempotent: it first looks the trace_id up, so a COMMIT - that landed but whose ack was lost on the wire - (PostgreSQL, connection dropped mid-ack) is not - inserted a second time — and the shared primary key - would reject a duplicate anyway. + """Persist the request-log row and charge the budget in ONE commit. + + The INSERT and the budget charge share a single transaction. If it + commits, both are durable; if it fails, both roll back and the + retry re-runs both. Because the charge lands in the same commit as + the row, a persisted trace_id proves the charge also landed — so a + retry returns without re-charging. The charge is therefore applied + exactly once per request: never doubled (on a commit-ack-loss + retry) and never dropped. """ + # Resolved before the row is built: a fail-closed charge + # raises `row_values["cost_microcents"]`, and the object + # inserted must carry the amount the counter will move by. + settlement = _settlement_amount() log = RequestLog(**row_values) if session_mod._session_factory is None: # Test-only fallback (the app always installs a # factory): the request-scoped session has to be # rolled back before a retry can reuse it. - if retry and await _already_persisted(db): + if retry and (await _already_persisted(db)): return db.add(log) try: + await _settle_budget(db, settlement, commit=False) await db.commit() except Exception: try: @@ -783,9 +1120,10 @@ async def _commit_row(*, retry: bool) -> None: return s = session_mod._session_factory() try: - if retry and await _already_persisted(s): + if retry and (await _already_persisted(s)): return s.add(log) + await _settle_budget(s, settlement, commit=False) await s.commit() finally: try: @@ -824,32 +1162,57 @@ async def _commit_row(*, retry: bool) -> None: except BaseException: # Cancelled during the backoff: nothing is # in flight, the row is given up on — say - # so, then propagate like the arm below. - logger.warning( - "request_log_commit_failed", - error=str(commit_err), attempts=attempt, + # so, then propagate. An exception raised in a + # handler is not caught by this try's other + # arms, so the give-up below does not run a + # second time for this settlement. + await _give_up_settlement( + kc, row_values["trace_id"], _settlement_amount(), + attempt, commit_err, _durable, ) raise continue - logger.warning( - "request_log_commit_failed", - error=str(commit_err), attempts=attempt, + await _give_up_settlement( + kc, row_values["trace_id"], _settlement_amount(), + attempt, commit_err, _durable, ) except BaseException: # CancelledError aimed at us, not at the commit — # wait the in-flight attempt out so a row about to # land isn't dropped, then let the cancellation - # propagate (no further attempts: we are being torn - # down). + # propagate. try: await commit_task except Exception as commit_err: - logger.warning( - "request_log_commit_failed", - error=str(commit_err), attempts=attempt, + await _give_up_settlement( + kc, row_values["trace_id"], _settlement_amount(), + attempt, commit_err, _durable, ) except BaseException: - pass + # The detached task itself was cancelled, so its + # write is gone and the charge died with it. Retrying + # outside a shield would be cancelled before it + # starts; inside one it runs to completion, and + # `_commit_row(retry=True)` is idempotent on + # trace_id — a commit that landed with a lost ack + # is recognised and not charged twice. + try: + with anyio.CancelScope(shield=True): + last_try = asyncio.ensure_future( + _commit_row(retry=True) + ) + await asyncio.shield(last_try) + except BaseException as retry_err: + # Nothing more can be written: the last-chance + # retry failed too, whether it was cancelled or + # simply refused by the database. The charge is + # parked rather than logged, because the cap is + # only honest if a delivered cost somewhere still + # counts against it. + await _give_up_settlement( + kc, row_values["trace_id"], _settlement_amount(), + attempt, retry_err, _durable, + ) raise last_d: dict = {} @@ -868,8 +1231,25 @@ async def _commit_row(*, retry: bool) -> None: agg_usage = d["usage"] if d.get("model"): agg_model = d["model"] + for choice in d.get("choices") or []: + if isinstance(choice, dict): + delta = choice.get("delta") or {} + chars = text_chars(delta.get("content")) or 0 + # Tool calls are delivered content the upstream bills + # for; count them like the blocking path does via + # tool_call_text so both paths price the same + # delivery consistently. Arguments sent as an object + # rather than a string have no honest character count. + chars += len(tool_call_text(delta.get("tool_calls"))) + agg_output_chars += chars last_d = d yield f"data: {json.dumps(d, separators=(',', ':'))}\n\n" + # The provider's stream is done: everything after this point is + # our own framing. Recording it HERE is what makes a client that + # hangs up at the trailing [DONE] frame pay the same charge as + # one that stays connected — where it disconnects is its choice, + # not a fact about the delivery. + stream_ending = _STREAM_COMPLETED # A trailing frame that already carries `usage` is the usage # frame, whether or not `choices` is empty. LiteLLM's # include_usage chunk uses @@ -918,6 +1298,16 @@ async def _commit_row(*, retry: bool) -> None: "chat_completion_stream_client_disconnect", served_model=agg_model, ) + # Record the ending; the charge is decided from it in + # `_settlement_amount`, together with every other ending. A + # stream that already completed keeps `completed`: hanging up + # at the trailing [DONE] frame changes nothing about what was + # delivered, and pricing it as a disconnect would let a client + # buy the estimate for a full answer by leaving at the last + # byte. A mid-stream bail keeps its own pricing (including the + # prompt-only charge for a hangup before the first byte). + if stream_ending != _STREAM_COMPLETED: + stream_ending = _STREAM_CLIENT_DISCONNECT # Cleanup MUST actually complete before we unwind, not just # be scheduled. `asyncio.shield()` here would NOT wait — the # outer await re-raises CancelledError immediately under @@ -953,6 +1343,11 @@ async def _commit_row(*, retry: bool) -> None: logger.warning( "chat_completion_stream_adapter_error", served_model=agg_model, ) + # An adapter fault is our bug, not a choice the caller made, so + # it prices like the other faults: from what reached the client. + # Nothing delivered settles at the 0 the row already records. + if stream_ending != _STREAM_COMPLETED: + stream_ending = _STREAM_ADAPTER_ERROR aclose = getattr(stream_obj, "aclose", None) with anyio.CancelScope(shield=True): if aclose is not None: @@ -987,6 +1382,13 @@ async def _commit_row(*, retry: bool) -> None: "chat_completion_stream_error", error=str(exc), error_type=error_type, ) + # Record the ending BEFORE yielding the error frame and the + # sentinel. If the client disconnects during that yield, + # GeneratorExit unwinds straight through finally, so anything + # decided after the yield would be skipped and this fault would + # be priced as a normal completion. + if stream_ending != _STREAM_COMPLETED: + stream_ending = _STREAM_UPSTREAM_ERROR err_body = { "error": { "message": f"Upstream provider error: {exc}", @@ -1088,11 +1490,166 @@ async def _commit_row(*, retry: bool) -> None: # _build_log_row would otherwise default to via requested_model). actual_resolved=actual_resolved or resolved_model, ) - db.add(log) - try: - await db.commit() - except Exception as commit_err: - logger.warning("request_log_commit_failed", error=str(commit_err)) + # Persist the log row and the budget charge atomically (same transaction), + # retrying transient commit failures so a budgeted key is never under- + # charged when the DB is stressed — mirroring the streaming path. A + # persisted trace_id proves both landed, so a retry skips rather than + # double-charging. + settle_amount = log.cost_microcents + # A budgeted key whose successful response carries no countable usage + # (a provider that answered without token counts, or with a frame that + # has none) has an unknown cost. The response is here in full, so what + # it actually returned can be priced: the same delivery-estimate every + # unmeasured ending uses. Charging the whole remaining allowance would + # bill a customer their lifetime budget for one ordinary request whose + # upstream omits a field, which is a far worse failure than the one the + # rule exists to prevent. + # + # Gated on a delivered completion dict AND delivered content: a request + # that failed before the upstream answered (response == {}, e.g. the + # re-raised HTTPException above, whose status_code never left 200) and + # a successful but empty completion both charge their recorded ~0 cost, + # mirroring the cache-hit and pre-stream-failure paths. A completion the + # provider priced through _orca_meta keeps that authoritative cost. + if ( + budget is not None + and status_code < 400 + and isinstance(response, dict) + and response + and ( + # Unmeasured: no countable usage at all. Priced from what the + # response returned, which is in full here. + ( + not countable_usage( + response.get("usage"), + delivered=blocking_delivery_chars(response)[0], + ) + and blocking_delivery_chars(response)[0] + ) + or ( + # Measured but unpriceable: tokens with no price we can + # honour (a custom upstream cannot cost, or the model is + # absent from the catalog). Fail closed — the counter must + # not stand still while the upstream still bills us. + (log.input_tokens or log.output_tokens) + and not (log.cost_microcents or 0) + and not _has_known_price( + litellm_cost_usd=(response.get("usage") or {}).get("cost_usd") + or (response.get("_orca_meta") or {}).get("cost_usd"), + model_id=log.model_resolved, + fallback_model=log.model_requested, + ) + ) + ) + ): + if countable_usage( + response.get("usage"), + delivered=blocking_delivery_chars(response)[0], + ): + # Measured but unpriceable. The counter must not stand still + # while the upstream still bills us, so charge the floor: one + # microcent per such request exhausts the cap eventually, which + # is the whole point, without one request consuming the entire + # remaining lifetime budget. + settle_amount = max(log.cost_microcents or 0, 1) + else: + delivered, completion_chars = blocking_delivery_chars(response) + charge, estimate = _unmeasured_charge( + delivered=delivered, + ending=( + _STREAM_COMPLETED if delivered else _STREAM_UPSTREAM_ERROR + ), + prompt_chars=sum( + message_chars(m.content, m.tool_calls) for m in body.messages + ), + completion_chars=completion_chars, + model_id=log.model_resolved, + fallback_model=log.model_requested, + ) + settle_amount = max(log.cost_microcents or 0, charge) + if estimate is not None: + # The row accounts for its own cost, as on the streaming path. + log.input_tokens = estimate["prompt_tokens"] + log.output_tokens = estimate["completion_tokens"] + log.cost_microcents = settle_amount + + # Values are snapshotted once (latency is measured in _build_log_row, + # before any commit attempt, so retry backoff never inflates it) and + # each attempt inserts a fresh ORM object carrying the same id/trace_id + # — mirroring the streaming path, so a retry works regardless of what + # the rollback left the old object as. + log_values = { + c.key: getattr(log, c.key) + for c in RequestLog.__table__.columns + if getattr(log, c.key) is not None + } + max_attempts = len(_LOG_COMMIT_BACKOFF_S) + 1 + + async def _durable() -> bool: + return await _settlement_is_durable(log_values["trace_id"]) + + for attempt in range(1, max_attempts + 1): + try: + if attempt > 1: + await db.rollback() + if attempt > 1 and ( + await db.scalar( + select(RequestLog.id).where(RequestLog.trace_id == log.trace_id) + ) + ) is not None: + break # already durable (log + charge committed) + db.add(RequestLog(**log_values)) + await _settle_budget(db, settle_amount, commit=False) + await db.commit() + break + except BaseException as commit_err: + # A cancellation (client disconnect) during the commit is not a + # retryable failure, but the charge must still be parked if the + # row did not land. Park, then let the cancellation propagate. + if isinstance(commit_err, asyncio.CancelledError): + with anyio.CancelScope(shield=True): + try: + await db.rollback() + except BaseException: + pass + await _give_up_settlement( + kc, log_values["trace_id"], settle_amount, + attempt, commit_err, _durable, + ) + raise + if not isinstance(commit_err, Exception): + raise + try: + await db.rollback() + except Exception: + pass + if attempt == max_attempts: + # Last attempt: probe whether the write landed and, if it + # did not, park the obligation. Nothing anywhere else + # remembers this cost — the row dies with the charge in one + # transaction — so a silent drop here would reopen the cap + # for exactly the key whose settlement failed. + await _give_up_settlement( + kc, log_values["trace_id"], settle_amount, + attempt, commit_err, _durable, + ) + break + logger.info( + "request_log_commit_retry", error=str(commit_err), attempt=attempt, + ) + try: + await asyncio.sleep(_LOG_COMMIT_BACKOFF_S[attempt - 1]) + except BaseException: + # Cancelled during the backoff: nothing is in flight and the + # row is given up on — say so, then propagate. An exception + # raised from a handler is not caught by this try's other + # arms, so the give-up below does not run a second time for + # this settlement. + await _give_up_settlement( + kc, log_values["trace_id"], settle_amount, + attempt, commit_err, _durable, + ) + raise hosted_fallback = _meta_hosted_fallback(response) if isinstance(response, dict) and "_orca_meta" in response: diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 27649968..82c9f2b5 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,507 @@ 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: + """Whether the durable row covers the entire expected obligation.""" + 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: + 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 False + + +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, `_park_is_durable` is called to probe + whether the write landed anyway. Returns False only when the obligation is + genuinely 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 concurrent give-up. A hold that + # merged further settlements into this one now carries more + # than the row does, so raise the row to it rather than + # reporting success and dropping the difference. Never lowered: + # the larger amount is the one the cap needs. + await s.rollback() + await s.execute( + update(BudgetPark) + .where( + BudgetPark.trace_id == trace_id, + BudgetPark.microcents < microcents, + ) + .values(microcents=microcents) + ) + await s.commit() + return True + except asyncio.CancelledError: + raise + except Exception: + return await _park_is_durable( + 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( + _park_is_durable( + 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 +541,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 +594,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 +620,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/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py new file mode 100644 index 00000000..f2237825 --- /dev/null +++ b/tests/integration/test_budget_enforcement.py @@ -0,0 +1,740 @@ +"""Budget enforcement in the request path — the cap actually stops traffic. + +`ApiKey.budget_limit_cents` is a hard lifetime cap: a key at its cap must be +rejected with 429, and a request that is served must leave the counter moved by +exactly what it cost. Both halves matter, because a cap that is checked but never +charged would re-grant every exhausted key on the next restart. + +Real app, real SQLite, a mocked router — the same shape as the log-durability +tests, so these cover the same paths (streaming and blocking) end to end. +""" + +from __future__ import annotations + +import time +from unittest.mock import AsyncMock + +import pytest +from sqlalchemy import event, select +from sqlalchemy.exc import OperationalError +from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker +from sqlalchemy.orm import Session + +from packages.auth.spend import MICROCENTS_PER_CENT +from packages.db import session as session_mod +from packages.db.engine import build_engine +from packages.db.models.api_key import ApiKey +from packages.db.models.base import Base +from packages.db.models.budget_park import BudgetPark +from packages.db.models.request_log import RequestLog + + +def _mark_log_flush(session, _flush_context, _instances) -> None: + """Flag a flush that carries a RequestLog, so its commit can be failed. + + Read from `.new` inside commit() instead would miss it: the budget charge + runs an UPDATE first and that autoflushes the pending INSERT, so by the time + commit() runs the row is no longer "new". + """ + if any(isinstance(o, RequestLog) for o in session.new): + session._carrying_log = True + + +class _FailingCommitSession(AsyncSession): + """AsyncSession that fails the commit of a session carrying a RequestLog. + + The auth middleware shares this factory and commits no log row, so it is + untouched. + """ + + async def commit(self): + if getattr(self.sync_session, "_carrying_log", False): + raise OperationalError("COMMIT", {}, Exception("database is locked")) + return await super().commit() + + +def _chunks() -> list[dict]: + now = int(time.time()) + return [ + { + "id": "chatcmpl-1", "object": "chat.completion.chunk", + "model": "gpt-4o-mini", "created": now, + "choices": [{"index": 0, "delta": {"content": "Hello"}, + "finish_reason": None}], + }, + { + "id": "chatcmpl-1", "object": "chat.completion.chunk", + "model": "gpt-4o-mini", "created": now, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6}, + }, + ] + + +async def _stream_iter(): + for chunk in _chunks(): + yield chunk + + +def _blocking_response() -> dict: + return { + "id": "chatcmpl-1", "object": "chat.completion", + "model": "gpt-4o-mini", "created": int(time.time()), + "choices": [{ + "index": 0, "finish_reason": "stop", + "message": {"role": "assistant", "content": "Hello there"}, + }], + "usage": {"prompt_tokens": 4, "completion_tokens": 2, "total_tokens": 6}, + "_orca_meta": {"cost_usd": 0.0001}, + } + + +@pytest.fixture +async def budget_app(tmp_sqlite_url, monkeypatch): + """App with a real DB, a mocked router, and one key with a lifetime cap.""" + monkeypatch.setenv("DATABASE_URL", tmp_sqlite_url) + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + + from app import config as cfg + cfg.get_settings.cache_clear() + + engine = build_engine(tmp_sqlite_url) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + factory = async_sessionmaker(engine, expire_on_commit=False) + session_mod._session_factory = factory + + from packages.auth.hashing import generate_api_key + full_key, key_hash, key_prefix = generate_api_key() + key = ApiKey( + workspace_id="default", + name="budgeted", + key_hash=key_hash, + key_prefix=key_prefix, + budget_limit_cents=1, # 1 cent = 10_000 microcents + ) + async with factory() as s: + s.add(key) + await s.commit() + + from app import router_cache + router_cache.invalidate_router() + + fake = AsyncMock() + + async def _acompletion(**kwargs): + if kwargs.get("stream"): + return _stream_iter() + return _blocking_response() + + fake.acompletion = AsyncMock(side_effect=_acompletion) + + async def _fake_get_router(_session): + return fake + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + monkeypatch.setattr("app.routes.chat._LOG_COMMIT_BACKOFF_S", (0.0, 0.0)) + + from app.main import create_app + app = create_app() + + # Scoped to this test: the listener is on the global Session class, so it + # must not outlive the fixture. + event.listen(Session, "before_flush", _mark_log_flush) + try: + from httpx import ASGITransport, AsyncClient + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://t", + headers={"Authorization": f"Bearer {full_key}"}, + ) as client: + yield client, factory, key.id + finally: + event.remove(Session, "before_flush", _mark_log_flush) + session_mod._session_factory = None + await engine.dispose() + + +async def _chat(client, *, stream: bool = False) -> dict: + r = await client.post( + "/v1/chat/completions", + json={ + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "stream": stream, + }, + ) + return {"status": r.status_code, "body": r.text} + + +async def _spent(factory, key_id) -> int: + async with factory() as s: + value = await s.scalar(select(ApiKey.spent_microcents).where(ApiKey.id == key_id)) + return int(value or 0) + + +async def _log_rows(factory) -> list[RequestLog]: + async with factory() as s: + return list((await s.execute(select(RequestLog))).scalars().all()) + + +async def _park_rows(factory) -> list[BudgetPark]: + async with factory() as s: + return list((await s.execute(select(BudgetPark))).scalars().all()) + + +async def test_request_under_the_cap_is_served_and_charged(budget_app): + """A blocking request records its measured cost against the cap.""" + client, factory, key_id = budget_app + + result = await _chat(client) + + assert result["status"] == 200 + assert await _spent(factory, key_id) > 0 + rows = await _log_rows(factory) + assert len(rows) == 1 + assert rows[0].cost_microcents > 0 + assert rows[0].input_tokens == 4 and rows[0].output_tokens == 2 + + +async def test_exhausted_key_is_rejected_with_429(budget_app): + """A key at its cap is refused before any upstream call is made.""" + client, factory, key_id = budget_app + + async with factory() as s: + await s.execute( + ApiKey.__table__.update() + .where(ApiKey.id == key_id) + .values(spent_microcents=1 * MICROCENTS_PER_CENT) + ) + await s.commit() + + result = await _chat(client) + + assert result["status"] == 429 + assert "budget exhausted" in result["body"] + # Nothing was served, so nothing new was recorded. + assert await _log_rows(factory) == [] + + +async def test_stream_under_the_cap_is_served_and_charged(budget_app): + """The streaming path charges the measured cost, same as the blocking path.""" + client, factory, key_id = budget_app + + result = await _chat(client, stream=True) + + assert result["status"] == 200 + assert "data: [DONE]" in result["body"] + spent = await _spent(factory, key_id) + assert spent > 0 + rows = await _log_rows(factory) + assert len(rows) == 1 + assert rows[0].cost_microcents == spent + assert rows[0].is_streaming is True + + +async def test_charge_and_log_row_land_in_one_transaction(budget_app): + """The counter and the row move together, so a logged cost is a charged cost.""" + client, factory, key_id = budget_app + + await _chat(client) + + spent = await _spent(factory, key_id) + rows = await _log_rows(factory) + assert rows[0].cost_microcents == spent + + +@pytest.fixture +async def failing_app(tmp_sqlite_url, monkeypatch): + """Same as budget_app, but every request-log commit fails.""" + monkeypatch.setenv("DATABASE_URL", tmp_sqlite_url) + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + + from app import config as cfg + cfg.get_settings.cache_clear() + + engine = build_engine(tmp_sqlite_url) + async with engine.begin() as conn: + await conn.run_sync(Base.metadata.create_all) + factory = async_sessionmaker(engine, class_=_FailingCommitSession, expire_on_commit=False) + session_mod._session_factory = factory + + from packages.auth.hashing import generate_api_key + full_key, key_hash, key_prefix = generate_api_key() + key = ApiKey( + workspace_id="default", + name="budgeted", + key_hash=key_hash, + key_prefix=key_prefix, + budget_limit_cents=1, + ) + async with factory() as s: + s.add(key) + await s.commit() + + from app import router_cache + router_cache.invalidate_router() + + fake = AsyncMock() + + async def _acompletion(**kwargs): + if kwargs.get("stream"): + return _stream_iter() + return _blocking_response() + + fake.acompletion = AsyncMock(side_effect=_acompletion) + + async def _fake_get_router(_session): + return fake + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + monkeypatch.setattr("app.routes.chat._LOG_COMMIT_BACKOFF_S", (0.0, 0.0)) + + from app.main import create_app + app = create_app() + + # Scoped to this test: the listener is on the global Session class, so it + # must not outlive the fixture. + event.listen(Session, "before_flush", _mark_log_flush) + try: + from httpx import ASGITransport, AsyncClient + async with AsyncClient( + transport=ASGITransport(app=app), base_url="http://t", + headers={"Authorization": f"Bearer {full_key}"}, + ) as client: + yield client, factory, key.id + finally: + event.remove(Session, "before_flush", _mark_log_flush) + session_mod._session_factory = None + await engine.dispose() + + +async def test_a_settlement_that_cannot_be_recorded_is_parked(failing_app): + """The row and the charge die together, so the cost is parked against the cap. + + Without this, a settlement that fails after every retry leaves the key's cap + open for a delivery that was served. + """ + client, factory, key_id = failing_app + + result = await _chat(client, stream=True) + + assert result["status"] == 200 + assert "data: [DONE]" in result["body"] + + parked = await _park_rows(factory) + assert len(parked) == 1 + assert parked[0].api_key_id == key_id + assert parked[0].microcents > 0 + # The row never landed, so the counter did not move. + assert await _spent(factory, key_id) == 0 + assert await _log_rows(factory) == [] + + +async def test_a_parked_charge_counts_against_the_cap_on_the_next_request(budget_app): + """A parked obligation is read by the pre-check and blocks dispatch. + + The park is created through the same public entry point the settlement path + uses (`record_unsettled_spend`), then a real request is sent through + `execute_chat`. If the pre-check did not fold the park into the key's spend, + this request would be served and the counter would move by only its own + cost — so a 200 here is the regression this test exists to catch. + """ + client, factory, key_id = budget_app + + from packages.auth.spend import record_unsettled_spend + + # Larger than the key's 10,000-microcent cap, so the debt alone must refuse + # the request. + await record_unsettled_spend( + trace_id="parked-larger-than-cap", + api_key_id=str(key_id), + microcents=50_000, + ) + parked = await _park_rows(factory) + assert len(parked) == 1 + assert parked[0].microcents == 50_000 + + result = await _chat(client) + + assert result["status"] == 429 + assert "budget exhausted" in result["body"] + # The request was never dispatched, so it recorded nothing of its own. + assert await _log_rows(factory) == [] + + +async def test_an_unpriceable_completion_is_not_charged_the_whole_cap( + budget_app, monkeypatch +): + """A model absent from the catalog has no price, but one request must not + consume the key's entire remaining lifetime budget.""" + client, factory, key_id = budget_app + + # Countable tokens with no price any tier can reach: the model is not in the + # catalog and no LiteLLM cost was attached. + unpriced = { + "id": "chatcmpl-1", "object": "chat.completion", + "model": "my-private-model", "created": int(time.time()), + "choices": [{ + "index": 0, "finish_reason": "stop", + "message": {"role": "assistant", "content": "Hello there"}, + }], + "usage": {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, + } + + from app import router_cache + + fake = AsyncMock() + fake.acompletion = AsyncMock(side_effect=lambda **kwargs: unpriced) + async def _fake_get_router(_session): + return fake + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + + result = await _chat(client) + + assert result["status"] == 200 + spent = await _spent(factory, key_id) + cap = 1 * MICROCENTS_PER_CENT + # Fail closed, but bounded: the cap is not consumed by a single request. + assert 0 < spent < cap + # The row records the same amount the counter moved by. + rows = await _log_rows(factory) + assert rows[0].cost_microcents == spent + + +def _unmeasured_chunks(*, content: str, tool_calls: list | None = None) -> list[dict]: + """A stream with no `usage` frame — the delivery is the only pricing signal. + + This is the path `_unmeasured_charge` owns: the provider sent characters + through `delta.content` (and optionally `delta.tool_calls`) and then ended + with `finish_reason: "stop"` — no token counts ever arrived, so the charge + is computed from what the client was delivered, exactly as the blocking + path does. + """ + now = int(time.time()) + frames: list[dict] = [] + if content: + frames.append({ + "id": "chatcmpl-1", "object": "chat.completion.chunk", + "model": "gpt-4o-mini", "created": now, + "choices": [{"index": 0, "delta": {"content": content}, + "finish_reason": None}], + }) + if tool_calls: + frames.append({ + "id": "chatcmpl-1", "object": "chat.completion.chunk", + "model": "gpt-4o-mini", "created": now, + "choices": [{"index": 0, "delta": {"tool_calls": tool_calls}, + "finish_reason": None}], + }) + frames.append({ + "id": "chatcmpl-1", "object": "chat.completion.chunk", + "model": "gpt-4o-mini", "created": now, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }) + return frames + + +async def _stream_iter_from(chunks): + for c in chunks: + yield c + + +@pytest.mark.parametrize( + "chunks_fixture", + [ + pytest.param( + lambda: _unmeasured_chunks(content="Hello there"), + id="content-only", + ), + pytest.param( + lambda: _unmeasured_chunks( + content="Tool call: ", + tool_calls=[{"index": 0, "id": "call_1", "type": "function", + "function": {"name": "who", "arguments": ""}}], + ), + id="content-plus-tool-call", + ), + ], +) +async def test_unmeasured_streaming_delivery_is_priced_not_wholesaled( + budget_app, monkeypatch, chunks_fixture +): + """A streamed delivery that carries no usage frame must be priced per the + characters the client received — never the whole remaining budget. + + Regression for the P1: the old whole-remainder charge (`cap - spent`) + would consume the entire key on a single request. The P2: streamed + tool-call arguments were not counted as delivered content. Both are fixed + in `app/routes/chat.py` (`_unmeasured_charge` + the streaming delta loop), + and this locks the behavior at the integration boundary. + """ + client, factory, key_id = budget_app + + from app import router_cache + + fake = AsyncMock() + fake.acompletion = AsyncMock( + return_value=_stream_iter_from(chunks_fixture()) + ) + + async def _fake_get_router(_session): + return fake + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + + result = await _chat(client, stream=True) + + assert result["status"] == 200 + spent = await _spent(factory, key_id) + cap = 1 * MICROCENTS_PER_CENT + # Bounded: the charge is proportional to delivery size, not cap size. + assert 0 < spent < cap + rows = await _log_rows(factory) + assert rows[0].cost_microcents == spent + + +async def test_measured_but_unpriceable_stream_is_floored_not_zero( + budget_app, monkeypatch +): + """A stream with countable tokens but no priceable model charges the floor. + + Regression for the P1: the streaming settlement returned the row's 0 cost + untouched when usage was countable, so a custom/unlisted model streamed at + zero recorded cost forever. The blocking twin floors such completions at 1 + microcent; the stream must do the same. + """ + client, factory, key_id = budget_app + + now = int(time.time()) + chunks = [ + { + "id": "chatcmpl-1", "object": "chat.completion.chunk", + "model": "my-private-model", "created": now, + "choices": [{"index": 0, "delta": {"content": "Hello there"}, + "finish_reason": None}], + }, + { + "id": "chatcmpl-1", "object": "chat.completion.chunk", + "model": "my-private-model", "created": now, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 10, "completion_tokens": 20, + "total_tokens": 30}, + }, + ] + + from app import router_cache + + fake = AsyncMock() + fake.acompletion = AsyncMock(return_value=_stream_iter_from(chunks)) + + async def _fake_get_router(_session): + return fake + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + + # Request the unlisted model directly: both the requested and the served + # id must miss the catalog, or Tier 2 prices the request normally and the + # floor is never exercised. + r = await client.post( + "/v1/chat/completions", + json={ + "model": "my-private-model", + "messages": [{"role": "user", "content": "hi"}], + "stream": True, + }, + ) + assert r.status_code == 200 + spent = await _spent(factory, key_id) + cap = 1 * MICROCENTS_PER_CENT + # Fail closed, but bounded: the cap advances without being consumed whole. + assert 0 < spent < cap + rows = await _log_rows(factory) + assert rows[0].cost_microcents == spent + + +def _blocking_response_with(*, model: str, content: str | None, usage: dict) -> dict: + now = int(time.time()) + return { + "id": "chatcmpl-1", "object": "chat.completion", + "model": model, "created": now, + "choices": [{ + "index": 0, "finish_reason": "stop", + "message": {"role": "assistant", "content": content}, + }], + "usage": usage, + } + + +async def _chat_with_model(client, model: str, *, stream: bool = False) -> dict: + r = await client.post( + "/v1/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": "hi"}], + "stream": stream, + }, + ) + return {"status": r.status_code, "body": r.text} + + +async def test_zero_usage_with_delivered_content_is_estimated(budget_app, monkeypatch): + """Zero usage frames with real content fall back to the estimator. + + Settlement matrix cell: a provider that reports `0` tokens for a delivered + completion must not settle at zero — the charge comes from the delivery. + """ + client, factory, key_id = budget_app + + from app import router_cache + + fake = AsyncMock() + fake.acompletion = AsyncMock( + return_value=_blocking_response_with( + model="gpt-4o-mini", content="Hello there", + usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + ) + ) + + async def _fake_get_router(_session): + return fake + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + + result = await _chat(client) + + assert result["status"] == 200 + spent = await _spent(factory, key_id) + cap = 1 * MICROCENTS_PER_CENT + assert 0 < spent < cap + rows = await _log_rows(factory) + assert rows[0].cost_microcents == spent + + +async def test_empty_completion_with_zero_usage_charges_nothing(budget_app, monkeypatch): + """Zero usage with a genuinely empty completion settles at zero. + + Settlement matrix cell: nothing delivered means nothing to charge. The + estimator must not invent a completion from an empty response. + """ + client, factory, key_id = budget_app + + from app import router_cache + + fake = AsyncMock() + fake.acompletion = AsyncMock( + return_value=_blocking_response_with( + model="gpt-4o-mini", content="", + usage={"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}, + ) + ) + + async def _fake_get_router(_session): + return fake + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + + result = await _chat(client) + + assert result["status"] == 200 + assert await _spent(factory, key_id) == 0 + + +async def test_upstream_failure_before_delivery_charges_nothing(budget_app, monkeypatch): + """A fault that delivered nothing charges nothing. + + Settlement matrix cell: failures before any content are not billable, for + either path. The prompt may have gone upstream, but with no delivery and + no measurement there is no honest charge. + """ + client, factory, key_id = budget_app + + from app import router_cache + + fake = AsyncMock() + fake.acompletion = AsyncMock(side_effect=RuntimeError("upstream down")) + + async def _fake_get_router(_session): + return fake + + monkeypatch.setattr(router_cache, "get_router", _fake_get_router) + + for stream in (False, True): + result = await _chat(client, stream=stream) + assert result["status"] != 200 + assert await _spent(factory, key_id) == 0 + + +async def test_zero_cap_rejects_everything(budget_app): + """A zero budget limit rejects all requests, including free models. + + Product contract: `budget_limit_cents=0` means no lifetime spend is + authorized, so dispatch is refused up front. A negative limit behaves the + same (it authorizes even less); an absent limit means uncapped. + """ + client, factory, key_id = budget_app + + async with factory() as s: + row = await s.get(ApiKey, key_id) + row.budget_limit_cents = 0 + await s.commit() + + assert (await _chat(client))["status"] == 429 + assert (await _chat(client, stream=True))["status"] == 429 + + async with factory() as s: + row = await s.get(ApiKey, key_id) + row.budget_limit_cents = -5 + await s.commit() + + assert (await _chat(client))["status"] == 429 + + async with factory() as s: + row = await s.get(ApiKey, key_id) + row.budget_limit_cents = None + await s.commit() + + assert (await _chat(client))["status"] == 200 + + +async def test_unmeasured_charge_covers_every_ending(): + """Unit matrix for `_unmeasured_charge`: every ending has a defined charge. + + - Client disconnect with nothing delivered: the prompt went upstream, so + the prompt estimate is charged. + - Upstream fault with nothing delivered: zero. + - Fault after a delivery: the delivery price, same as a clean completion. + """ + from app.routes.chat import ( + _STREAM_CLIENT_DISCONNECT, + _STREAM_COMPLETED, + _STREAM_UPSTREAM_ERROR, + _unmeasured_charge, + ) + + charge, estimate = _unmeasured_charge( + delivered=False, ending=_STREAM_CLIENT_DISCONNECT, + prompt_chars=400, completion_chars=0, + model_id="gpt-4o-mini", fallback_model="gpt-4o-mini", + ) + assert charge > 0 + assert estimate == {"prompt_tokens": 100, "completion_tokens": 0} + + charge, estimate = _unmeasured_charge( + delivered=False, ending=_STREAM_UPSTREAM_ERROR, + prompt_chars=400, completion_chars=0, + model_id="gpt-4o-mini", fallback_model="gpt-4o-mini", + ) + assert (charge, estimate) == (0, None) + + faulted, _ = _unmeasured_charge( + delivered=True, ending=_STREAM_UPSTREAM_ERROR, + prompt_chars=400, completion_chars=40, + model_id="gpt-4o-mini", fallback_model="gpt-4o-mini", + ) + clean, _ = _unmeasured_charge( + delivered=True, ending=_STREAM_COMPLETED, + prompt_chars=400, completion_chars=40, + model_id="gpt-4o-mini", fallback_model="gpt-4o-mini", + ) + assert faulted == clean > 0 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..b669c5ce --- /dev/null +++ b/tests/unit/test_budget_park.py @@ -0,0 +1,439 @@ +"""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 select +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} + + +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_second_cancel_during_probe_with_a_durable_row_leaves_no_hold( + factory, monkeypatch +): + """A second cancel during the probe must still honor a durable row. + + Twin of the absent-row case: the probe runs to its decision either way, + so a landed commit never leaves a hold behind even under repeated + cancellation. + """ + 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=600)) + await s.commit() + + 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 + assert spend._unsettled.get("k") in (None, {}) + await settle_parked_spend("k", cap_microcents=10_000) + assert await _spent(factory, "k") == 600 diff --git a/tests/unit/test_budget_pricing.py b/tests/unit/test_budget_pricing.py new file mode 100644 index 00000000..502aaa54 --- /dev/null +++ b/tests/unit/test_budget_pricing.py @@ -0,0 +1,310 @@ +"""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 + + +@pytest.mark.parametrize( + "content", + [ + "hello", + ["hello", "world"], + [{"type": "text", "text": "ab"}, "cd"], + [{"type": "tool_use", "name": "get_weather", "input": {"city": "SF"}}], + [{"type": "tool_result", "tool_use_id": "t1", "content": "tool output here"}], + [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}], + ], +) +def test_prompt_and_delivery_price_the_same_characters(content): + """Blocking and streaming price content through the same part reader. + + The prompt side (`text_chars`) and the delivery side + (`blocking_delivery_chars`) must agree for every part shape, so one path + cannot undercharge what the other bills. + """ + _, delivered = blocking_delivery_chars({"choices": [{"message": {"content": content}}]}) + assert text_chars(content) == delivered + + +def test_whitespace_divergence_between_prompt_and_delivery_is_pinned(): + """Whitespace is the one documented divergence: the prompt counts every + character (streamed deltas are priced through `text_chars`), while the + delivery path skips whitespace-only values when deciding anything arrived. + """ + assert text_chars(" ") == 3 + assert blocking_delivery_chars({"choices": [{"message": {"content": " "}}]}) == (False, 0) 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 == []