From 2aa2788a24f106cfb29e44856b608ffc44be0608 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 9 Oct 2026 10:49:59 +0000 Subject: [PATCH 01/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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/20] 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 3ad6a2b83f4b03a826fb9409320ed3c161448162 Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 16 Oct 2026 04:42:49 +0000 Subject: [PATCH 19/20] fix(budget/models): hold only on a definitive absent probe; never raise a trimmed park - spend: `_park_is_durable` answers None when the probe itself cannot read the database, and both the general-failure and cancellation paths decide through `_decide_park_durability`, which retries until the probe answers True (row durable) or False (row confirmed absent). A hold is created only on a confirmed-absent row: a landed park write whose ack was lost can no longer gain a memory hold beside its own durable row, which another worker's fold turns into a second billing of the same obligation. - spend: the IntegrityError collision path no longer raises the row to the incoming amount. A row can only be lower than its settlement when the fold trimmed it after moving the difference into spent_microcents, so raising it back re-added that billed share (200 row trimmed to 60 after billing 140, then raised to 200 = 340 counted for a 200 obligation). The existing row is the authoritative remaining debt and is accepted exactly as it is. - tests: ack-lost write with a retried probe leaves no hold and bills once; an unlanded write still holds after the probe retries to absent; a second give-up after a partial fold leaves the row at its remainder and totals the obligation once. --- packages/auth/spend.py | 77 +++++++++++------ tests/unit/test_budget_park.py | 145 ++++++++++++++++++++++++++++++++- 2 files changed, 198 insertions(+), 24 deletions(-) diff --git a/packages/auth/spend.py b/packages/auth/spend.py index 9b25cbab..61bcd426 100644 --- a/packages/auth/spend.py +++ b/packages/auth/spend.py @@ -83,12 +83,22 @@ class _FoldConflict(Exception): 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.""" +) -> bool | None: + """Whether the durable row covers the entire expected obligation. + + ``None`` means the probe could not read the database at all — a different + answer from ``False`` (the row is not there). Callers deciding whether to + hold the obligation in memory must never treat an unreadable probe as + absence: a hold beside its own durable row lets another worker fold the + row and this worker re-file the same obligation later, billing it twice. + They retry until this probe answers (`_decide_park_durability`). + """ from packages.db import session as session_mod factory = session_mod._session_factory if factory is None: + # No database is configured, so no row can be durable here; the write + # path does not even attempt one, which makes "absent" a real answer. return False try: async with factory() as s: @@ -103,7 +113,34 @@ async def _park_is_durable( and int(amount) >= expected_microcents ) except Exception: - return False + return None + + +# How long to wait between durability-probe retries while the database is +# unreadable. The probe is cheap and the caller is already unwinding a failure, +# so the delay only keeps an outage from spinning the retry loop. +_PARK_PROBE_RETRY_SECONDS = 0.1 + + +async def _decide_park_durability( + *, trace_id: str, api_key_id: str, expected_microcents: int +) -> bool: + """Answer whether the row is durable, retrying a probe that cannot read. + + The answer decides whether the obligation is held in memory, and it cannot + be guessed: holding a durable obligation double-charges it (the row and the + hold), while dropping an absent one loses the spend. So an unreadable probe + (``None``) is retried until the database can answer. + """ + while True: + durable = await _park_is_durable( + trace_id=trace_id, + api_key_id=api_key_id, + expected_microcents=expected_microcents, + ) + if durable is not None: + return durable + await asyncio.sleep(_PARK_PROBE_RETRY_SECONDS) async def _insert_park(*, trace_id: str, api_key_id: str, microcents: int) -> bool: @@ -111,10 +148,11 @@ async def _insert_park(*, trace_id: str, api_key_id: str, microcents: int) -> bo 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. + second time. On general failure, `_decide_park_durability` probes whether + the write landed anyway, retrying until the database can answer. Returns + False only when the obligation is confirmed not durable, leaving the caller + to hold the amount in memory. A cancellation propagates with the memory + copy still held. """ from packages.db import session as session_mod @@ -132,26 +170,19 @@ async def _insert_park(*, trace_id: str, api_key_id: str, microcents: int) -> bo 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. + # write after an ack loss, or a repeated give-up for the same + # settlement. The row is the authoritative remaining debt for + # the trace: the only other writer that lowers it is the fold's + # trim, and the difference a trim removes moved into + # `spent_microcents` in that transaction. Raising the row back + # to `microcents` would add that already-billed difference a + # second time; accept the row exactly as it is. await s.rollback() - 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( + return await _decide_park_durability( trace_id=trace_id, api_key_id=api_key_id, expected_microcents=microcents, @@ -188,7 +219,7 @@ async def record_unsettled_spend( # 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( + _decide_park_durability( trace_id=str(trace_id), api_key_id=key, expected_microcents=microcents, diff --git a/tests/unit/test_budget_park.py b/tests/unit/test_budget_park.py index 4d094312..6853b21f 100644 --- a/tests/unit/test_budget_park.py +++ b/tests/unit/test_budget_park.py @@ -11,7 +11,8 @@ from datetime import datetime, timedelta, timezone import pytest -from sqlalchemy import select +from sqlalchemy import delete, select, update +from sqlalchemy.exc import OperationalError from sqlalchemy.ext.asyncio import async_sessionmaker from packages.auth import spend @@ -331,3 +332,145 @@ 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} + + +class _CommitFailingSession: + """A session whose `commit` fails, simulating a write outage. + + ``land=True`` commits for real and then raises — an acknowledgement loss: + the row is durable but the writer cannot know it. Otherwise the commit + raises before anything lands. + """ + + def __init__(self, session, *, land: bool): + self._session = session + self._land = land + + def __getattr__(self, name): + return getattr(self._session, name) + + async def __aenter__(self): + await self._session.__aenter__() + return self + + async def __aexit__(self, *exc): + return await self._session.__aexit__(*exc) + + async def commit(self): + if self._land: + await self._session.commit() + raise OperationalError("COMMIT", None, Exception("commit failed")) + + +def _commit_failing_factory(maker, *, land: bool): + """A session factory whose sessions fail their commit as described above.""" + + def factory(): + return _CommitFailingSession(maker(), land=land) + + return factory + + +async def test_ack_lost_write_is_not_held_when_the_probe_has_to_retry( + factory, monkeypatch +): + """A park that landed must never gain a hold on a failed probe. + + Regression for the P1 where the general-failure path probed once and read + an unreadable probe as "row absent": the obligation was held on top of its + own durable row, another worker folded the row, and this worker's hold + re-filed the same obligation and billed it a second time. + """ + await _make_key(factory, key_id="k", spent=0) + real_probe = spend._park_is_durable + probes = {"n": 0} + + async def unreadable_once(**kw): + probes["n"] += 1 + if probes["n"] == 1: + return None # the database is unreachable at probe time + return await real_probe(**kw) + + monkeypatch.setattr(spend, "_park_is_durable", unreadable_once) + monkeypatch.setattr(spend, "_PARK_PROBE_RETRY_SECONDS", 0) + previous = session_mod._session_factory + session_mod._session_factory = _commit_failing_factory(factory, land=True) + try: + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=600) + finally: + session_mod._session_factory = previous + + # The probe was retried instead of being read as "absent", so no hold. + assert probes["n"] == 2 + assert spend._unsettled.get("k") in (None, {}) + assert await _park_rows(factory, "k") == [("t", 600)] + + # Another worker folds the durable row while this worker holds nothing. + async with factory() as s: + await s.execute( + update(ApiKey).where(ApiKey.id == "k").values(spent_microcents=600) + ) + await s.execute(delete(BudgetPark).where(BudgetPark.trace_id == "t")) + await s.commit() + + # Nothing is left to re-file, so the obligation is billed exactly once. + assert await settle_parked_spend("k", cap_microcents=10_000) == 0 + assert await _spent(factory, "k") == 600 + assert await _park_rows(factory, "k") == [] + + +async def test_unlanded_write_with_a_retried_probe_still_holds_the_amount( + factory, monkeypatch +): + """A park that never landed is held — but only once the probe says absent.""" + await _make_key(factory, key_id="k", spent=0) + real_probe = spend._park_is_durable + probes = {"n": 0} + + async def unreadable_once(**kw): + probes["n"] += 1 + if probes["n"] == 1: + return None + return await real_probe(**kw) + + monkeypatch.setattr(spend, "_park_is_durable", unreadable_once) + monkeypatch.setattr(spend, "_PARK_PROBE_RETRY_SECONDS", 0) + previous = session_mod._session_factory + session_mod._session_factory = _commit_failing_factory(factory, land=False) + try: + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=600) + finally: + session_mod._session_factory = previous + + assert probes["n"] == 2 + assert await _park_rows(factory, "k") == [] + assert spend._unsettled.get("k") == {"t": 600} + + +async def test_second_give_up_after_a_partial_fold_does_not_raise_the_row(factory): + """A repeated park write must not resurrect debt the fold already billed. + + Regression for the P1 where the collision path raised the row back to the + incoming full amount: after the fold billed part of a parked obligation + and trimmed the row to the remainder, a repeated give-up for the same + settlement raised the row to the full amount again, double-counting the + share already moved into `spent_microcents`. + """ + await _make_key(factory, key_id="k", spent=0) + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=200) + + # Room for 140 of the 200: the fold bills 140 and trims the row to 60. + assert await settle_parked_spend("k", cap_microcents=140) == 140 + assert await _spent(factory, "k") == 140 + assert await _park_rows(factory, "k") == [("t", 60)] + + # The settlement is retried and gives up again for the same trace. + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=200) + + assert spend._unsettled.get("k") in (None, {}) + assert await _park_rows(factory, "k") == [("t", 60)] + + # The row is the remaining debt: 140 + 60 = the 200 obligation, once. + assert await settle_parked_spend("k", cap_microcents=10_000) == 60 + assert await _spent(factory, "k") == 200 + assert await _park_rows(factory, "k") == [] From 166f1b2f6cbdfc8318de63df50ffea606312a4cb Mon Sep 17 00:00:00 2001 From: Hasit Bhatt Date: Fri, 16 Oct 2026 05:31:16 +0000 Subject: [PATCH 20/20] fix(budget): tombstone folded traces so a repeated give-up cannot re-bill MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The park's idempotency was one-shot: a full bill deleted the row in the same transaction that moved its microcents into `spent_microcents`, which freed the unique `trace_id`. A repeated give-up for the same settlement (a retried ack-lost write, or a settlement retry) then re-inserted the full amount, and the next fold — finding no request-log row for a charge that never landed, so the `already_charged` guard cannot fire — billed the same obligation a second time (spent 100 -> 200 for one settlement). - spend: folding a fully-billed park row writes a `budget_folded` tombstone in the same transaction. `pending_parked_spend` excludes tombstoned traces from rows and holds, the re-file loop drops a hold whose trace another worker already folded, and a row re-filed in the race window is deleted unbilled. The obligation is counted once. - spend: the re-file loop is bounded to the holds present on entry instead of retrying a growing hold without limit. - migrate/main/models: `budget_folded` is created only through the guarded _apply_ddl path, like `budget_parks`; the boot-time create_all skips it. - tests: the repeated give-up after a full fold, the cross-worker hold drop, the pending exclusion, the stale-row deletion, and the upgrade path that creates both ledger tables. --- 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/unit/test_budget_migration.py | 13 ++- tests/unit/test_budget_park.py | 90 +++++++++++++++++++++ 7 files changed, 263 insertions(+), 22 deletions(-) create mode 100644 packages/db/models/budget_folded.py 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 61bcd426..9a999463 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 @@ -264,6 +265,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. @@ -271,11 +285,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 @@ -308,6 +327,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 @@ -317,20 +342,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] = {} @@ -373,6 +428,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 @@ -394,6 +477,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: @@ -437,6 +525,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/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 6853b21f..6871515a 100644 --- a/tests/unit/test_budget_park.py +++ b/tests/unit/test_budget_park.py @@ -474,3 +474,93 @@ async def test_second_give_up_after_a_partial_fold_does_not_raise_the_row(factor assert await settle_parked_spend("k", cap_microcents=10_000) == 60 assert await _spent(factory, "k") == 200 assert await _park_rows(factory, "k") == [] + + +async def test_folded_trace_is_never_rebilled_across_workers(factory): + """Worker A holds T while worker B folds T's row: no double charge. + + The commit landed but the ack was lost, so A's hold coexists with its + durable row. B's fold bills the row and leaves a tombstone; A's next + pre-check must drop the stale hold instead of re-filing it. + """ + from packages.db.models.budget_folded import BudgetFolded + + await _make_key(factory, key_id="k", spent=0) + async with factory() as s: + s.add(BudgetPark(trace_id="T", api_key_id="k", microcents=100)) + await s.commit() + # A's view: the hold survived beside its durable row. + spend._hold("k", "T", 100) + # B's fold (B never saw the hold): bills 100, deletes the row. + holds = spend._unsettled.pop("k") + assert await settle_parked_spend("k", cap_microcents=10_000) == 100 + spend._unsettled["k"] = holds + # A's pre-check: the tombstone drops the hold, nothing is re-filed. + assert await settle_parked_spend("k", cap_microcents=10_000) == 0 + assert await _spent(factory, "k") == 100 + assert await _park_rows(factory, "k") == [] + assert spend._unsettled.get("k") in (None, {}) + async with factory() as s: + tombstoned = ( + await s.execute( + select(BudgetFolded.trace_id).where(BudgetFolded.api_key_id == "k") + ) + ).scalars().all() + assert list(tombstoned) == ["T"] + + +async def test_tombstoned_hold_is_not_counted_in_pending(factory): + """A hold whose trace was already folded is spent, not outstanding.""" + from packages.db.models.budget_folded import BudgetFolded + + await _make_key(factory, key_id="k", spent=100) + async with factory() as s: + s.add(BudgetFolded(trace_id="T", api_key_id="k", microcents=100)) + await s.commit() + spend._hold("k", "T", 100) + assert await pending_parked_spend("k") == 0 + + +async def test_stale_refiled_row_is_deleted_unbilled(factory): + """A row re-filed after its trace was tombstoned is removed, not billed. + + The race window: A re-inserts T after B's fold tombstoned it but before + A's fold reads. The fold must delete the stale row without moving the + counter a second time. + """ + from packages.db.models.budget_folded import BudgetFolded + + await _make_key(factory, key_id="k", spent=100) + async with factory() as s: + s.add(BudgetFolded(trace_id="T", api_key_id="k", microcents=100)) + s.add(BudgetPark(trace_id="T", api_key_id="k", microcents=100)) + await s.commit() + assert await settle_parked_spend("k", cap_microcents=10_000) == 0 + assert await _spent(factory, "k") == 100 + assert await _park_rows(factory, "k") == [] + + +async def test_repeated_give_up_after_a_full_fold_is_not_rebilled(factory): + """The unique trace key cannot be freed by the fold: re-parking is not re-billing. + + Regression for the P1 where a full bill deleted the park row and left the + `trace_id` free: a repeated give-up for the same settlement re-inserted the + full amount, and the next fold — finding no request-log row for a charge + that never landed — billed the same obligation a second time. + """ + await _make_key(factory, key_id="k", spent=0) + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=100) + assert await settle_parked_spend("k", cap_microcents=10_000) == 100 + assert await _spent(factory, "k") == 100 + assert await _park_rows(factory, "k") == [] + + # The settlement is retried and gives up again for the same trace: the + # write lands, because the fold freed the unique key when it deleted the + # row. It must not read as outstanding debt. + await record_unsettled_spend(trace_id="t", api_key_id="k", microcents=100) + assert await pending_parked_spend("k") == 0 + + # The next fold deletes the stale re-file unbilled: one obligation, one bill. + assert await settle_parked_spend("k", cap_microcents=10_000) == 0 + assert await _spent(factory, "k") == 100 + assert await _park_rows(factory, "k") == []