From 47bffe38a1c7afdda040e80709079bf7624dd875 Mon Sep 17 00:00:00 2001 From: soroush5 Date: Thu, 3 Sep 2026 23:09:22 +0330 Subject: [PATCH 1/4] fix(budget): pass messages to token_counter in projected_cost --- litellm/budget_manager.py | 3 +- tests/test_litellm/test_budget_manager.py | 48 +++++++++++++++++++++++ 2 files changed, 49 insertions(+), 2 deletions(-) create mode 100644 tests/test_litellm/test_budget_manager.py diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index dcb5a7cc183..49060117670 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -100,8 +100,7 @@ class BudgetManager: return self.user_dict[user] def projected_cost(self, model: str, messages: list, user: str): - text: Final = "".join(message["content"] for message in messages) - prompt_tokens: Final = litellm.token_counter(model=model, text=text) + prompt_tokens: Final = litellm.token_counter(model=model, messages=messages) prompt_cost, _ = litellm.cost_per_token(model=model, prompt_tokens=prompt_tokens, completion_tokens=0) current_cost: Final = self.user_dict[user].get("current_cost", 0) projected_cost: Final = prompt_cost + current_cost diff --git a/tests/test_litellm/test_budget_manager.py b/tests/test_litellm/test_budget_manager.py new file mode 100644 index 00000000000..fd6d82b6a62 --- /dev/null +++ b/tests/test_litellm/test_budget_manager.py @@ -0,0 +1,48 @@ +import pytest + +from litellm.budget_manager import BudgetManager + + +@pytest.fixture() +def manager() -> BudgetManager: + bm = BudgetManager(project_name="test", client_type="local") + bm.create_budget(total_budget=10, user="u", duration="daily") + return bm + + +def test_projected_cost_string_content(manager: BudgetManager): + cost = manager.projected_cost( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hello"}], + user="u", + ) + assert cost >= 0 + + +def test_projected_cost_vision_content(manager: BudgetManager): + cost = manager.projected_cost( + model="gpt-4o-mini", + messages=[ + { + "role": "user", + "content": [ + {"type": "text", "text": "hi"}, + {"type": "image_url", "image_url": {"url": "http://x/y.png"}}, + ], + } + ], + user="u", + ) + assert cost >= 0 + + +def test_projected_cost_none_and_missing_content(manager: BudgetManager): + assert ( + manager.projected_cost( + model="gpt-4o-mini", + messages=[{"role": "assistant", "content": None}], + user="u", + ) + >= 0 + ) + assert manager.projected_cost(model="gpt-4o-mini", messages=[{"role": "user"}], user="u") >= 0 From 81f19176fc5861a7c64319330823f47f1d4b3c07 Mon Sep 17 00:00:00 2001 From: soroush5 Date: Fri, 4 Sep 2026 18:53:01 +0330 Subject: [PATCH 2/4] test(budget): isolate user_cost.json writes and pin positive projected cost --- tests/test_litellm/test_budget_manager.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/test_litellm/test_budget_manager.py b/tests/test_litellm/test_budget_manager.py index fd6d82b6a62..255b22e34c9 100644 --- a/tests/test_litellm/test_budget_manager.py +++ b/tests/test_litellm/test_budget_manager.py @@ -4,7 +4,10 @@ from litellm.budget_manager import BudgetManager @pytest.fixture() -def manager() -> BudgetManager: +def manager(tmp_path, monkeypatch) -> BudgetManager: + # BudgetManager persists to ./user_cost.json via a background thread; + # isolate cwd so the suite never litters the repo or races parallel workers. + monkeypatch.chdir(tmp_path) bm = BudgetManager(project_name="test", client_type="local") bm.create_budget(total_budget=10, user="u", duration="daily") return bm @@ -16,7 +19,7 @@ def test_projected_cost_string_content(manager: BudgetManager): messages=[{"role": "user", "content": "hello"}], user="u", ) - assert cost >= 0 + assert cost > 0 def test_projected_cost_vision_content(manager: BudgetManager): @@ -33,7 +36,7 @@ def test_projected_cost_vision_content(manager: BudgetManager): ], user="u", ) - assert cost >= 0 + assert cost > 0 def test_projected_cost_none_and_missing_content(manager: BudgetManager): From 11addf2de0622f092ec4b8a4e1f8f7aac6edeb16 Mon Sep 17 00:00:00 2001 From: soroush5 Date: Fri, 4 Sep 2026 19:56:26 +0330 Subject: [PATCH 3/4] fix(budget): avoid remote image fetch in projected cost --- litellm/budget_manager.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/budget_manager.py b/litellm/budget_manager.py index 49060117670..742f648d224 100644 --- a/litellm/budget_manager.py +++ b/litellm/budget_manager.py @@ -100,7 +100,10 @@ class BudgetManager: return self.user_dict[user] def projected_cost(self, model: str, messages: list, user: str): - prompt_tokens: Final = litellm.token_counter(model=model, messages=messages) + # Fixed image estimate: a budget check must never fetch remote + # image URLs (unbounded/chunked bodies exhaust memory, and a check + # should not do network I/O at all). + prompt_tokens: Final = litellm.token_counter(model=model, messages=messages, use_default_image_token_count=True) prompt_cost, _ = litellm.cost_per_token(model=model, prompt_tokens=prompt_tokens, completion_tokens=0) current_cost: Final = self.user_dict[user].get("current_cost", 0) projected_cost: Final = prompt_cost + current_cost From e66267ae38bee3e0d26ff0a0608e8d2cd4366244 Mon Sep 17 00:00:00 2001 From: soroush5 Date: Sat, 5 Sep 2026 14:20:02 +0330 Subject: [PATCH 4/4] test(budget): cover tool-call and provider-field message shapes in projected cost --- tests/test_litellm/test_budget_manager.py | 42 +++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tests/test_litellm/test_budget_manager.py b/tests/test_litellm/test_budget_manager.py index 255b22e34c9..ac645b55035 100644 --- a/tests/test_litellm/test_budget_manager.py +++ b/tests/test_litellm/test_budget_manager.py @@ -49,3 +49,45 @@ def test_projected_cost_none_and_missing_content(manager: BudgetManager): >= 0 ) assert manager.projected_cost(model="gpt-4o-mini", messages=[{"role": "user"}], user="u") >= 0 + + +def test_projected_cost_tool_calls_with_null_content(manager: BudgetManager): + # The expensive shape: assistant turn carrying tool calls and no text. + cost = manager.projected_cost( + model="gpt-4o-mini", + messages=[ + {"role": "user", "content": "what is the weather?"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + }, + ], + user="u", + ) + assert cost > 0 + + +def test_projected_cost_provider_specific_fields(manager: BudgetManager): + # Provider extras (name, cache_control, reasoning_content) must not break counting. + cost = manager.projected_cost( + model="gpt-4o-mini", + messages=[ + {"role": "user", "content": "hi", "name": "soroush"}, + { + "role": "assistant", + "content": [ + {"type": "text", "text": "hello", "cache_control": {"type": "ephemeral"}}, + ], + "reasoning_content": "thinking...", + }, + ], + user="u", + ) + assert cost > 0