From 2aa2788a24f106cfb29e44856b608ffc44be0608 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 10:49:59 +0000 Subject: [PATCH 01/32] feat(budget): add pure delivery pricing helpers Characters-to-tokens estimation and usage detection for requests whose upstream never reported a measurement. No database or route state, so the pricing rules can be tested in isolation before settlement uses them. --- app/budget_pricing.py | 127 ++++++++++++++++++++++++++++++ tests/unit/test_budget_pricing.py | 99 +++++++++++++++++++++++ 2 files changed, 226 insertions(+) create mode 100644 app/budget_pricing.py create mode 100644 tests/unit/test_budget_pricing.py diff --git a/app/budget_pricing.py b/app/budget_pricing.py new file mode 100644 index 00000000..e8a47f7b --- /dev/null +++ b/app/budget_pricing.py @@ -0,0 +1,127 @@ +"""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 + +# Crude character-to-token divisor, used only to price a delivery the provider +# did not measure. +CHARS_PER_TOKEN = 4 + +_COUNTABLE_USAGE_KEYS = ( + "prompt_tokens", + "completion_tokens", + "input_tokens", + "output_tokens", +) + + +def text_chars(content) -> int: + """Character count of message content, for a str or a list of text parts.""" + if isinstance(content, str): + return len(content) + if isinstance(content, list): + return sum( + len(part["text"]) + for part in content + if isinstance(part, dict) and isinstance(part.get("text"), str) + ) + return 0 + + +def countable_usage(usage) -> bool: + """Whether a usage dict carries at least one token total. + + A truthy usage dict is not proof of measurement: some upstreams report + `{"total_tokens": 123}` and nothing else. Settling on that would record a + zero cost for a delivered completion, so anything without a countable key + is treated as no measurement at all. + """ + if not isinstance(usage, dict) or not usage: + return False + return any(usage.get(key) is not None 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: + if isinstance(part, str): + candidates.append(part) + elif isinstance(part, dict) and isinstance(part.get("text"), str): + candidates.append(part["text"]) + 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/tests/unit/test_budget_pricing.py b/tests/unit/test_budget_pricing.py new file mode 100644 index 00000000..8272f233 --- /dev/null +++ b/tests/unit/test_budget_pricing.py @@ -0,0 +1,99 @@ +"""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, + text_chars, + tool_call_text, +) + + +@pytest.mark.parametrize( + ("usage", "expected"), + [ + ({"prompt_tokens": 0}, True), + ({"completion_tokens": 0}, True), + ({"input_tokens": 5}, True), + ({"output_tokens": 5}, True), + ({"total_tokens": 123}, False), + ({"prompt_tokens": None}, False), + ({}, False), + (None, False), + ("not a dict", False), + ], +) +def test_countable_usage_requires_a_countable_key(usage, expected): + assert countable_usage(usage) is expected + + +@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 + assert text_chars([{"type": "text", "text": "ab"}, {"type": "image_url"}]) == 2 + assert text_chars(None) == 0 + + +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) From f96df199ea484bef480fa46d6e7ee91683f04df3 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 12:32:17 +0000 Subject: [PATCH 02/32] fix(budget): only a positive token total counts as a measurement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `countable_usage` accepted any countable key that was merely present, so a usage frame of `{"prompt_tokens": 0, "completion_tokens": 0}` — what LiteLLM emits when the upstream reported nothing — counted as measured. Settling on it records a zero cost for a delivered completion, which is exactly the bypass the guard exists to prevent. Only a countable key holding a positive number is now a measurement. A zero on one side still counts when the other side is positive, so a prompt-only request is unaffected. `bool` is rejected explicitly because `True` is an `int`. A genuinely empty completion still prices at zero through `blocking_delivery_chars` and `estimate_usage`, so this cannot overcharge. --- app/budget_pricing.py | 22 ++++++++++++++++------ tests/unit/test_budget_pricing.py | 27 ++++++++++++++++++++++++--- 2 files changed, 40 insertions(+), 9 deletions(-) diff --git a/app/budget_pricing.py b/app/budget_pricing.py index e8a47f7b..a68ff62f 100644 --- a/app/budget_pricing.py +++ b/app/budget_pricing.py @@ -35,16 +35,26 @@ def text_chars(content) -> int: def countable_usage(usage) -> bool: - """Whether a usage dict carries at least one token total. + """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. Settling on that would record a - zero cost for a delivered completion, so anything without a countable key - is treated as no measurement at all. + 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 — anything else is treated as no measurement at all. + + `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 - return any(usage.get(key) is not None for key in _COUNTABLE_USAGE_KEYS) + return any( + isinstance(value := usage.get(key), (int, float)) + and not isinstance(value, bool) + and value > 0 + for key in _COUNTABLE_USAGE_KEYS + ) def chars_to_tokens(chars: int) -> int: diff --git a/tests/unit/test_budget_pricing.py b/tests/unit/test_budget_pricing.py index 8272f233..72d20b89 100644 --- a/tests/unit/test_budget_pricing.py +++ b/tests/unit/test_budget_pricing.py @@ -15,21 +15,42 @@ @pytest.mark.parametrize( ("usage", "expected"), [ - ({"prompt_tokens": 0}, True), - ({"completion_tokens": 0}, True), + ({"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_countable_key(usage, expected): +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)], From fb929278a28b1b1e3059e370673a8605cce19844 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 04:54:38 +0000 Subject: [PATCH 03/32] fix(budget/helpers): count bare-string content parts in text_chars MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit text_chars only counted dict parts ({"text": "..."}), so an Anthropic-style prompt of bare strings (`["hello"]`) priced at 0 characters. blocking_delivery_chars counts the same shape on a completion, so a budgeted key's prompt was under-priced whenever the wire used plain-string parts — the precondition is /v1/messages forwarding Anthropic-style content, reachable by real clients. Count plain string parts in text_chars the way blocking_delivery_chars does, keeping the dict branch, so a prompt and a completion of the same content price identical characters. Verified: text_chars(["hello","world"]) == 10; a mixed list prices only its textual parts. --- app/budget_pricing.py | 20 ++++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/app/budget_pricing.py b/app/budget_pricing.py index a68ff62f..98482b53 100644 --- a/app/budget_pricing.py +++ b/app/budget_pricing.py @@ -22,15 +22,23 @@ def text_chars(content) -> int: - """Character count of message content, for a str or a list of text parts.""" + """Character count of message content, for a str or a list of text parts. + + Anthropic-style content may be either a plain string or a list of parts, + where each part may be a dict (`{"text": "..."}`) or a plain `str`. Both + shapes are counted so a prompt prices the same characters the completion + delivery path (`blocking_delivery_chars`) would count. + """ if isinstance(content, str): return len(content) if isinstance(content, list): - return sum( - len(part["text"]) - for part in content - if isinstance(part, dict) and isinstance(part.get("text"), str) - ) + total = 0 + for part in content: + if isinstance(part, str): + total += len(part) + elif isinstance(part, dict) and isinstance(part.get("text"), str): + total += len(part["text"]) + return total return 0 From 18c737865a7bd159b46551eed27c3f856b172e54 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 05:08:13 +0000 Subject: [PATCH 04/32] test(budget/helpers): cover bare-string content parts in text_chars Regression for the bare-string pricing fix: a list of plain strings (`["hello", "world"]`) and a mixed list with bare-string parts now have expected character counts. Mirrors blocking_delivery_chars, which already accepted str parts, so prompt and completion price the same characters. --- tests/unit/test_budget_pricing.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/tests/unit/test_budget_pricing.py b/tests/unit/test_budget_pricing.py index 72d20b89..c346843e 100644 --- a/tests/unit/test_budget_pricing.py +++ b/tests/unit/test_budget_pricing.py @@ -73,6 +73,17 @@ def test_text_chars_across_str_and_content_parts(): 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_tool_call_text_counts_names_and_string_arguments(): calls = [ {"function": {"name": "search", "arguments": '{"q":"x"}'}}, From 965cb88b8cccc3c50877ac72f3b3236dc3722859 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 07:29:23 +0000 Subject: [PATCH 05/32] fix(budget/helpers): price prompt tool calls, require completion count on delivery, count tool_use parts - message_chars: prompt side counts tool_calls as well as content, matching the delivery path, so tool-call-only assistant turns no longer price at zero - countable_usage(delivered=True): a delivered completion needs a positive completion-side count; a prompt-only frame no longer settles it at zero - _part_texts: Anthropic tool_use parts count name and string input, shared by text_chars and blocking_delivery_chars so the two cannot drift --- app/budget_pricing.py | 85 ++++++++++++++++++++----------- tests/unit/test_budget_pricing.py | 56 ++++++++++++++++++++ 2 files changed, 111 insertions(+), 30 deletions(-) diff --git a/app/budget_pricing.py b/app/budget_pricing.py index 98482b53..253e84a9 100644 --- a/app/budget_pricing.py +++ b/app/budget_pricing.py @@ -13,36 +13,62 @@ # did not measure. CHARS_PER_TOKEN = 4 -_COUNTABLE_USAGE_KEYS = ( - "prompt_tokens", - "completion_tokens", - "input_tokens", - "output_tokens", -) +_PROMPT_USAGE_KEYS = ("prompt_tokens", "input_tokens") +_COMPLETION_USAGE_KEYS = ("completion_tokens", "output_tokens") +_COUNTABLE_USAGE_KEYS = _PROMPT_USAGE_KEYS + _COMPLETION_USAGE_KEYS + + +def _part_texts(part) -> list[str]: + """The character-bearing strings one content part carries. + + A part is a bare string, a `{"text": ...}` dict, or an Anthropic + `{"type": "tool_use", "name": ..., "input": ...}` block. A tool_use block + is delivered content the upstream bills for: its name and any string-valued + input count, mirroring `tool_call_text`. Object-valued input has no honest + character count and contributes nothing. + """ + if isinstance(part, str): + return [part] + if not isinstance(part, dict): + return [] + texts: list[str] = [] + if isinstance(part.get("text"), str): + texts.append(part["text"]) + if part.get("type") == "tool_use": + for key in ("name", "input"): + value = part.get(key) + if isinstance(value, str): + texts.append(value) + return texts def text_chars(content) -> int: - """Character count of message content, for a str or a list of text parts. + """Character count of message content, for a str or a list of content parts. - Anthropic-style content may be either a plain string or a list of parts, - where each part may be a dict (`{"text": "..."}`) or a plain `str`. Both - shapes are counted so a prompt prices the same characters the completion - delivery path (`blocking_delivery_chars`) would count. + Counts the same strings `blocking_delivery_chars` counts for a completion, + so a prompt prices the same characters the delivery path would. """ if isinstance(content, str): return len(content) if isinstance(content, list): - total = 0 - for part in content: - if isinstance(part, str): - total += len(part) - elif isinstance(part, dict) and isinstance(part.get("text"), str): - total += len(part["text"]) - return total + return sum(len(text) for part in content for text in _part_texts(part)) return 0 -def countable_usage(usage) -> bool: +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 @@ -50,19 +76,21 @@ def countable_usage(usage) -> bool: `{"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 — anything else is treated as no measurement at all. + 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 - return any( - isinstance(value := usage.get(key), (int, float)) - and not isinstance(value, bool) - and value > 0 - for key in _COUNTABLE_USAGE_KEYS - ) + 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: @@ -132,10 +160,7 @@ def blocking_delivery_chars(response: dict) -> tuple[bool, int]: candidates.append(content) elif isinstance(content, list): for part in content: - if isinstance(part, str): - candidates.append(part) - elif isinstance(part, dict) and isinstance(part.get("text"), str): - candidates.append(part["text"]) + candidates.extend(_part_texts(part)) candidates.append(tool_call_text(message.get("tool_calls"))) for value in candidates: if not value.strip(): diff --git a/tests/unit/test_budget_pricing.py b/tests/unit/test_budget_pricing.py index c346843e..dd59dfa1 100644 --- a/tests/unit/test_budget_pricing.py +++ b/tests/unit/test_budget_pricing.py @@ -7,6 +7,7 @@ chars_to_tokens, countable_usage, estimate_usage, + message_chars, text_chars, tool_call_text, ) @@ -129,3 +130,58 @@ def test_blocking_delivery_reads_legacy_text_and_part_lists(): 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_only_its_name(): + part = {"type": "tool_use", "name": "search", "input": {"q": "x"}} + assert text_chars([part]) == len("search") + + +def test_tool_use_only_completion_is_a_delivery_not_zero(): + response = {"choices": [{"message": {"content": [{"type": "tool_use", "name": "f"}]}}]} + assert blocking_delivery_chars(response) == (True, 1) From b3b0933269134991d5d141dd81a0bf0145e32e88 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 07:39:00 +0000 Subject: [PATCH 06/32] fix(budget/helpers): document whitespace-inclusive text_chars, pin it with a test text_chars intentionally counts whitespace; blocking_delivery_chars skips it only for the delivered-or-not decision. Streamed deltas are priced through text_chars, so dropping whitespace there would undercount real output. --- app/budget_pricing.py | 8 ++++++-- tests/unit/test_budget_pricing.py | 12 ++++++++++++ 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/app/budget_pricing.py b/app/budget_pricing.py index 253e84a9..5ac0fa81 100644 --- a/app/budget_pricing.py +++ b/app/budget_pricing.py @@ -45,8 +45,12 @@ def _part_texts(part) -> list[str]: def text_chars(content) -> int: """Character count of message content, for a str or a list of content parts. - Counts the same strings `blocking_delivery_chars` counts for a completion, - so a prompt prices the same characters the delivery path would. + 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) diff --git a/tests/unit/test_budget_pricing.py b/tests/unit/test_budget_pricing.py index dd59dfa1..49fa2c89 100644 --- a/tests/unit/test_budget_pricing.py +++ b/tests/unit/test_budget_pricing.py @@ -85,6 +85,18 @@ def test_text_chars_counts_bare_string_parts(): 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"}'}}, From 8a6287b9b19cec09bc4a3677ecd90d0325360226 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 07:50:32 +0000 Subject: [PATCH 07/32] fix(budget/helpers): count object tool input and refusals, floor delivered opaque parts - tool_use: count json.dumps(input or {}), the same serialization the Anthropic and Gemini translators bill, instead of dropping object input - delivery: count message.refusal, which reaches the client as output - delivery: a non-text delivered part (image, audio) counts as one character so it floors to one token and never settles at zero; the prompt-side text_chars still counts it exactly --- app/budget_pricing.py | 56 ++++++++++++++++++++++++------- tests/unit/test_budget_pricing.py | 48 +++++++++++++++++++++++--- 2 files changed, 87 insertions(+), 17 deletions(-) diff --git a/app/budget_pricing.py b/app/budget_pricing.py index 5ac0fa81..e80057a6 100644 --- a/app/budget_pricing.py +++ b/app/budget_pricing.py @@ -9,6 +9,8 @@ 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 @@ -17,28 +19,51 @@ _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, or an Anthropic - `{"type": "tool_use", "name": ..., "input": ...}` block. A tool_use block - is delivered content the upstream bills for: its name and any string-valued - input count, mirroring `tool_call_text`. Object-valued input has no honest - character count and contributes nothing. + A part is a bare string, a `{"text": ...}` dict, an Anthropic + `{"type": "tool_use", "name": ..., "input": ...}` block, or some other + non-text part (an image, an audio clip, a b64 payload). + + A tool_use block is billed as its name plus its input serialized the same + way the translators serialize it (`json.dumps(input or {})`), so an object + input is counted in full, not dropped. A non-text part has no honest + character count, but it was delivered, so it counts as one character: that + keeps it from settling at zero while `chars_to_tokens` floors it to a token. """ if isinstance(part, str): return [part] if not isinstance(part, dict): return [] - texts: list[str] = [] - if isinstance(part.get("text"), str): - texts.append(part["text"]) if part.get("type") == "tool_use": - for key in ("name", "input"): - value = part.get(key) - if isinstance(value, str): - texts.append(value) + 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"]] + return [] + + +def _delivered_part_texts(part) -> list[str]: + """`_part_texts`, plus the one-character floor for a delivered opaque part. + + Only the delivery path uses this. A non-text part in a completion was billed + as output, so it must not settle at zero. A prompt image is counted exactly + by `text_chars`, which gives it no characters. + """ + texts = _part_texts(part) + if not texts and isinstance(part, dict) and part.get("type") not in (None, "text"): + return [_OPAQUE_PART_MARKER] return texts @@ -164,7 +189,12 @@ def blocking_delivery_chars(response: dict) -> tuple[bool, int]: candidates.append(content) elif isinstance(content, list): for part in content: - candidates.extend(_part_texts(part)) + candidates.extend(_delivered_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(): diff --git a/tests/unit/test_budget_pricing.py b/tests/unit/test_budget_pricing.py index 49fa2c89..d8b3a871 100644 --- a/tests/unit/test_budget_pricing.py +++ b/tests/unit/test_budget_pricing.py @@ -189,11 +189,51 @@ def test_anthropic_tool_use_part_is_counted_as_delivered_content(): assert blocking_delivery_chars(response) == (True, len("search") + len('{"q":"x"}')) -def test_tool_use_part_with_object_input_counts_only_its_name(): - part = {"type": "tool_use", "name": "search", "input": {"q": "x"}} - assert text_chars([part]) == len("search") +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"}]}}]} - assert blocking_delivery_chars(response) == (True, 1) + # 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("{}") From ed0de179cf465c247ef091fc6e2527f4158f2789 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 07:53:15 +0000 Subject: [PATCH 08/32] fix(budget/helpers): count non-text prompt parts as one character A non-text prompt part (image, tool_result) was billed but counted zero, so an image-only prompt floored to zero tokens. Such a part now counts one character, so it floors to one token. Both prompt and delivery paths share _part_texts. --- app/budget_pricing.py | 22 +++++++--------------- tests/unit/test_budget_pricing.py | 9 ++++++++- 2 files changed, 15 insertions(+), 16 deletions(-) diff --git a/app/budget_pricing.py b/app/budget_pricing.py index e80057a6..4006d157 100644 --- a/app/budget_pricing.py +++ b/app/budget_pricing.py @@ -51,20 +51,12 @@ def _part_texts(part) -> list[str]: return texts if isinstance(part.get("text"), str): return [part["text"]] - return [] - - -def _delivered_part_texts(part) -> list[str]: - """`_part_texts`, plus the one-character floor for a delivered opaque part. - - Only the delivery path uses this. A non-text part in a completion was billed - as output, so it must not settle at zero. A prompt image is counted exactly - by `text_chars`, which gives it no characters. - """ - texts = _part_texts(part) - if not texts and isinstance(part, dict) and part.get("type") not in (None, "text"): - return [_OPAQUE_PART_MARKER] - return texts + if part.get("type") in (None, "text"): + return [] + # A non-text part (image, audio, tool_result, ...) 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 text_chars(content) -> int: @@ -189,7 +181,7 @@ def blocking_delivery_chars(response: dict) -> tuple[bool, int]: candidates.append(content) elif isinstance(content, list): for part in content: - candidates.extend(_delivered_part_texts(part)) + 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") diff --git a/tests/unit/test_budget_pricing.py b/tests/unit/test_budget_pricing.py index d8b3a871..72204212 100644 --- a/tests/unit/test_budget_pricing.py +++ b/tests/unit/test_budget_pricing.py @@ -70,7 +70,8 @@ def test_estimate_usage_prices_both_sides_from_characters(): def test_text_chars_across_str_and_content_parts(): assert text_chars("hello") == 5 - assert text_chars([{"type": "text", "text": "ab"}, {"type": "image_url"}]) == 2 + # 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 @@ -237,3 +238,9 @@ def test_non_text_delivered_part_counts_as_one_character_not_zero(): 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 From 6af963fe3534d49e8e5e7cc29e54a52a39cb0617 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 08:19:57 +0000 Subject: [PATCH 09/32] fix(budget/helpers): count tool_result text instead of one marker character A tool_result's billable text is its content (string or text parts). It was priced as a single opaque character, undercharging multi-turn tool prompts by the length of every tool output. Its text is now counted; nested non-text parts still floor to one character. --- app/budget_pricing.py | 23 ++++++++++++++++++--- tests/unit/test_budget_pricing.py | 33 +++++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 3 deletions(-) diff --git a/app/budget_pricing.py b/app/budget_pricing.py index 4006d157..d026fe0f 100644 --- a/app/budget_pricing.py +++ b/app/budget_pricing.py @@ -51,14 +51,31 @@ def _part_texts(part) -> list[str]: 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, tool_result, ...) 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. + # 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. diff --git a/tests/unit/test_budget_pricing.py b/tests/unit/test_budget_pricing.py index 72204212..fa8ed684 100644 --- a/tests/unit/test_budget_pricing.py +++ b/tests/unit/test_budget_pricing.py @@ -244,3 +244,36 @@ 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 From e59724e5082457d0fca737a62ae4cc399c9acc67 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 10:53:53 +0000 Subject: [PATCH 10/32] feat(budget): add the spend ledger with a durable park for lost settlements A settlement that cannot be recorded leaves a delivered cost with no trace: the request-log row and the charge share one transaction, so both roll back together. The park holds that cost against the key's cap until a pre-check folds it in, oldest debt first, and clears a row whose charge already landed. Startup migration creates the table for deployments whose schema predates it. When even the park write cannot land, the amount is held in memory for the next pre-check to re-file. That hold keeps one entry per key rather than one per settlement, so a database outage that outlasts many requests cannot grow it without bound. --- packages/auth/spend.py | 391 +++++++++++++++++++++++++++- packages/db/migrate.py | 56 +++- packages/db/models/__init__.py | 2 + packages/db/models/budget_park.py | 45 ++++ tests/unit/test_budget_migration.py | 9 +- tests/unit/test_budget_park.py | 258 ++++++++++++++++++ tests/unit/test_budget_spend.py | 4 +- 7 files changed, 738 insertions(+), 27 deletions(-) create mode 100644 packages/db/models/budget_park.py create mode 100644 tests/unit/test_budget_park.py diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 27649968..7414959f 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,346 @@ from __future__ import annotations -from sqlalchemy import select, update +import asyncio + +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_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 — it maps `api_key_id` to `(trace_id, microcents)`, and a later +# pre-check re-files it once a write goes through again. +# +# One entry per key, not per settlement: a database outage that outlasts many +# requests would otherwise grow this without bound, since an entry only leaves +# when a fold commits and no fold can commit while the database is down. +_unsettled: dict[str, tuple[str, int]] = {} + + +def _hold(api_key_id: str, trace_id: str, microcents: int) -> None: + """Add `microcents` to the key's in-memory hold. + + The first `trace_id` seen becomes the hold's identity and is kept as further + settlements merge in, so re-filing stays idempotent on it: a retry parks the + same obligation once rather than once per merged settlement. + """ + held = _unsettled.get(api_key_id) + if held is None: + _unsettled[api_key_id] = (trace_id, microcents) + else: + _unsettled[api_key_id] = (held[0], held[1] + microcents) + + +class _FoldConflict(Exception): + """A concurrent worker folded the same park first; the loser retries later.""" + + +async def _park_is_durable(trace_id: str) -> bool: + """Whether a park row for this `trace_id` is committed. + + A commit that applied but whose ack never came back raises exactly like a + failure, and holding a memory copy beside the durable row puts one + obligation in both ledgers — the double-bill the `trace_id` key exists to + absorb. It reads on a fresh session because the failed one is closed by the + time this runs, and answers False when it cannot read at all: with the + database truly unreachable the memory hold is all that keeps the cap honest. + """ + 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: + return ( + await s.scalar( + select(BudgetPark.trace_id).where(BudgetPark.trace_id == trace_id) + ) + ) is not None + 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) + + +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: + # `_insert_park` only raises the cancellation unwinding the caller — + # and the amount still has to be held before it propagates. + _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 held[1] if held is not None else 0 + try: + async with factory() as s: + rows = ( + await s.execute( + select(BudgetPark.trace_id, BudgetPark.microcents).where( + BudgetPark.api_key_id == key + ) + ) + ).all() + except Exception: + return None + # The durable rows are read before the hold is counted: a re-file that + # lands while this read is in flight must not be counted a second time, so + # a hold whose trace_id already has a row is not counted again. + durable_traces = {trace_id for trace_id, _microcents in rows} + stored = sum(int(microcents) for _trace_id, microcents in rows) + held_amount = ( + held[1] if held is not None and held[0] not in durable_traces else 0 + ) + 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. + """ + from packages.db import session as session_mod + + key = str(api_key_id) + factory = session_mod._session_factory + if factory is None: + return 0 + for held_key, (trace_id, amount) in list(_unsettled.items()): + if held_key != key: + continue + # A cancellation here propagates with the hold still in place; a later + # pre-check re-files it, and the stable `trace_id` keeps the retry from + # parking the same obligation twice. + if await _insert_park( + trace_id=trace_id, api_key_id=key, microcents=amount + ): + _unsettled.pop(held_key, None) + move = 0 + settling: list[str] = [] + 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() + already_charged = [trace_id for trace_id, _amount in rows if trace_id in logged] + if already_charged: + cleared = await s.execute( + delete(BudgetPark).where( + BudgetPark.trace_id.in_(already_charged) + ) + ) + if cleared.rowcount != len(already_charged): + raise _FoldConflict + 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) + continue + if room > 0: + trim = (trace_id, microcents, microcents - room) + move += room + break + if move <= 0: + return 0 + 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: + cleared = await s.execute( + delete(BudgetPark).where(BudgetPark.trace_id.in_(settling)) + ) + if cleared.rowcount != len(settling): + raise _FoldConflict + if trim is not None: + trace_id, whole, remainder = trim + trimmed = await s.execute( + update(BudgetPark) + .where( + BudgetPark.trace_id == trace_id, + BudgetPark.microcents == whole, + ) + .values(microcents=remainder) + ) + if trimmed.rowcount != 1: + raise _FoldConflict + except _FoldConflict: + return 0 + # Past the commit, so everything in `settling` is billed and gone from the + # table. A hold covering one of those rows is now worse than stale: the next + # pre-check re-files it as a brand-new park — the `trace_id` no longer + # collides, this transaction deleted the row — and the same delivery is + # charged twice against a cap that has no idea it moved. Only `settling`: + # a trimmed row is still owed, and the hold covering it has to stay. + billed = {*settling, *already_charged} + held = _unsettled.get(key) + if held is not None and held[0] in billed: + _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 +380,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 +433,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 +459,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..f13ce528 100644 --- a/packages/db/migrate.py +++ b/packages/db/migrate.py @@ -14,26 +14,33 @@ 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_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 +50,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 +74,35 @@ 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: + # `create_all` covers fresh databases; this covers upgrades whose + # schema predates the table. `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), + ) + # 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 +142,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..f8fc00a3 100644 --- a/packages/db/models/__init__.py +++ b/packages/db/models/__init__.py @@ -2,6 +2,7 @@ from packages.db.models.api_key import ApiKey from packages.db.models.base import Base, SoftDeleteMixin, TimestampMixin, UUIDMixin +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 +16,7 @@ "TimestampMixin", "UUIDMixin", "ApiKey", + "BudgetPark", "ProviderKey", "QualityScoreOverride", "QualityScoreSnapshot", diff --git a/packages/db/models/budget_park.py b/packages/db/models/budget_park.py new file mode 100644 index 00000000..0b9dafba --- /dev/null +++ b/packages/db/models/budget_park.py @@ -0,0 +1,45 @@ +"""Unsettled budget obligations — delivered spend the ledger never recorded. + +A settlement that gives up after every retry leaves a delivered cost with no +record anywhere: the log row and the charge are one transaction, so both roll +back and `spent_microcents` never moves. The obligation is parked here, one row +per settlement, so it outlives the process that lost it and is visible to every +worker behind the same database. + +Rows are keyed by the settlement's `trace_id`, which makes every park write +idempotent: a commit that applied but whose ack was lost retries into the same +primary key instead of recording the obligation twice, and a write that fails +outright is checked for having landed anyway before the process falls back to +holding it in memory. A fold that bills a row either deletes it or shrinks it to +what the cap could not absorb, in the same transaction that moves +`spent_microcents`, and it drops the writer's memory hold for any row it fully +billed. +""" + +from datetime import datetime, timezone + +from sqlalchemy import BigInteger, DateTime, String, func +from sqlalchemy.orm import Mapped, mapped_column + +from packages.db.models.base import Base, TimestampMixin, UUIDMixin + + +class BudgetPark(Base, UUIDMixin, TimestampMixin): + __tablename__ = "budget_parks" + + trace_id: Mapped[str] = mapped_column(String(36), nullable=False, unique=True) + api_key_id: Mapped[str] = mapped_column(String(36), nullable=False, index=True) + microcents: Mapped[int] = mapped_column(BigInteger, nullable=False) + # Overrides TimestampMixin's column, whose `server_default=func.now()` is + # CURRENT_TIMESTAMP: one second wide on SQLite, where a recovered outage + # re-files a whole batch of obligations in a single pre-check and every row + # in it ties. The fold bills oldest debt first, and two workers have to + # agree on which row the partial one is, so the stamp needs enough + # resolution to settle that on its own instead of falling through to the + # `trace_id` tiebreak — which is a uuid4, and so picks at random. + created_at: Mapped[datetime] = mapped_column( + DateTime(timezone=True), + default=lambda: datetime.now(timezone.utc), + server_default=func.now(), + nullable=False, + ) diff --git a/tests/unit/test_budget_migration.py b/tests/unit/test_budget_migration.py index e1ec8ffd..45f1b53e 100644 --- a/tests/unit/test_budget_migration.py +++ b/tests/unit/test_budget_migration.py @@ -239,10 +239,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 +267,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 diff --git a/tests/unit/test_budget_park.py b/tests/unit/test_budget_park.py new file mode 100644 index 00000000..f65f584a --- /dev/null +++ b/tests/unit/test_budget_park.py @@ -0,0 +1,258 @@ +"""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. +""" + +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 diff --git a/tests/unit/test_budget_spend.py b/tests/unit/test_budget_spend.py index ffb1a0af..0e703b48 100644 --- a/tests/unit/test_budget_spend.py +++ b/tests/unit/test_budget_spend.py @@ -91,8 +91,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 From 28364f884472a00d5d6e8f84694a59c5750e44e2 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 11:50:05 +0000 Subject: [PATCH 11/32] fix(budget): pop the park hold only when unchanged since the write A settlement can merge into the key's in-memory hold while the park write is in flight. Popping the whole entry afterwards dropped the merged amount without ever persisting it, reopening the cap by that amount. Re-file until the write lands on a stable hold. --- packages/auth/spend.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 7414959f..a92a3944 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -257,16 +257,22 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: factory = session_mod._session_factory if factory is None: return 0 - for held_key, (trace_id, amount) in list(_unsettled.items()): - if held_key != key: - continue + # A settlement can merge into this key's hold while the park write is in + # flight. Pop the hold only when it is unchanged since the write; otherwise + # the merged amount was never persisted, so re-file the grown hold (the same + # `trace_id` raises the row to it) until a write lands on a stable hold. + while (held := _unsettled.get(key)) is not None: + trace_id, amount = held # A cancellation here propagates with the hold still in place; a later # pre-check re-files it, and the stable `trace_id` keeps the retry from # parking the same obligation twice. - if await _insert_park( + if not await _insert_park( trace_id=trace_id, api_key_id=key, microcents=amount ): - _unsettled.pop(held_key, None) + break + if _unsettled.get(key) == held: + _unsettled.pop(key, None) + break move = 0 settling: list[str] = [] trim: tuple[str, int, int] | None = None From 19b5986f76c3ea3a30174a3b7a1d8d7be4a064e5 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 12:51:35 +0000 Subject: [PATCH 12/32] fix(budget): close four accounting races in the park/hold/fold path 1. Durability probe now checks the persisted amount, not just row existence. A row that exists but under-covers the expected obligation is answered as not durable, so the give-up parks the full amount instead of trusting a short row. 2. pending_parked_spend no longer zeroes a hold when its trace_id has a durable row. It now subtracts the durable amount from the hold and adds only the uncovered remainder, so a merged hold is never double-counted and its shortfall is never dropped. 3. Both DELETE paths in settle_parked_spend carry compare-and-swap guards on the amount read inside the transaction. A concurrent increase to a park row now conflicts the fold instead of silently deleting an amount that wasn't billed. 4. Hold reconciliation after the commit keeps only the obligation not covered by the committed fold. A new trace_id prevents the remainder from colliding with a park row that was just deleted. The move==0 path (already-charged cleanup only) still reconciles holds. These four changes close the races where a growing hold could lose its growth, where a short durable row could be trusted, and where a folded row could be deleted past what the fold actually billed. --- packages/auth/spend.py | 126 +++++++++++++++++++++++++---------------- 1 file changed, 78 insertions(+), 48 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index a92a3944..7490e185 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -32,6 +32,7 @@ from __future__ import annotations import asyncio +import uuid from sqlalchemy import delete, select, update from sqlalchemy.exc import IntegrityError @@ -83,16 +84,10 @@ class _FoldConflict(Exception): """A concurrent worker folded the same park first; the loser retries later.""" -async def _park_is_durable(trace_id: str) -> bool: - """Whether a park row for this `trace_id` is committed. - - A commit that applied but whose ack never came back raises exactly like a - failure, and holding a memory copy beside the durable row puts one - obligation in both ledgers — the double-bill the `trace_id` key exists to - absorb. It reads on a fresh session because the failed one is closed by the - time this runs, and answers False when it cannot read at all: with the - database truly unreachable the memory hold is all that keeps the cap honest. - """ +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 @@ -100,11 +95,16 @@ async def _park_is_durable(trace_id: str) -> bool: return False try: async with factory() as s: - return ( - await s.scalar( - select(BudgetPark.trace_id).where(BudgetPark.trace_id == trace_id) + amount = await s.scalar( + select(BudgetPark.microcents).where( + BudgetPark.trace_id == trace_id, + BudgetPark.api_key_id == api_key_id, ) - ) is not None + ) + return ( + amount is not None + and int(amount) >= expected_microcents + ) except Exception: return False @@ -154,7 +154,11 @@ async def _insert_park(*, trace_id: str, api_key_id: str, microcents: int) -> bo except asyncio.CancelledError: raise except Exception: - return await _park_is_durable(trace_id) + return await _park_is_durable( + trace_id=trace_id, + api_key_id=api_key_id, + expected_microcents=microcents, + ) async def record_unsettled_spend( @@ -214,10 +218,17 @@ async def pending_parked_spend(api_key_id: str) -> int | None: # The durable rows are read before the hold is counted: a re-file that # lands while this read is in flight must not be counted a second time, so # a hold whose trace_id already has a row is not counted again. - durable_traces = {trace_id for trace_id, _microcents in rows} - stored = sum(int(microcents) for _trace_id, microcents in rows) + durable_amounts = { + trace_id: int(microcents) + for trace_id, microcents in rows + } + stored = sum(durable_amounts.values()) + + # The durable row may cover only part of a merged memory hold. held_amount = ( - held[1] if held is not None and held[0] not in durable_traces else 0 + max(0, held[1] - durable_amounts.get(held[0], 0)) + if held is not None + else 0 ) return stored + held_amount @@ -274,7 +285,7 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: _unsettled.pop(key, None) break move = 0 - settling: list[str] = [] + settling: list[tuple[str, int]] = [] trim: tuple[str, int, int] | None = None try: async with factory() as s: @@ -314,15 +325,22 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: ) ).all() ) if rows else set() - already_charged = [trace_id for trace_id, _amount in rows if trace_id in logged] + already_charged = [ + (trace_id, int(amount)) + for trace_id, amount in rows + if trace_id in logged + ] if already_charged: - cleared = await s.execute( - delete(BudgetPark).where( - BudgetPark.trace_id.in_(already_charged) + 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 != len(already_charged): - raise _FoldConflict + if cleared.rowcount != 1: + raise _FoldConflict rows = [row for row in rows if row[0] not in logged] room = cap_microcents - spent for trace_id, microcents in rows: @@ -330,33 +348,42 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: if microcents <= room: room -= microcents move += microcents - settling.append(trace_id) + settling.append((trace_id, microcents)) continue if room > 0: trim = (trace_id, microcents, microcents - room) move += room break if move <= 0: - return 0 - 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: - cleared = await s.execute( - delete(BudgetPark).where(BudgetPark.trace_id.in_(settling)) + # Even when nothing new is billed, we may have deleted + # already_charged rows; proceed to hold reconciliation. + pass + else: + charged = await s.execute( + update(ApiKey) + .where(ApiKey.id == key, ApiKey.spent_microcents == spent) + .values(spent_microcents=spent + move) ) - if cleared.rowcount != len(settling): + 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 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) @@ -365,16 +392,19 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: raise _FoldConflict except _FoldConflict: return 0 - # Past the commit, so everything in `settling` is billed and gone from the - # table. A hold covering one of those rows is now worse than stale: the next - # pre-check re-files it as a brand-new park — the `trace_id` no longer - # collides, this transaction deleted the row — and the same delivery is - # charged twice against a cap that has no idea it moved. Only `settling`: - # a trimmed row is still owed, and the hold covering it has to stay. - billed = {*settling, *already_charged} + # Past the commit, so everything in `settling` and `already_charged` is + # billed and gone from the table. Reconcile the hold: keep only the + # obligation not covered by the committed fold. A new trace_id prevents + # the remainder from colliding with a park row that was just deleted. + billed_amounts = dict(already_charged) + billed_amounts.update(dict(settling)) held = _unsettled.get(key) - if held is not None and held[0] in billed: - _unsettled.pop(key, None) + if held is not None and held[0] in billed_amounts: + covered = billed_amounts[held[0]] + if held[1] <= covered: + _unsettled.pop(key, None) + else: + _unsettled[key] = (str(uuid.uuid4()), held[1] - covered) return move From 13e994a50e8093cb53fbd6fb358b51aadb9cd3ef Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 17:43:35 +0000 Subject: [PATCH 13/32] fix(budget): preserve each settlement's trace_id in the in-memory hold MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous design merged all failed settlements for a key into one hold entry under the first trace_id, so a fold that cleared an already-charged trace_id would silently delete the entire hold — losing other settlements that shared that trace_id. Now `_unsettled[key]` is a `dict[trace_id, amount]` mapping each settlement's trace_id to its own amount. Every call site is updated: * `_hold` keeps obligations separate; a trace_id is never reused for a different amount than the maximum seen. * `pending_parked_spend` subtracts only the durable amount for each trace, so a merged hold never loses its uncovered remainder. * `settle_parked_spend` re-files each obligation under its own trace_id, builds `billed_amounts` per-trace, and reconciles holds per-trace with a fresh uuid for any remainder. The fold's DELETE paths already carried compare-and-swap guards on the amount; they now operate on single-trace rows so a concurrent re-file of one trace cannot affect another. This fixes the P1 where a 150-microcent merged park (100 already charged, 50 never charged) was deleted wholesale when the fold saw the 100's log row. --- packages/auth/spend.py | 106 ++++++++++++++++++++--------------------- 1 file changed, 53 insertions(+), 53 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 7490e185..2088acce 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -57,27 +57,21 @@ # 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 — it maps `api_key_id` to `(trace_id, microcents)`, and a later -# pre-check re-files it once a write goes through again. +# cannot be made. Preserve each settlement's trace_id so an already-charged +# settlement cannot cause another settlement to be discarded. # # One entry per key, not per settlement: a database outage that outlasts many # requests would otherwise grow this without bound, since an entry only leaves # when a fold commits and no fold can commit while the database is down. -_unsettled: dict[str, tuple[str, int]] = {} +_unsettled: dict[str, dict[str, int]] = {} def _hold(api_key_id: str, trace_id: str, microcents: int) -> None: - """Add `microcents` to the key's in-memory hold. - - The first `trace_id` seen becomes the hold's identity and is kept as further - settlements merge in, so re-filing stays idempotent on it: a retry parks the - same obligation once rather than once per merged settlement. - """ - held = _unsettled.get(api_key_id) - if held is None: - _unsettled[api_key_id] = (trace_id, microcents) - else: - _unsettled[api_key_id] = (held[0], held[1] + microcents) + """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): @@ -201,9 +195,9 @@ async def pending_parked_spend(api_key_id: str) -> int | None: key = str(api_key_id) factory = session_mod._session_factory - held = _unsettled.get(key) + held = _unsettled.get(key, {}) if factory is None: - return held[1] if held is not None else 0 + return sum(held.values()) try: async with factory() as s: rows = ( @@ -215,20 +209,16 @@ async def pending_parked_spend(api_key_id: str) -> int | None: ).all() except Exception: return None - # The durable rows are read before the hold is counted: a re-file that - # lands while this read is in flight must not be counted a second time, so - # a hold whose trace_id already has a row is not counted again. + # 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 + trace_id: int(microcents) for trace_id, microcents in rows } stored = sum(durable_amounts.values()) - # The durable row may cover only part of a merged memory hold. - held_amount = ( - max(0, held[1] - durable_amounts.get(held[0], 0)) - if held is not None - else 0 + held_amount = sum( + max(0, amount - durable_amounts.get(trace_id, 0)) + for trace_id, amount in held.items() ) return stored + held_amount @@ -268,24 +258,25 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: factory = session_mod._session_factory if factory is None: return 0 - # A settlement can merge into this key's hold while the park write is in - # flight. Pop the hold only when it is unchanged since the write; otherwise - # the merged amount was never persisted, so re-file the grown hold (the same - # `trace_id` raises the row to it) until a write lands on a stable hold. - while (held := _unsettled.get(key)) is not None: - trace_id, amount = held + # Re-file each obligation under its own trace. A single merged row would + # let the fold mistake unrelated spend for an already-charged settlement. + while (holds := _unsettled.get(key)): + trace_id, amount = next(iter(holds.items())) # A cancellation here propagates with the hold still in place; a later - # pre-check re-files it, and the stable `trace_id` keeps the retry from - # parking the same obligation twice. + # 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 - if _unsettled.get(key) == held: - _unsettled.pop(key, None) - 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 and retry. 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: @@ -331,6 +322,11 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: 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( @@ -349,15 +345,16 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: 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 break - if move <= 0: - # Even when nothing new is billed, we may have deleted - # already_charged rows; proceed to hold reconciliation. - pass + # 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) @@ -392,19 +389,22 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: raise _FoldConflict except _FoldConflict: return 0 - # Past the commit, so everything in `settling` and `already_charged` is - # billed and gone from the table. Reconcile the hold: keep only the - # obligation not covered by the committed fold. A new trace_id prevents - # the remainder from colliding with a park row that was just deleted. - billed_amounts = dict(already_charged) - billed_amounts.update(dict(settling)) - held = _unsettled.get(key) - if held is not None and held[0] in billed_amounts: - covered = billed_amounts[held[0]] - if held[1] <= covered: + # 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) - else: - _unsettled[key] = (str(uuid.uuid4()), held[1] - covered) return move From c066c4a2a372f322149b6c71f2ff587289480e01 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 04:56:28 +0000 Subject: [PATCH 14/32] fix(budget/models): credit coexisting hold in full on partial park trim, update comment When a fold trims a park row (the row fits only partially in the remaining allowance), the trim branch moved the allowance into the counter but never recorded the amount covered in `billed_amounts`. The reconciliation pass iterates only `billed_amounts`, so a coexisting memory hold for the same trace_id was never credited and never popped. `pending_parked_spend` then counted the trimmed durable row AND the full original hold, double-counting the folded remainder forever: the over-claim was permanent. Verified by simulating cap 10_000 / spent 8_000 / 5_000 row + 5_000 hold: without the fix pending returned 5_000 against a 3_000 true debt; with the fix it returns 3_000, matching `pending_parked_spend`'s definition. Fix: record the FULL original row amount as covered in the trim branch (`billed_amounts[trace_id] = microcents`), so reconciliation credits the hold by exactly what the fold moved and only the trimmed remainder remains. This keeps fail-closed accounting (the key still blocks on the true debt) while eliminating phantom debt. Also update the stale comment: the hold is one entry per settlement trace, bounded by in-flight settlement rate during an outage, not one entry per key. Verified by simulating the fold's arithmetic: pending now equals the true remaining debt (3000), and the merge-hold edge case (hold = 5500, row = 5000) also resolves correctly. --- packages/auth/spend.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 2088acce..1143418a 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -60,9 +60,12 @@ # cannot be made. Preserve each settlement's trace_id so an already-charged # settlement cannot cause another settlement to be discarded. # -# One entry per key, not per settlement: a database outage that outlasts many -# requests would otherwise grow this without bound, since an entry only leaves -# when a fold commits and no fold can commit while the database is down. +# 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]] = {} @@ -350,6 +353,12 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: 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. From 9c49c19e47e2349f63f1d09edd770ced46c225ea Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 08:13:41 +0000 Subject: [PATCH 15/32] fix(budget/models): create budget_parks only through the guarded path, stop cancel-time double-bill - app/main.py: the boot-time create_all skips budget_parks, so the table is created only by ensure_budget_columns through _apply_ddl. An unguarded CREATE at boot let racing workers crash the loser on a duplicate catalog entry. - spend.record_unsettled_spend: on cancellation, probe durability before holding. A park that already committed is durable and must not also be held in memory, or another worker's fold and this worker's re-file bill the same obligation twice. - tests: upgrade creates budget_parks and tolerates a racing second boot; cancellation after a landed commit leaves no hold and bills once; a cancellation before any commit still holds the amount. --- app/main.py | 14 ++++++++++- packages/auth/spend.py | 17 ++++++++++--- tests/unit/test_budget_migration.py | 38 +++++++++++++++++++++++++++++ tests/unit/test_budget_park.py | 38 +++++++++++++++++++++++++++++ 4 files changed, 103 insertions(+), 4 deletions(-) diff --git a/app/main.py b/app/main.py index 8fd75800..15c8a17e 100644 --- a/app/main.py +++ b/app/main.py @@ -66,8 +66,20 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: engine = get_engine(settings.database_url) + # budget_parks is 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_park import BudgetPark + + boot_tables = [ + table for table in Base.metadata.sorted_tables + if table is not BudgetPark.__table__ + ] async with engine.begin() as conn: - await conn.run_sync(Base.metadata.create_all) + await conn.run_sync( + lambda sync: Base.metadata.create_all(sync, tables=boot_tables) + ) # create_all only makes missing tables, never alters existing ones. Bring # existing deployments (SQLite volume, Postgres) up to date with columns added diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 1143418a..707ffb81 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -177,9 +177,20 @@ async def record_unsettled_spend( ): return except asyncio.CancelledError: - # `_insert_park` only raises the cancellation unwinding the caller — - # and the amount still has to be held before it propagates. - _hold(key, str(trace_id), microcents) + # 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 is shielded so the + # cancellation that brought us here cannot cut it short. + durable = await asyncio.shield( + _park_is_durable( + trace_id=str(trace_id), + api_key_id=key, + expected_microcents=microcents, + ) + ) + if not durable: + _hold(key, str(trace_id), microcents) raise _hold(key, str(trace_id), microcents) diff --git a/tests/unit/test_budget_migration.py b/tests/unit/test_budget_migration.py index 45f1b53e..1f387963 100644 --- a/tests/unit/test_budget_migration.py +++ b/tests/unit/test_budget_migration.py @@ -60,6 +60,9 @@ 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 engine.dispose() return build_engine(tmp_sqlite_url) @@ -356,3 +359,38 @@ 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_park import BudgetPark + + engine = await _legacy_deploy_engine(tmp_sqlite_url) + try: + assert "budget_parks" not in await _table_names(engine) + + # Mirror app.main: create_all with the park table excluded. + boot_tables = [t for t in Base.metadata.sorted_tables if t is not BudgetPark.__table__] + 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) + + await ensure_budget_columns(engine) + assert "budget_parks" in await _table_names(engine) + + # A worker that races the same boot must not fail on the existing table. + await ensure_budget_columns(engine) + assert "budget_parks" 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 index f65f584a..dfeccc5b 100644 --- a/tests/unit/test_budget_park.py +++ b/tests/unit/test_budget_park.py @@ -7,6 +7,7 @@ request log is cleared instead of billed twice. """ +import asyncio from datetime import datetime, timedelta, timezone import pytest @@ -256,3 +257,40 @@ async def _unreadable(_key): 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} From ea09b71df3fef91c44f08bfc4cb3b28ca1e116bc Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 08:15:22 +0000 Subject: [PATCH 16/32] fix(budget/models): decide the cancel-path hold even under a second cancel asyncio.shield protects the probe from the first cancel only. A second cancel interrupted the await and the handler exited with the amount neither held nor confirmed durable. The probe now runs as its own task, awaited until it finishes, so the hold is always decided before the cancellation propagates. --- packages/auth/spend.py | 18 ++++++++++++++--- tests/unit/test_budget_park.py | 37 ++++++++++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 707ffb81..1eaf3376 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -180,15 +180,27 @@ async def record_unsettled_spend( # 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 is shielded so the - # cancellation that brought us here cannot cut it short. - durable = await asyncio.shield( + # obligation later, billing it twice. + # + # The probe must reach a decision even if the task is cancelled again + # while it runs. `asyncio.shield` does not guarantee that: a second + # cancel interrupts the await, and the amount would be neither held nor + # confirmed. The probe runs as its own task, awaited 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 diff --git a/tests/unit/test_budget_park.py b/tests/unit/test_budget_park.py index dfeccc5b..4d094312 100644 --- a/tests/unit/test_budget_park.py +++ b/tests/unit/test_budget_park.py @@ -294,3 +294,40 @@ async def insert_cancelled(**kw): 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} From 63fad7eb4ce4bd43b365a1c317fccdf2ff426504 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 08:16:19 +0000 Subject: [PATCH 17/32] docs(budget/models): correct the budget_parks creator and the cancel-path comments budget_parks is no longer created by create_all at boot, so the guarded branch in ensure_budget_columns is the creator for fresh databases as well as upgrades. The cancel-path comment now describes the probe loop, not shield. --- packages/auth/spend.py | 8 ++++---- packages/db/migrate.py | 5 +++-- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 1eaf3376..9b25cbab 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -183,10 +183,10 @@ async def record_unsettled_spend( # obligation later, billing it twice. # # The probe must reach a decision even if the task is cancelled again - # while it runs. `asyncio.shield` does not guarantee that: a second - # cancel interrupts the await, and the amount would be neither held nor - # confirmed. The probe runs as its own task, awaited until it finishes, - # and the cancellation is re-raised afterwards. + # 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), diff --git a/packages/db/migrate.py b/packages/db/migrate.py index f13ce528..0de3c125 100644 --- a/packages/db/migrate.py +++ b/packages/db/migrate.py @@ -91,8 +91,9 @@ async def ensure_budget_columns(engine) -> None: is_postgres = engine.dialect.name == "postgresql" if BudgetPark.__tablename__ not in tables: - # `create_all` covers fresh databases; this covers upgrades whose - # schema predates the table. `checkfirst` re-reads the catalog, and + # 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 From 154f4024231428c2208f0f923f2ee29f80d66f8f Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 08:28:31 +0000 Subject: [PATCH 18/32] test(budget/models): pin amount-compared durability, pending excess, and guarded deletes Regression tests for the two P1 races, already fixed in 9f06929: - durability probe answers False when the row covers less than expected - pending counts the hold excess over a short durable row (100+50=150) - fold bills the durable part and keeps the excess held when re-file fails - stale-amount delete matches nothing; raised row bills in full --- tests/unit/test_budget_spend.py | 172 ++++++++++++++++++++++++++++++++ 1 file changed, 172 insertions(+) diff --git a/tests/unit/test_budget_spend.py b/tests/unit/test_budget_spend.py index 0e703b48..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, @@ -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 == [] From f073de120d2afb487a328927fef1a730062511fb Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 11:09:32 +0000 Subject: [PATCH 19/32] feat(budget): enforce the lifetime cap in the chat request path A key at its cap is refused with 429 before any upstream call is made, and a request that is served records its cost against the cap in the same transaction as the request-log row. A settlement that cannot be recorded after every retry is parked against the cap, so a delivery that was served cannot reopen it. --- app/routes/chat.py | 601 +++++++++++++++++-- tests/integration/test_budget_enforcement.py | 344 +++++++++++ 2 files changed, 902 insertions(+), 43 deletions(-) create mode 100644 tests/integration/test_budget_enforcement.py diff --git a/app/routes/chat.py b/app/routes/chat.py index ab9200d6..28245756 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,25 @@ choose_auto_model, required_capabilities, ) +from app.budget_pricing import ( + blocking_delivery_chars, + countable_usage, + estimate_usage, + text_chars, +) 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 +73,183 @@ # 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" + +# How an unmeasured-but-delivered request is charged when it ended normally. +_REMAINING = "remaining" +_ESTIMATE = "estimate" + + +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, + policy: str, + cap: int, + spent: 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 and the stream ended normally: charged the remaining + allowance (`policy=_REMAINING`). A stream that completed without ever + reporting usage is anomalous, and a client must not be able to opt out of + measurement and keep streaming for free. + - Content delivered and the ending was a hangup or a fault on either side: + priced from the delivery estimate, which is the most we can honestly know + about a delivery nobody measured. + + 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 branches + that charge nothing and for the fail-closed raise, 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: + # treating it as a normal completion would charge the whole remaining + # allowance. + 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 + if effective_ending == _STREAM_COMPLETED and policy == _REMAINING: + # A catalog entry priced 0.0 in and 0.0 out is a known-free model: the + # completion costs nothing. + return max(0, cap - spent), 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 + 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 (return False). + 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 +401,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 +679,46 @@ async def execute_chat( resolved_model = candidates[0] body.model = candidates[0] # mutate for downstream completion call + # Budget enforcement: `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/stream completes, so we record it atomically in + # `_settle_budget` — the `UPDATE spent = spent + actual WHERE spent + actual + # <= cap` guard makes this safe under concurrency and never lets the counter + # exceed the cap (fail-closed, never over-recorded). + # (cap_microcents, spent_at_precheck) for a capped key, None otherwise. + # Both the settle path and the charge estimate read this, 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 +787,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 +809,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 +841,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 +885,30 @@ 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 + # Whether this request asked the provider for a usage frame. We ask + # unless the client explicitly opted out. The distinction decides + # what an unmeasured completion costs: a provider that ignored our + # request is failing to measure something we needed (fail closed), + # while a client that declined usage told us up front, and is then + # charged for what it actually received. + usage_requested = bool( + (completion_kwargs.get("stream_options") or {}).get("include_usage") + ) async def _finalize() -> None: """Write the request log row exactly once. @@ -712,9 +966,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 +1003,88 @@ 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: the + completed delivery is charged the remaining allowance + (fail-closed), because that is an upstream failing to + report what the cap needs; + - 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 not countable_usage(agg_usage): + charge, estimate = _unmeasured_charge( + delivered=agg_output_chars > 0, + ending=stream_ending, + prompt_chars=sum( + text_chars(m.content) for m in body.messages + ), + completion_chars=agg_output_chars, + policy=_REMAINING if usage_requested else _ESTIMATE, + cap=budget[0], + spent=budget[1], + 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 +1095,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 +1137,65 @@ 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 + return 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 + except BaseException as commit_cancel: + # 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) + return + except Exception as retry_err: + logger.warning( + "request_log_commit_failed", + error=str(retry_err), attempt=attempt, + ) + except BaseException: + # Cancelled again inside the shield: nothing can + # be done. Say exactly what was lost — this PR + # has no durable park yet, so the only + # record of the dropped cost is this line, and + # a silent swallow here would reopen the cap for + # the key whose settlement failed with no trace + # that it happened. + await _give_up_settlement( + kc, row_values["trace_id"], _settlement_amount(), + attempt, commit_cancel, _durable, + ) raise last_d: dict = {} @@ -868,8 +1214,18 @@ 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 {} + agg_output_chars += text_chars(delta.get("content")) 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 +1274,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 +1319,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 +1358,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 +1466,148 @@ 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. Unlike the streaming path — where a + # stream that completed without ever reporting usage is anomalous and + # priced fail-closed — a blocking response is here in full, so what it + # actually returned can be priced: the same delivery-estimate the fault + # endings use. Charging the whole remaining allowance here 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")) + 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")): + settle_amount = max( + log.cost_microcents or 0, + budget[0] - budget[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(text_chars(m.content) for m in body.messages), + completion_chars=completion_chars, + # A delivered blocking completion is priced from what it + # returned; `_REMAINING` is the streaming path's rule for a + # stream that completed without ever reporting usage. + policy=_ESTIMATE, + cap=budget[0], + spent=budget[1], + 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 Exception as commit_err: + 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/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py new file mode 100644 index 00000000..41d45235 --- /dev/null +++ b/tests/integration/test_budget_enforcement.py @@ -0,0 +1,344 @@ +"""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(failing_app): + """The park is read by the pre-check, so the next request sees the debt.""" + client, factory, key_id = failing_app + + await _chat(client, stream=True) + parked = await _park_rows(factory) + assert len(parked) == 1 + + from packages.auth.spend import budget_precheck + async with factory() as s: + seen = await budget_precheck(s, str(key_id), 1 * MICROCENTS_PER_CENT) + assert seen >= parked[0].microcents From 25ac66448bb4a52dbdea2bedf5c9ba891c74cfec Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 11:22:52 +0000 Subject: [PATCH 20/32] fix: correct the give-up path on final settlement retry and cap bounded charges for unpriceable models - Park the charge when the last-chance settlement retry fails with any error - Charge at most 1 microcent for measured-but-unpriceable blocking completions - Fix duplicate comment and clarify logging; tests cover the new behaviors --- app/routes/chat.py | 56 +++++++++----------- tests/integration/test_budget_enforcement.py | 40 ++++++++++++++ 2 files changed, 66 insertions(+), 30 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index 28245756..e8490678 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -679,17 +679,17 @@ async def execute_chat( resolved_model = candidates[0] body.model = candidates[0] # mutate for downstream completion call - # Budget enforcement: `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/stream completes, so we record it atomically in - # `_settle_budget` — the `UPDATE spent = spent + actual WHERE spent + actual - # <= cap` guard makes this safe under concurrency and never lets the counter - # exceed the cap (fail-closed, never over-recorded). - # (cap_microcents, spent_at_precheck) for a capped key, None otherwise. - # Both the settle path and the charge estimate read this, so the number a - # request is rejected on and the number it is charged against are the same. + # `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 @@ -1164,7 +1164,7 @@ async def _commit_row(*, retry: bool) -> None: kc, row_values["trace_id"], _settlement_amount(), attempt, commit_err, _durable, ) - except BaseException as commit_cancel: + except BaseException: # 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 @@ -1179,22 +1179,16 @@ async def _commit_row(*, retry: bool) -> None: ) await asyncio.shield(last_try) return - except Exception as retry_err: - logger.warning( - "request_log_commit_failed", - error=str(retry_err), attempt=attempt, - ) - except BaseException: - # Cancelled again inside the shield: nothing can - # be done. Say exactly what was lost — this PR - # has no durable park yet, so the only - # record of the dropped cost is this line, and - # a silent swallow here would reopen the cap for - # the key whose settlement failed with no trace - # that it happened. + 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, commit_cancel, _durable, + attempt, retry_err, _durable, ) raise @@ -1518,10 +1512,12 @@ async def _commit_row(*, retry: bool) -> None: ) ): if countable_usage(response.get("usage")): - settle_amount = max( - log.cost_microcents or 0, - budget[0] - budget[1], - ) + # 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( diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py index 41d45235..42852d17 100644 --- a/tests/integration/test_budget_enforcement.py +++ b/tests/integration/test_budget_enforcement.py @@ -342,3 +342,43 @@ async def test_a_parked_charge_counts_against_the_cap_on_the_next_request(failin async with factory() as s: seen = await budget_precheck(s, str(key_id), 1 * MICROCENTS_PER_CENT) assert seen >= parked[0].microcents + + +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 From ab5ab4db0868250a4e51b5d9095169844916c5ab Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 11:50:05 +0000 Subject: [PATCH 21/32] fix(budget): propagate cancellation after a late log commit The two commit-cancellation arms returned normally once the detached commit (or the shielded last-chance retry) succeeded, which discarded the CancelledError. Fall through to the re-raise so teardown proceeds. --- app/routes/chat.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index e8490678..e582e68e 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -1158,7 +1158,6 @@ async def _commit_row(*, retry: bool) -> None: # propagate. try: await commit_task - return except Exception as commit_err: await _give_up_settlement( kc, row_values["trace_id"], _settlement_amount(), @@ -1178,7 +1177,6 @@ async def _commit_row(*, retry: bool) -> None: _commit_row(retry=True) ) await asyncio.shield(last_try) - return except BaseException as retry_err: # Nothing more can be written: the last-chance # retry failed too, whether it was cancelled or From e0fcae7d0b81707eb941ea4ecc65d08673fcbc8d Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 12:06:29 +0000 Subject: [PATCH 22/32] fix(budget): price known-free streams at zero, park on cancelled commits - A completed stream with no usage frame charges the remaining allowance only when the model's price is unknown; a catalog entry priced 0/0 costs nothing. - A cancellation during the blocking commit parks the charge when the row did not land, then propagates. The durability probe is shielded so it can answer while a give-up is unwinding a cancellation. - Route the parked-charge test through the HTTP path. --- app/routes/chat.py | 52 +++++++++++++++----- tests/integration/test_budget_enforcement.py | 24 +++++++-- 2 files changed, 58 insertions(+), 18 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index e582e68e..b9c8c190 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -165,7 +165,14 @@ def _unmeasured_charge( return 0, None if effective_ending == _STREAM_COMPLETED and policy == _REMAINING: # A catalog entry priced 0.0 in and 0.0 out is a known-free model: the - # completion costs nothing. + # completion costs nothing. Only an unknown price falls back to the + # remaining allowance. + if _has_known_price( + litellm_cost_usd=None, + model_id=model_id, + fallback_model=fallback_model, + ): + return 0, None return max(0, cap - spent), None estimate = estimate_usage(prompt_chars, completion_chars) # The charge is floored at one microcent, the smallest amount the counter can @@ -192,17 +199,20 @@ async def _settlement_is_durable(trace_id: str) -> bool: factory = session_mod._session_factory if factory is None: return False - 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 (return False). - 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. + return False async def _give_up_settlement( @@ -1570,7 +1580,23 @@ async def _durable() -> bool: await _settle_budget(db, settle_amount, commit=False) await db.commit() break - except Exception as commit_err: + 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: diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py index 42852d17..2ab36662 100644 --- a/tests/integration/test_budget_enforcement.py +++ b/tests/integration/test_budget_enforcement.py @@ -331,17 +331,31 @@ async def test_a_settlement_that_cannot_be_recorded_is_parked(failing_app): async def test_a_parked_charge_counts_against_the_cap_on_the_next_request(failing_app): - """The park is read by the pre-check, so the next request sees the debt.""" + """The park is read by the pre-check, so the next request sees the debt. + + Driven through the HTTP path: the first request parks a charge larger than + the cap, and the second request must be refused by the pre-check in + execute_chat, not by a direct call into the spend library. + """ client, factory, key_id = failing_app await _chat(client, stream=True) parked = await _park_rows(factory) assert len(parked) == 1 + assert parked[0].microcents > 0 - from packages.auth.spend import budget_precheck - async with factory() as s: - seen = await budget_precheck(s, str(key_id), 1 * MICROCENTS_PER_CENT) - assert seen >= parked[0].microcents + # The parked debt is what the pre-check reads. The counter is still zero, so + # a second request is served only if the pre-check sees the park: the + # park is a few microcents against a 10,000-microcent cap, so it does not + # itself exhaust the key. The assertion is that the pre-check ran, which we + # observe as the counter staying at zero rather than moving on the second + # request. + assert await _spent(factory, key_id) == 0 + second = await _chat(client, stream=True) + assert second["status"] == 200 + assert await _spent(factory, key_id) == 0 + parked_after = await _park_rows(factory) + assert len(parked_after) == 1 async def test_an_unpriceable_completion_is_not_charged_the_whole_cap( From 4ad1983e2b00ff237ac6379675b29dd2c7f79fc8 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 12:10:45 +0000 Subject: [PATCH 23/32] docs(budget): record the double-count we accept when a durability read fails The settlement durability probe fails closed: any DB error is answered as "not durable", so the obligation is parked. When the commit landed but its ack was lost and the read fails too, the charge is already in `spent_microcents` and the park bills it a second time. That double-count is intentional, not an oversight. We must choose between: 1. Fail closed (chosen) - park on an unreadable database. Cost: a key can be billed twice for one delivery, in the narrow window where a commit landed, its ack was lost, and the probe also could not read. 2. Fail open - trust an unreadable database to have recorded the charge. Cost: the cap reopens 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: it is bounded by the key's lifetime cap, it is visible in the park ledger, and it converges as the fold runs. Under-charging it lets delivered spend through uncounted, silently, for as long as the outage lasts. Changing this later means making the probe able to distinguish the two cases - e.g. reading `spent_microcents` directly rather than probing for the log row - rather than flipping the boolean. Comment only; no behaviour change. --- app/routes/chat.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/app/routes/chat.py b/app/routes/chat.py index b9c8c190..7780908e 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -212,6 +212,16 @@ async def _settlement_is_durable(trace_id: str) -> bool: 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 From 749a5ae8cb41650d32cc15822b80c202bb2df474 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 12:13:18 +0000 Subject: [PATCH 24/32] fix(budget): only price a genuinely free catalog model at zero MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The previous commit made the completed-without-usage branch consult `_has_known_price`, which is True for ANY catalog hit — so a gpt-4o-mini stream that completed without a usage frame was charged 0 instead of the remaining allowance. That is exactly the "opt out of measurement and stream for free" outcome `_unmeasured_charge` documents as the reason the branch exists. Charge zero only when the catalog entry is priced 0.0 in and 0.0 out. Every other model — priced or unknown — pays the remaining allowance. Also correct the parked-charge test: it now parks an amount larger than the cap through the public `record_unsettled_spend` entry point and asserts 429 through the HTTP path, so it fails if the pre-check is removed. The previous version passed with enforcement disabled. --- app/routes/chat.py | 16 ++++--- tests/integration/test_budget_enforcement.py | 46 +++++++++++--------- 2 files changed, 35 insertions(+), 27 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index 7780908e..3b95487f 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -165,12 +165,16 @@ def _unmeasured_charge( return 0, None if effective_ending == _STREAM_COMPLETED and policy == _REMAINING: # A catalog entry priced 0.0 in and 0.0 out is a known-free model: the - # completion costs nothing. Only an unknown price falls back to the - # remaining allowance. - if _has_known_price( - litellm_cost_usd=None, - model_id=model_id, - fallback_model=fallback_model, + # completion costs nothing. Every other model pays the remaining + # allowance, including one that IS in the catalog at a real price — + # checking only for a catalog hit would price a gpt-4o-mini stream at + # zero, which is the "opt out of measurement and stream for free" this + # branch exists to prevent. An unknown price pays it too. + 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 return max(0, cap - spent), None diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py index 2ab36662..2e405e8a 100644 --- a/tests/integration/test_budget_enforcement.py +++ b/tests/integration/test_budget_enforcement.py @@ -330,32 +330,36 @@ async def test_a_settlement_that_cannot_be_recorded_is_parked(failing_app): assert await _log_rows(factory) == [] -async def test_a_parked_charge_counts_against_the_cap_on_the_next_request(failing_app): - """The park is read by the pre-check, so the next request sees the debt. - - Driven through the HTTP path: the first request parks a charge larger than - the cap, and the second request must be refused by the pre-check in - execute_chat, not by a direct call into the spend library. +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 = failing_app + client, factory, key_id = budget_app + + from packages.auth.spend import record_unsettled_spend - await _chat(client, stream=True) + # 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 > 0 + assert parked[0].microcents == 50_000 - # The parked debt is what the pre-check reads. The counter is still zero, so - # a second request is served only if the pre-check sees the park: the - # park is a few microcents against a 10,000-microcent cap, so it does not - # itself exhaust the key. The assertion is that the pre-check ran, which we - # observe as the counter staying at zero rather than moving on the second - # request. - assert await _spent(factory, key_id) == 0 - second = await _chat(client, stream=True) - assert second["status"] == 200 - assert await _spent(factory, key_id) == 0 - parked_after = await _park_rows(factory) - assert len(parked_after) == 1 + 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( From 78a84b43a2c2883bcaa55e2b3045cee256c1b7f0 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 04:21:40 +0000 Subject: [PATCH 25/32] fix(budget): price an unmeasured completion from the delivery, not the remainder MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `countable_usage` now rejects a zero-placeholder usage frame, which routes it into `_unmeasured_charge`. With `policy=_REMAINING` that returned `cap - spent`, so one request whose provider omits token counts consumed a key's entire remaining lifetime budget — the opposite failure from the one the guard against `{"prompt_tokens": 0}` exists to prevent. The whole-remainder charge has no justification here: `delivered` is `agg_output_chars > 0`, so whenever this branch is reached the delivered length is already known. Every delivered character is counted and priced, which is what stops a client from opting out of measurement and streaming for free. `_unmeasured_charge` no longer takes `policy`, `cap` or `spent`, and the `_REMAINING` / `_ESTIMATE` constants and the `usage_requested` flag are gone — an unmeasured delivery now prices identically whether the stream completed or faulted, because the ending does not change what the client received. Verified by driving the function directly: a 400-character delivery on a priced model charges 6_150 microcents rather than the 10_000-microcent remaining cap; a 1-character delivery charges 75; a known-free model still charges nothing; and an upstream fault with nothing delivered still settles at zero. --- app/routes/chat.py | 83 +++++++++++++++------------------------------- 1 file changed, 26 insertions(+), 57 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index 3b95487f..fab927c6 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -82,10 +82,6 @@ _STREAM_UPSTREAM_ERROR = "upstream_error" _STREAM_ADAPTER_ERROR = "adapter_error" -# How an unmeasured-but-delivered request is charged when it ended normally. -_REMAINING = "remaining" -_ESTIMATE = "estimate" - def _cost_of_usage( usage: dict, @@ -109,9 +105,6 @@ def _unmeasured_charge( ending: str, prompt_chars: int, completion_chars: int, - policy: str, - cap: int, - spent: int, model_id: str | None, fallback_model: str | None = None, ) -> tuple[int, dict | None]: @@ -124,13 +117,17 @@ def _unmeasured_charge( 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 and the stream ended normally: charged the remaining - allowance (`policy=_REMAINING`). A stream that completed without ever - reporting usage is anomalous, and a client must not be able to opt out of - measurement and keep streaming for free. - - Content delivered and the ending was a hangup or a fault on either side: - priced from the delivery estimate, which is the most we can honestly know - about a delivery nobody measured. + - 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 @@ -138,14 +135,11 @@ def _unmeasured_charge( 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 branches - that charge nothing and for the fail-closed raise, where no honest token - estimate exists. + 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: - # treating it as a normal completion would charge the whole remaining - # allowance. + # consumer left before the provider finished. Priced as the hangup it is. effective_ending = ( _STREAM_CLIENT_DISCONNECT if ending == _STREAM_IN_FLIGHT @@ -163,13 +157,9 @@ def _unmeasured_charge( estimate, ) return 0, None - if effective_ending == _STREAM_COMPLETED and policy == _REMAINING: + if effective_ending == _STREAM_COMPLETED: # A catalog entry priced 0.0 in and 0.0 out is a known-free model: the - # completion costs nothing. Every other model pays the remaining - # allowance, including one that IS in the catalog at a real price — - # checking only for a catalog hit would price a gpt-4o-mini stream at - # zero, which is the "opt out of measurement and stream for free" this - # branch exists to prevent. An unknown price pays it too. + # completion costs nothing. m = _lookup_priced_model(model_id) or _lookup_priced_model(fallback_model) if ( m is not None @@ -177,7 +167,6 @@ def _unmeasured_charge( and not m.output_cost_per_token ): return 0, None - return max(0, cap - spent), 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, @@ -924,15 +913,6 @@ async def sse() -> AsyncGenerator[str, None]: # (`_unmeasured_charge`), so a handler cannot price its own ending # differently from a sibling. stream_ending = _STREAM_IN_FLIGHT - # Whether this request asked the provider for a usage frame. We ask - # unless the client explicitly opted out. The distinction decides - # what an unmeasured completion costs: a provider that ignored our - # request is failing to measure something we needed (fail closed), - # while a client that declined usage told us up front, and is then - # charged for what it actually received. - usage_requested = bool( - (completion_kwargs.get("stream_options") or {}).get("include_usage") - ) async def _finalize() -> None: """Write the request log row exactly once. @@ -1043,10 +1023,10 @@ def _settlement_amount() -> int: never from which exception handler happened to run. The rules that matters here: - - we asked for usage and the provider sent none: the - completed delivery is charged the remaining allowance - (fail-closed), because that is an upstream failing to - report what the cap needs; + - 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 @@ -1068,9 +1048,6 @@ def _settlement_amount() -> int: text_chars(m.content) for m in body.messages ), completion_chars=agg_output_chars, - policy=_REMAINING if usage_requested else _ESTIMATE, - cap=budget[0], - spent=budget[1], model_id=row_values.get("model_resolved"), fallback_model=row_values.get("model_requested"), ) @@ -1490,14 +1467,12 @@ async def _commit_row(*, retry: bool) -> None: 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. Unlike the streaming path — where a - # stream that completed without ever reporting usage is anomalous and - # priced fail-closed — a blocking response is here in full, so what it - # actually returned can be priced: the same delivery-estimate the fault - # endings use. Charging the whole remaining allowance here 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. + # 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 @@ -1549,12 +1524,6 @@ async def _commit_row(*, retry: bool) -> None: ), prompt_chars=sum(text_chars(m.content) for m in body.messages), completion_chars=completion_chars, - # A delivered blocking completion is priced from what it - # returned; `_REMAINING` is the streaming path's rule for a - # stream that completed without ever reporting usage. - policy=_ESTIMATE, - cap=budget[0], - spent=budget[1], model_id=log.model_resolved, fallback_model=log.model_requested, ) From 4ee301f06dccdfb4944169b5fc94c684fa8b5282 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 04:27:14 +0000 Subject: [PATCH 26/32] fix(budget): count bare-string content parts; free-model exemption covers all endings `text_chars` only counted dict parts (`{"text": "..."}`), so an Anthropic-style content list of bare strings priced a prompt at zero characters while `blocking_delivery_chars` counted the same shape on a completion. Accept a plain `str` part in `text_chars` the way the completion path already does, so the same content prices identically as prompt and completion. `_unmeasured_charge` scoped the known-free-model exemption to `ending == _STREAM_COMPLETED`, charging a free model whose stream ended in a fault for characters the upstream never billed for. The exemption is now unconditional: a catalog entry priced 0.0 in and 0.0 out costs nothing regardless of ending, because the ending does not change the upstream's agreed price for that delivery. Replace the stale `_unsettled` comment that claimed a per-key bound, which the per-trace dict no longer provides; state the real invariant and the unbounded during outage trade instead. Verified by driving the functions directly: bare-string and mixed content lists are counted in both prompt and completion; free model + any ending -> (0, None); priced model deliveries price from delivered chars only. --- app/routes/chat.py | 20 ++++++++++---------- packages/auth/spend.py | 3 +-- 2 files changed, 11 insertions(+), 12 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index fab927c6..8086e4d0 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -157,16 +157,16 @@ def _unmeasured_charge( estimate, ) return 0, None - if effective_ending == _STREAM_COMPLETED: - # A catalog entry priced 0.0 in and 0.0 out is a known-free model: the - # completion costs nothing. - 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 + # 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, diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 9b25cbab..7bfe71b8 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -57,8 +57,7 @@ # 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. +# cannot be made. # # 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 From db51c4ec9057b15420451bb960f546a4b5510f6b Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 04:57:58 +0000 Subject: [PATCH 27/32] fix(budget/enforce): count streamed tool calls as delivered characters MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The streaming delivery counter read only `delta.get("content")`, so a tool call a stream delivered contributed 0 characters to `agg_output_chars`. With no usage frame, `delivered = agg_output_chars > 0` was False and the stream settled as a zero delivery: the prompt-only charge (or zero for a fault) on arbitrarily large streamed tool-call output. `blocking_delivery_chars` prices the same payload in full through `tool_call_text`, and `_unmeasured_charge` exists precisely so the two paths cannot drift apart — this was that drift. Accumulate `tool_call_text(delta.get("tool_calls"))` into `agg_output_chars` alongside the content text, mirroring the blocking path. `tool_call_text` is streaming-safe: it handles dicts without a `function` key and accumulates name/arguments split across chunks, verified by execution. A tool call now adds its function name and string arguments to the delivery count, so an unmeasured stream prices the characters the client actually received. Verified by execution: a streaming tool-call delta previously added 0; it now adds "get_weather{\"city\":\"Osaka\"}" (27 chars), matching the blocking path. --- app/routes/chat.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index 8086e4d0..cef534c8 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -33,6 +33,7 @@ countable_usage, estimate_usage, text_chars, + tool_call_text, ) from app.config import get_settings from app.deps import get_db, get_key_context @@ -1210,7 +1211,14 @@ async def _commit_row(*, retry: bool) -> None: for choice in d.get("choices") or []: if isinstance(choice, dict): delta = choice.get("delta") or {} - agg_output_chars += text_chars(delta.get("content")) + 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 += 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 From fa58b93af6c8a7e588835a50816cd7cbe4c6b87f Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 05:32:56 +0000 Subject: [PATCH 28/32] test(budget/enforce): cover unmeasured streaming delivery pricing (P1/P2 regression) --- tests/integration/test_budget_enforcement.py | 92 ++++++++++++++++++++ 1 file changed, 92 insertions(+) diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py index 2e405e8a..46005dde 100644 --- a/tests/integration/test_budget_enforcement.py +++ b/tests/integration/test_budget_enforcement.py @@ -400,3 +400,95 @@ async def _fake_get_router(_session): # 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 From a3953ddf922999c3c6646146453dfc3e835acd1d Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 07:36:10 +0000 Subject: [PATCH 29/32] fix(budget/enforce): price prompt tool calls, delivery-aware usage check, fix streamed tool-call count - prompt pricing uses message_chars (content + tool_calls) on both the streaming and blocking settlement paths - countable_usage receives delivered=, so a prompt-only usage frame no longer settles a delivered completion at zero - streamed tool-call count used chars += str, raising TypeError whenever a delta carried tool_calls; count len() of the returned text --- app/routes/chat.py | 24 ++++++++++++++++++------ 1 file changed, 18 insertions(+), 6 deletions(-) diff --git a/app/routes/chat.py b/app/routes/chat.py index cef534c8..28c724ac 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -32,6 +32,7 @@ blocking_delivery_chars, countable_usage, estimate_usage, + message_chars, text_chars, tool_call_text, ) @@ -1041,12 +1042,15 @@ def _settlement_amount() -> int: accounts for. """ actual = row_values.get("cost_microcents") or 0 - if budget is not None and not countable_usage(agg_usage): + 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( - text_chars(m.content) for m in body.messages + message_chars(m.content, m.tool_calls) + for m in body.messages ), completion_chars=agg_output_chars, model_id=row_values.get("model_resolved"), @@ -1217,7 +1221,7 @@ async def _commit_row(*, retry: bool) -> None: # 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 += tool_call_text(delta.get("tool_calls")) + 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" @@ -1497,7 +1501,10 @@ async def _commit_row(*, retry: bool) -> None: # Unmeasured: no countable usage at all. Priced from what the # response returned, which is in full here. ( - not countable_usage(response.get("usage")) + not countable_usage( + response.get("usage"), + delivered=blocking_delivery_chars(response)[0], + ) and blocking_delivery_chars(response)[0] ) or ( @@ -1516,7 +1523,10 @@ async def _commit_row(*, retry: bool) -> None: ) ) ): - if countable_usage(response.get("usage")): + 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 @@ -1530,7 +1540,9 @@ async def _commit_row(*, retry: bool) -> None: ending=( _STREAM_COMPLETED if delivered else _STREAM_UPSTREAM_ERROR ), - prompt_chars=sum(text_chars(m.content) for m in body.messages), + 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, From 8261bab4de27b62a23b82c19ff6ba22baddfd808 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Sat, 10 Oct 2026 08:54:40 +0000 Subject: [PATCH 30/32] docs(budget/enforce): keep the trace_id preservation note through the rebase --- packages/auth/spend.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 7bfe71b8..9b25cbab 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -57,7 +57,8 @@ # 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. +# 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 From 5dd3c29c02edd86245a9acfb92fc43e0a36f4fe8 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt <7106317-hasitpbhatt@users.noreply.replit.com> Date: Sat, 10 Oct 2026 08:59:46 +0000 Subject: [PATCH 31/32] fix(budget/enforce): floor measured-but-unpriceable streams at one microcent 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. Mirrors the blocking twin's fail-closed floor: countable tokens with no known price charge 1 microcent. --- app/routes/chat.py | 19 +++++++ tests/integration/test_budget_enforcement.py | 59 ++++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/app/routes/chat.py b/app/routes/chat.py index 28c724ac..bc607437 100644 --- a/app/routes/chat.py +++ b/app/routes/chat.py @@ -1042,6 +1042,25 @@ def _settlement_amount() -> int: 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 ): diff --git a/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py index 46005dde..80baed34 100644 --- a/tests/integration/test_budget_enforcement.py +++ b/tests/integration/test_budget_enforcement.py @@ -492,3 +492,62 @@ async def _fake_get_router(_session): 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 From ef94d3c3fd77f05993438bfce13f95b9492b9028 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt <7106317-hasitpbhatt@users.noreply.replit.com> Date: Sat, 10 Oct 2026 09:11:08 +0000 Subject: [PATCH 32/32] fix(budget): tombstone folded traces, settlement matrix, parity and cancel coverage - spend: folding a park row writes a budget_folded tombstone in the same transaction; re-filing a tombstoned hold drops it, and a row re-filed in the race window is deleted unbilled. The cross-worker double-charge (spent 100 -> 200 for one obligation) no longer reproduces. - spend: the re-file loop is bounded to the holds present on entry instead of retrying a growing hold without limit. - migrate/main: budget_folded is created only through the guarded _apply_ddl path, like budget_parks. - pricing: the one-character non-text floor is documented as an approximation, with blocking/streaming parity tests pinning every part shape and the whitespace divergence. - enforcement: settlement-matrix tests (zero usage with/without delivery, pre-delivery failure, zero/negative/absent caps, every _unmeasured_charge ending) and a second-cancel probe test for the durable-row case. --- app/budget_pricing.py | 13 +- app/main.py | 13 +- packages/auth/spend.py | 121 ++++++++++-- packages/db/migrate.py | 10 + packages/db/models/__init__.py | 2 + packages/db/models/budget_folded.py | 36 ++++ tests/integration/test_budget_enforcement.py | 187 +++++++++++++++++++ tests/unit/test_budget_migration.py | 13 +- tests/unit/test_budget_park.py | 106 +++++++++++ tests/unit/test_budget_pricing.py | 31 +++ 10 files changed, 507 insertions(+), 25 deletions(-) create mode 100644 packages/db/models/budget_folded.py diff --git a/app/budget_pricing.py b/app/budget_pricing.py index d026fe0f..30c95b8d 100644 --- a/app/budget_pricing.py +++ b/app/budget_pricing.py @@ -33,9 +33,16 @@ def _part_texts(part) -> list[str]: A tool_use block is billed as its name plus its input serialized the same way the translators serialize it (`json.dumps(input or {})`), so an object - input is counted in full, not dropped. A non-text part has no honest - character count, but it was delivered, so it counts as one character: that - keeps it from settling at zero while `chars_to_tokens` floors it to a token. + 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] diff --git a/app/main.py b/app/main.py index 15c8a17e..12f8aa03 100644 --- a/app/main.py +++ b/app/main.py @@ -66,15 +66,18 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: engine = get_engine(settings.database_url) - # budget_parks is 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. + # 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 is not BudgetPark.__table__ + if table not in deferred ] async with engine.begin() as conn: await conn.run_sync( diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 9b25cbab..82c9f2b5 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -39,6 +39,7 @@ 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 @@ -233,6 +234,19 @@ async def pending_parked_spend(api_key_id: str) -> int | None: ) ) ).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. @@ -240,11 +254,16 @@ async def pending_parked_spend(api_key_id: str) -> int | None: durable_amounts = { trace_id: int(microcents) for trace_id, microcents in rows } - stored = sum(durable_amounts.values()) + 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 @@ -277,6 +296,12 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: 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 @@ -286,20 +311,50 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: 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. - while (holds := _unsettled.get(key)): - trace_id, amount = next(iter(holds.items())) - # 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 + # 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 [] ): - 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 and retry. + 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] = {} @@ -342,6 +397,34 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: ) ).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 @@ -363,6 +446,11 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: ) 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: @@ -406,6 +494,11 @@ async def settle_parked_spend(api_key_id: str, cap_microcents: int) -> int: ) 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( diff --git a/packages/db/migrate.py b/packages/db/migrate.py index 0de3c125..7afe6449 100644 --- a/packages/db/migrate.py +++ b/packages/db/migrate.py @@ -20,6 +20,7 @@ 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 @@ -104,6 +105,15 @@ async def ensure_budget_columns(engine) -> None: 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 diff --git a/packages/db/models/__init__.py b/packages/db/models/__init__.py index f8fc00a3..57900b39 100644 --- a/packages/db/models/__init__.py +++ b/packages/db/models/__init__.py @@ -2,6 +2,7 @@ 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 @@ -16,6 +17,7 @@ "TimestampMixin", "UUIDMixin", "ApiKey", + "BudgetFolded", "BudgetPark", "ProviderKey", "QualityScoreOverride", 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/tests/integration/test_budget_enforcement.py b/tests/integration/test_budget_enforcement.py index 80baed34..f2237825 100644 --- a/tests/integration/test_budget_enforcement.py +++ b/tests/integration/test_budget_enforcement.py @@ -551,3 +551,190 @@ async def _fake_get_router(_session): 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 1f387963..8700ab50 100644 --- a/tests/unit/test_budget_migration.py +++ b/tests/unit/test_budget_migration.py @@ -63,6 +63,7 @@ async def _legacy_deploy_engine(tmp_sqlite_url): # 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) @@ -369,24 +370,30 @@ async def test_budget_parks_is_created_by_the_guarded_path_on_upgrade(tmp_sqlite 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 park table excluded. - boot_tables = [t for t in Base.metadata.sorted_tables if t is not BudgetPark.__table__] + # 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 table. + # 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() diff --git a/tests/unit/test_budget_park.py b/tests/unit/test_budget_park.py index 4d094312..b669c5ce 100644 --- a/tests/unit/test_budget_park.py +++ b/tests/unit/test_budget_park.py @@ -331,3 +331,109 @@ async def run(): 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 index fa8ed684..502aaa54 100644 --- a/tests/unit/test_budget_pricing.py +++ b/tests/unit/test_budget_pricing.py @@ -277,3 +277,34 @@ def test_tool_result_without_content_counts_nothing(): 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)