From 2286bf3eca414cc24e0a03b008a7a4e6b9647c44 Mon Sep 17 00:00:00 2001
From: mynkyu <mynkyu@bucketplace.net>
Date: Thu, 27 Aug 2026 18:30:16 +0900
Subject: [PATCH 01/14] fix(router): stamp model_group when retrieving a batch

Batch token usage is accounted on the retrieve call, not on create: a provider
only reports token counts once the job finishes, so the usage arrives on
aretrieve_batch and that is the spend log row the tokens land on.

Router.acreate_batch stamps the requested model group into its metadata, but
Router.aretrieve_batch never did. A batch is retrieved by id, so the request
carries no model, and the router fans the lookup out over its deployments -
leaving model_group unset on the one record that carries the tokens.
/global/activity/model groups the spend logs by model_group, so every batch's
tokens were bucketed under an empty group.

Stamp the model group inside the per-deployment retrieve attempt, preferring an
explicitly requested group and otherwise using the model_name of the deployment
that answered, which is unambiguous even when the request named no model. An
existing model_group in the metadata is left untouched, so nothing that already
resolves a group changes.

Scope is limited to aretrieve_batch: acompletion, aresponses and acreate_batch
logging are untouched, and cost/spend attribution by model is unchanged.

Signed-off-by: mynkyu <mynkyu@bucketplace.net>
---
 litellm/router.py                             |   9 ++
 .../test_router_batch_retrieve_model_group.py | 118 ++++++++++++++++++
 2 files changed, 127 insertions(+)
 create mode 100644 tests/test_litellm/test_router_batch_retrieve_model_group.py

diff --git a/litellm/router.py b/litellm/router.py
index 3f450661946..c6c25b6be17 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -6137,6 +6137,8 @@ class Router:
         """
         try:
             parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs)
+            requested_model_group: Final = model
+            metadata_variable_name: Final = _get_router_metadata_variable_name(function_name="aretrieve_batch")
             if model is not None:
                 filtered_model_list: (
                     list[DeploymentTypedDict] | list[dict] | dict | None
@@ -6173,6 +6175,13 @@ class Router:
                         kwargs=new_kwargs,
                         function_name="aretrieve_batch",
                     )
+                    ## STAMP THE MODEL GROUP FOR SPEND TRACKING ##
+                    # A batch is retrieved by id, so the request carries no model group of its
+                    # own - only the deployment that answered knows it. Batch token usage lands
+                    # on this retrieve call (the provider reports counts once the job finishes),
+                    # so without this the tokens are logged under an empty model_group.
+                    model_group: Final = requested_model_group or model_name["model_name"]
+                    new_kwargs[metadata_variable_name].setdefault("model_group", model_group)
                     new_kwargs.pop("custom_llm_provider", None)
                     data.pop("custom_llm_provider", None)
                     return await litellm.aretrieve_batch(
diff --git a/tests/test_litellm/test_router_batch_retrieve_model_group.py b/tests/test_litellm/test_router_batch_retrieve_model_group.py
new file mode 100644
index 00000000000..ef8a23e4917
--- /dev/null
+++ b/tests/test_litellm/test_router_batch_retrieve_model_group.py
@@ -0,0 +1,118 @@
+"""
+model_group attribution on router batch retrieval.
+
+Batch token usage is accounted on the *retrieve* call, not on create: the
+provider only knows the token counts once the job finishes, so
+`LiteLLMBatch.usage` arrives on `aretrieve_batch` and that is the record the
+spend log tokens land on.
+
+`aretrieve_batch` is addressed by batch_id, so the request carries no model,
+and the router fans the lookup out across its deployments. These tests lock
+that the winning deployment's model group is stamped on the emitted
+StandardLoggingPayload, so `/global/activity/model` - which groups the spend
+logs by `model_group` - can attribute those tokens instead of bucketing every
+batch under "".
+"""
+
+import asyncio
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import litellm
+import litellm.batches.main as bm
+from litellm import Router
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.types.utils import LiteLLMBatch, Usage
+
+MODEL_GROUP = "vertex-gemini-2.5-flash-lite-dev"
+DEPLOYMENT_MODEL = "vertex_ai/gemini-2.5-flash-lite"
+
+
+class _PayloadCollector(CustomLogger):
+    def __init__(self):
+        super().__init__()
+        self.payloads = []
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self.payloads.append(kwargs.get("standard_logging_object"))
+
+
+@pytest.fixture
+def router():
+    return Router(
+        model_list=[
+            {
+                "model_name": MODEL_GROUP,
+                "litellm_params": {
+                    "model": DEPLOYMENT_MODEL,
+                    "vertex_project": "fake-project",
+                    "vertex_location": "us-central1",
+                    "vertex_credentials": "fake-creds",
+                },
+            }
+        ]
+    )
+
+
+@pytest.fixture
+def collector():
+    logger = _PayloadCollector()
+    previous = litellm.callbacks
+    litellm.callbacks = [logger]
+    try:
+        yield logger
+    finally:
+        litellm.callbacks = previous
+
+
+@pytest.fixture
+def vertex_retrieve():
+    """Mock the vertex provider seam - the only real network boundary."""
+    batch = LiteLLMBatch(
+        id="batch-1",
+        completion_window="24h",
+        created_at=0,
+        endpoint="/v1/chat/completions",
+        input_file_id="file-1",
+        object="batch",
+        status="completed",
+        usage=Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200),
+    )
+    seam = MagicMock(name="vertex_ai_batches_instance")
+    seam.retrieve_batch.return_value = batch
+    with patch.object(bm, "vertex_ai_batches_instance", seam):
+        yield seam
+
+
+async def _collected_payload(collector) -> dict:
+    for _ in range(50):  # the success handler runs as a background task
+        payloads = [p for p in collector.payloads if p is not None]
+        if payloads:
+            return payloads[-1]
+        await asyncio.sleep(0.05)
+    raise AssertionError(f"no StandardLoggingPayload was emitted: {collector.payloads}")
+
+
+@pytest.mark.asyncio
+async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, vertex_retrieve):
+    """
+    The proxy retrieves a managed batch by id only - no `model` in the request.
+    The router fans out over its deployments, so the model group is only known
+    from the deployment that answered.
+    """
+    response = await router.aretrieve_batch(batch_id="batch-1")
+
+    assert response.usage.total_tokens == 1200
+    payload = await _collected_payload(collector)
+    assert payload["model"] == DEPLOYMENT_MODEL
+    assert payload["model_group"] == MODEL_GROUP
+
+
+@pytest.mark.asyncio
+async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, vertex_retrieve):
+    """An explicitly requested model group is what gets logged."""
+    await router.aretrieve_batch(model=MODEL_GROUP, batch_id="batch-1")
+
+    payload = await _collected_payload(collector)
+    assert payload["model_group"] == MODEL_GROUP

From e630f21d16b10b78e22c28a974dee73009749167 Mon Sep 17 00:00:00 2001
From: mynkyu <mynkyu@bucketplace.net>
Date: Thu, 27 Aug 2026 19:01:18 +0900
Subject: [PATCH 02/14] test: fake the provider at the HTTP boundary in the
 batch model_group test

The test-quality gate flagged the first version for patching an SDK internal
(litellm.batches.main.vertex_ai_batches_instance) and for writing
litellm.callbacks directly.

Drive an openai-compatible deployment through respx instead, so the retrieve
call and the usage accounting that reads the completed batch's output file both
run for real, and install the collector with monkeypatch so nothing leaks into
the next test.

Signed-off-by: mynkyu <mynkyu@bucketplace.net>
---
 .../test_router_batch_retrieve_model_group.py | 146 +++++++++++-------
 1 file changed, 89 insertions(+), 57 deletions(-)

diff --git a/tests/test_litellm/test_router_batch_retrieve_model_group.py b/tests/test_litellm/test_router_batch_retrieve_model_group.py
index ef8a23e4917..b99ec50e041 100644
--- a/tests/test_litellm/test_router_batch_retrieve_model_group.py
+++ b/tests/test_litellm/test_router_batch_retrieve_model_group.py
@@ -1,35 +1,81 @@
 """
 model_group attribution on router batch retrieval.
 
-Batch token usage is accounted on the *retrieve* call, not on create: the
-provider only knows the token counts once the job finishes, so
-`LiteLLMBatch.usage` arrives on `aretrieve_batch` and that is the record the
-spend log tokens land on.
+Batch token usage is accounted on the *retrieve* call, not on create: a provider
+only reports token counts once the job finishes, so the usage is read off the
+completed batch's output file during retrieve logging and that is the spend log
+row the tokens land on.
 
-`aretrieve_batch` is addressed by batch_id, so the request carries no model,
-and the router fans the lookup out across its deployments. These tests lock
-that the winning deployment's model group is stamped on the emitted
-StandardLoggingPayload, so `/global/activity/model` - which groups the spend
-logs by `model_group` - can attribute those tokens instead of bucketing every
-batch under "".
+A batch is retrieved by id, so the request carries no model and the router fans
+the lookup out across its deployments. These tests lock that the answering
+deployment's model group is stamped on the emitted StandardLoggingPayload, so
+`/global/activity/model` - which groups the spend logs by `model_group` - can
+attribute those tokens instead of bucketing every batch under "".
+
+The provider is faked at the HTTP boundary, so the whole retrieve + usage
+accounting path runs for real.
 """
 
 import asyncio
-from unittest.mock import MagicMock, patch
+import json
 
+import httpx
 import pytest
+import respx
 
 import litellm
-import litellm.batches.main as bm
 from litellm import Router
 from litellm.integrations.custom_logger import CustomLogger
-from litellm.types.utils import LiteLLMBatch, Usage
 
-MODEL_GROUP = "vertex-gemini-2.5-flash-lite-dev"
-DEPLOYMENT_MODEL = "vertex_ai/gemini-2.5-flash-lite"
+MODEL_GROUP = "gemini-batch-group"
+DEPLOYMENT_MODEL = "openai/gpt-4o-mini"
+API_BASE = "http://localhost:4001/v1"
+BATCH_ID = "batch-1"
+ROWS = 2
+TOKENS_PER_ROW = 600
+
+COMPLETED_BATCH = {
+    "id": BATCH_ID,
+    "object": "batch",
+    "endpoint": "/v1/chat/completions",
+    "errors": None,
+    "input_file_id": "file-in-1",
+    "completion_window": "24h",
+    "status": "completed",
+    "output_file_id": "file-out-1",
+    "error_file_id": None,
+    "created_at": 0,
+    "completed_at": 1,
+    "request_counts": {"total": ROWS, "completed": ROWS, "failed": 0},
+    "metadata": None,
+}
+
+OUTPUT_JSONL = "\n".join(
+    json.dumps(
+        {
+            "id": f"req-{row}",
+            "custom_id": f"row-{row}",
+            "response": {
+                "status_code": 200,
+                "body": {
+                    "id": f"chatcmpl-{row}",
+                    "object": "chat.completion",
+                    "model": "gpt-4o-mini",
+                    "choices": [
+                        {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}
+                    ],
+                    "usage": {"prompt_tokens": 500, "completion_tokens": 100, "total_tokens": TOKENS_PER_ROW},
+                },
+            },
+        }
+    )
+    for row in range(ROWS)
+)
 
 
 class _PayloadCollector(CustomLogger):
+    """Captures the StandardLoggingPayload the spend log is built from."""
+
     def __init__(self):
         super().__init__()
         self.payloads = []
@@ -37,6 +83,14 @@ class _PayloadCollector(CustomLogger):
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
         self.payloads.append(kwargs.get("standard_logging_object"))
 
+    async def retrieve_batch_payload(self) -> dict:
+        for _ in range(100):  # the success handler runs as a background task
+            for payload in self.payloads:
+                if payload and payload.get("call_type") == "aretrieve_batch":
+                    return payload
+            await asyncio.sleep(0.05)
+        raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}")
+
 
 @pytest.fixture
 def router():
@@ -46,9 +100,8 @@ def router():
                 "model_name": MODEL_GROUP,
                 "litellm_params": {
                     "model": DEPLOYMENT_MODEL,
-                    "vertex_project": "fake-project",
-                    "vertex_location": "us-central1",
-                    "vertex_credentials": "fake-creds",
+                    "api_base": API_BASE,
+                    "api_key": "sk-fake",
                 },
             }
         ]
@@ -56,63 +109,42 @@ def router():
 
 
 @pytest.fixture
-def collector():
+def collector(monkeypatch):
     logger = _PayloadCollector()
-    previous = litellm.callbacks
-    litellm.callbacks = [logger]
-    try:
-        yield logger
-    finally:
-        litellm.callbacks = previous
+    monkeypatch.setattr(litellm, "callbacks", [logger])
+    return logger
 
 
 @pytest.fixture
-def vertex_retrieve():
-    """Mock the vertex provider seam - the only real network boundary."""
-    batch = LiteLLMBatch(
-        id="batch-1",
-        completion_window="24h",
-        created_at=0,
-        endpoint="/v1/chat/completions",
-        input_file_id="file-1",
-        object="batch",
-        status="completed",
-        usage=Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200),
-    )
-    seam = MagicMock(name="vertex_ai_batches_instance")
-    seam.retrieve_batch.return_value = batch
-    with patch.object(bm, "vertex_ai_batches_instance", seam):
-        yield seam
-
-
-async def _collected_payload(collector) -> dict:
-    for _ in range(50):  # the success handler runs as a background task
-        payloads = [p for p in collector.payloads if p is not None]
-        if payloads:
-            return payloads[-1]
-        await asyncio.sleep(0.05)
-    raise AssertionError(f"no StandardLoggingPayload was emitted: {collector.payloads}")
+def provider():
+    """Fake the provider at the HTTP boundary: the completed batch plus the
+    output file the usage accounting reads."""
+    with respx.mock(assert_all_called=True) as respx_mock:
+        respx_mock.get(f"{API_BASE}/batches/{BATCH_ID}").mock(return_value=httpx.Response(200, json=COMPLETED_BATCH))
+        respx_mock.get(f"{API_BASE}/files/file-out-1/content").mock(return_value=httpx.Response(200, text=OUTPUT_JSONL))
+        yield respx_mock
 
 
 @pytest.mark.asyncio
-async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, vertex_retrieve):
+async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, provider):
     """
     The proxy retrieves a managed batch by id only - no `model` in the request.
     The router fans out over its deployments, so the model group is only known
     from the deployment that answered.
     """
-    response = await router.aretrieve_batch(batch_id="batch-1")
+    response = await router.aretrieve_batch(batch_id=BATCH_ID)
 
-    assert response.usage.total_tokens == 1200
-    payload = await _collected_payload(collector)
+    assert response.id == BATCH_ID
+    payload = await collector.retrieve_batch_payload()
+    assert payload["total_tokens"] == ROWS * TOKENS_PER_ROW
     assert payload["model"] == DEPLOYMENT_MODEL
     assert payload["model_group"] == MODEL_GROUP
 
 
 @pytest.mark.asyncio
-async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, vertex_retrieve):
+async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, provider):
     """An explicitly requested model group is what gets logged."""
-    await router.aretrieve_batch(model=MODEL_GROUP, batch_id="batch-1")
+    await router.aretrieve_batch(model=MODEL_GROUP, batch_id=BATCH_ID)
 
-    payload = await _collected_payload(collector)
+    payload = await collector.retrieve_batch_payload()
     assert payload["model_group"] == MODEL_GROUP

From df6990c7127a0c30c77b96323e8311d9200a23be Mon Sep 17 00:00:00 2001
From: mynkyu <mynkyu@bucketplace.net>
Date: Sun, 6 Sep 2026 10:10:07 +0900
Subject: [PATCH 03/14] test: move the batch model_group regression into
 test_router.py

CLAUDE.md asks bug fixes to extend the existing mapped test file rather than
add a new one, and tests/test_litellm/test_router.py already covers
Router.aretrieve_batch. Fold the two cases in next to that coverage and drop
the standalone file.

The helpers are prefixed so they read unambiguously in a shared file, and the
respx context stays open while the payload is awaited, since the usage
accounting reads the batch's output file from the success handler.

Signed-off-by: mynkyu <mynkyu@bucketplace.net>
---
 tests/test_litellm/test_router.py             | 153 ++++++++++++++++++
 .../test_router_batch_retrieve_model_group.py | 150 -----------------
 2 files changed, 153 insertions(+), 150 deletions(-)
 delete mode 100644 tests/test_litellm/test_router_batch_retrieve_model_group.py

diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 7c044310e14..9b146092927 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -958,6 +958,159 @@ async def test_arouter_aretrieve_batch():
         assert mock_aretrieve_batch.call_args.kwargs["api_base"] == "my-custom-base"
 
 
+# ---------------------------------------------------------------------------
+# Batch retrieval has to attribute its tokens to a model group.
+#
+# Batch token usage is accounted on the *retrieve* call, not on create: a
+# provider only reports token counts once the job finishes, so the usage is read
+# off the completed batch's output file during retrieve logging, and that is the
+# spend log row the tokens land on. A batch is retrieved by id, so the request
+# carries no model and the router fans the lookup out across its deployments -
+# the group of the deployment that answered is the only one there is to stamp.
+# Leaving it unset files every batch's tokens under an empty model_group, which
+# is what /global/activity/model groups the spend logs by.
+#
+# The provider is faked at the HTTP boundary, so the retrieve call and the usage
+# accounting that reads the output file both run for real.
+# ---------------------------------------------------------------------------
+
+_BATCH_GROUP = "gemini-batch-group"
+_BATCH_DEPLOYMENT_MODEL = "openai/gpt-4o-mini"
+_BATCH_API_BASE = "http://localhost:4001/v1"
+_BATCH_ID = "batch-1"
+_BATCH_ROWS = 2
+_BATCH_TOKENS_PER_ROW = 600
+
+_BATCH_COMPLETED = {
+    "id": _BATCH_ID,
+    "object": "batch",
+    "endpoint": "/v1/chat/completions",
+    "errors": None,
+    "input_file_id": "file-in-1",
+    "completion_window": "24h",
+    "status": "completed",
+    "output_file_id": "file-out-1",
+    "error_file_id": None,
+    "created_at": 0,
+    "completed_at": 1,
+    "request_counts": {"total": _BATCH_ROWS, "completed": _BATCH_ROWS, "failed": 0},
+    "metadata": None,
+}
+
+_BATCH_OUTPUT_JSONL = "\n".join(
+    json.dumps(
+        {
+            "id": f"req-{row}",
+            "custom_id": f"row-{row}",
+            "response": {
+                "status_code": 200,
+                "body": {
+                    "id": f"chatcmpl-{row}",
+                    "object": "chat.completion",
+                    "model": "gpt-4o-mini",
+                    "choices": [
+                        {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}
+                    ],
+                    "usage": {
+                        "prompt_tokens": 500,
+                        "completion_tokens": 100,
+                        "total_tokens": _BATCH_TOKENS_PER_ROW,
+                    },
+                },
+            },
+        }
+    )
+    for row in range(_BATCH_ROWS)
+)
+
+
+class _BatchPayloadCollector(CustomLogger):
+    """Captures the StandardLoggingPayload the spend log row is built from."""
+
+    def __init__(self):
+        super().__init__()
+        self.payloads = []
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self.payloads.append(kwargs.get("standard_logging_object"))
+
+    async def retrieve_batch_payload(self):
+        for _ in range(100):  # the success handler runs as a background task
+            for payload in self.payloads:
+                if payload and payload.get("call_type") == "aretrieve_batch":
+                    return payload
+            await asyncio.sleep(0.05)
+        raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}")
+
+
+def _batch_model_group_router():
+    return litellm.Router(
+        model_list=[
+            {
+                "model_name": _BATCH_GROUP,
+                "litellm_params": {
+                    "model": _BATCH_DEPLOYMENT_MODEL,
+                    "api_base": _BATCH_API_BASE,
+                    "api_key": "sk-fake",
+                },
+            }
+        ]
+    )
+
+
+def _mock_batch_provider(respx_mock):
+    """The completed batch, plus the output file the usage accounting reads."""
+    respx_mock.get(f"{_BATCH_API_BASE}/batches/{_BATCH_ID}").mock(
+        return_value=httpx.Response(200, json=_BATCH_COMPLETED)
+    )
+    respx_mock.get(f"{_BATCH_API_BASE}/files/file-out-1/content").mock(
+        return_value=httpx.Response(200, text=_BATCH_OUTPUT_JSONL)
+    )
+
+
+@pytest.mark.asyncio
+async def test_arouter_aretrieve_batch_without_model_stamps_model_group(monkeypatch: pytest.MonkeyPatch):
+    """
+    The proxy retrieves a managed batch by id only - no `model` in the request.
+    The router fans out over its deployments, so the model group is only known
+    from the deployment that answered.
+    """
+    import respx
+
+    collector = _BatchPayloadCollector()
+    monkeypatch.setattr(litellm, "callbacks", [collector])
+    router = _batch_model_group_router()
+
+    with respx.mock(assert_all_called=True) as respx_mock:
+        _mock_batch_provider(respx_mock)
+        response = await router.aretrieve_batch(batch_id=_BATCH_ID)
+        # the usage accounting reads the output file from the success handler,
+        # so the provider has to stay faked until that payload lands
+        payload = await collector.retrieve_batch_payload()
+
+    assert response.id == _BATCH_ID
+    assert payload["total_tokens"] == _BATCH_ROWS * _BATCH_TOKENS_PER_ROW
+    assert payload["model"] == _BATCH_DEPLOYMENT_MODEL
+    assert payload["model_group"] == _BATCH_GROUP
+
+
+@pytest.mark.asyncio
+async def test_arouter_aretrieve_batch_with_model_stamps_requested_model_group(monkeypatch: pytest.MonkeyPatch):
+    """An explicitly requested model group is what gets logged."""
+    import respx
+
+    collector = _BatchPayloadCollector()
+    monkeypatch.setattr(litellm, "callbacks", [collector])
+    router = _batch_model_group_router()
+
+    with respx.mock(assert_all_called=True) as respx_mock:
+        _mock_batch_provider(respx_mock)
+        await router.aretrieve_batch(model=_BATCH_GROUP, batch_id=_BATCH_ID)
+        payload = await collector.retrieve_batch_payload()
+
+    assert payload["model_group"] == _BATCH_GROUP
+
+
 @pytest.mark.asyncio
 async def test_arouter_aretrieve_file_content():
     """
diff --git a/tests/test_litellm/test_router_batch_retrieve_model_group.py b/tests/test_litellm/test_router_batch_retrieve_model_group.py
deleted file mode 100644
index b99ec50e041..00000000000
--- a/tests/test_litellm/test_router_batch_retrieve_model_group.py
+++ /dev/null
@@ -1,150 +0,0 @@
-"""
-model_group attribution on router batch retrieval.
-
-Batch token usage is accounted on the *retrieve* call, not on create: a provider
-only reports token counts once the job finishes, so the usage is read off the
-completed batch's output file during retrieve logging and that is the spend log
-row the tokens land on.
-
-A batch is retrieved by id, so the request carries no model and the router fans
-the lookup out across its deployments. These tests lock that the answering
-deployment's model group is stamped on the emitted StandardLoggingPayload, so
-`/global/activity/model` - which groups the spend logs by `model_group` - can
-attribute those tokens instead of bucketing every batch under "".
-
-The provider is faked at the HTTP boundary, so the whole retrieve + usage
-accounting path runs for real.
-"""
-
-import asyncio
-import json
-
-import httpx
-import pytest
-import respx
-
-import litellm
-from litellm import Router
-from litellm.integrations.custom_logger import CustomLogger
-
-MODEL_GROUP = "gemini-batch-group"
-DEPLOYMENT_MODEL = "openai/gpt-4o-mini"
-API_BASE = "http://localhost:4001/v1"
-BATCH_ID = "batch-1"
-ROWS = 2
-TOKENS_PER_ROW = 600
-
-COMPLETED_BATCH = {
-    "id": BATCH_ID,
-    "object": "batch",
-    "endpoint": "/v1/chat/completions",
-    "errors": None,
-    "input_file_id": "file-in-1",
-    "completion_window": "24h",
-    "status": "completed",
-    "output_file_id": "file-out-1",
-    "error_file_id": None,
-    "created_at": 0,
-    "completed_at": 1,
-    "request_counts": {"total": ROWS, "completed": ROWS, "failed": 0},
-    "metadata": None,
-}
-
-OUTPUT_JSONL = "\n".join(
-    json.dumps(
-        {
-            "id": f"req-{row}",
-            "custom_id": f"row-{row}",
-            "response": {
-                "status_code": 200,
-                "body": {
-                    "id": f"chatcmpl-{row}",
-                    "object": "chat.completion",
-                    "model": "gpt-4o-mini",
-                    "choices": [
-                        {"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}
-                    ],
-                    "usage": {"prompt_tokens": 500, "completion_tokens": 100, "total_tokens": TOKENS_PER_ROW},
-                },
-            },
-        }
-    )
-    for row in range(ROWS)
-)
-
-
-class _PayloadCollector(CustomLogger):
-    """Captures the StandardLoggingPayload the spend log is built from."""
-
-    def __init__(self):
-        super().__init__()
-        self.payloads = []
-
-    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
-        self.payloads.append(kwargs.get("standard_logging_object"))
-
-    async def retrieve_batch_payload(self) -> dict:
-        for _ in range(100):  # the success handler runs as a background task
-            for payload in self.payloads:
-                if payload and payload.get("call_type") == "aretrieve_batch":
-                    return payload
-            await asyncio.sleep(0.05)
-        raise AssertionError(f"no aretrieve_batch payload was emitted: {self.payloads}")
-
-
-@pytest.fixture
-def router():
-    return Router(
-        model_list=[
-            {
-                "model_name": MODEL_GROUP,
-                "litellm_params": {
-                    "model": DEPLOYMENT_MODEL,
-                    "api_base": API_BASE,
-                    "api_key": "sk-fake",
-                },
-            }
-        ]
-    )
-
-
-@pytest.fixture
-def collector(monkeypatch):
-    logger = _PayloadCollector()
-    monkeypatch.setattr(litellm, "callbacks", [logger])
-    return logger
-
-
-@pytest.fixture
-def provider():
-    """Fake the provider at the HTTP boundary: the completed batch plus the
-    output file the usage accounting reads."""
-    with respx.mock(assert_all_called=True) as respx_mock:
-        respx_mock.get(f"{API_BASE}/batches/{BATCH_ID}").mock(return_value=httpx.Response(200, json=COMPLETED_BATCH))
-        respx_mock.get(f"{API_BASE}/files/file-out-1/content").mock(return_value=httpx.Response(200, text=OUTPUT_JSONL))
-        yield respx_mock
-
-
-@pytest.mark.asyncio
-async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, provider):
-    """
-    The proxy retrieves a managed batch by id only - no `model` in the request.
-    The router fans out over its deployments, so the model group is only known
-    from the deployment that answered.
-    """
-    response = await router.aretrieve_batch(batch_id=BATCH_ID)
-
-    assert response.id == BATCH_ID
-    payload = await collector.retrieve_batch_payload()
-    assert payload["total_tokens"] == ROWS * TOKENS_PER_ROW
-    assert payload["model"] == DEPLOYMENT_MODEL
-    assert payload["model_group"] == MODEL_GROUP
-
-
-@pytest.mark.asyncio
-async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, provider):
-    """An explicitly requested model group is what gets logged."""
-    await router.aretrieve_batch(model=MODEL_GROUP, batch_id=BATCH_ID)
-
-    payload = await collector.retrieve_batch_payload()
-    assert payload["model_group"] == MODEL_GROUP

From 128cb114bdac6b8cf41a9d689f0a573a2e27eced Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 20:45:47 -0700
Subject: [PATCH 04/14] style: trim comments on batch retrieve model group
 stamp

---
 litellm/router.py                 |  7 ++-----
 tests/test_litellm/test_router.py | 16 ----------------
 2 files changed, 2 insertions(+), 21 deletions(-)

diff --git a/litellm/router.py b/litellm/router.py
index c6c25b6be17..20c9018abee 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -6175,11 +6175,8 @@ class Router:
                         kwargs=new_kwargs,
                         function_name="aretrieve_batch",
                     )
-                    ## STAMP THE MODEL GROUP FOR SPEND TRACKING ##
-                    # A batch is retrieved by id, so the request carries no model group of its
-                    # own - only the deployment that answered knows it. Batch token usage lands
-                    # on this retrieve call (the provider reports counts once the job finishes),
-                    # so without this the tokens are logged under an empty model_group.
+                    # A batch is retrieved by id, so only the deployment that answered knows the
+                    # group, and batch token usage is logged on this retrieve call.
                     model_group: Final = requested_model_group or model_name["model_name"]
                     new_kwargs[metadata_variable_name].setdefault("model_group", model_group)
                     new_kwargs.pop("custom_llm_provider", None)
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 9b146092927..258d1c973d0 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -958,22 +958,6 @@ async def test_arouter_aretrieve_batch():
         assert mock_aretrieve_batch.call_args.kwargs["api_base"] == "my-custom-base"
 
 
-# ---------------------------------------------------------------------------
-# Batch retrieval has to attribute its tokens to a model group.
-#
-# Batch token usage is accounted on the *retrieve* call, not on create: a
-# provider only reports token counts once the job finishes, so the usage is read
-# off the completed batch's output file during retrieve logging, and that is the
-# spend log row the tokens land on. A batch is retrieved by id, so the request
-# carries no model and the router fans the lookup out across its deployments -
-# the group of the deployment that answered is the only one there is to stamp.
-# Leaving it unset files every batch's tokens under an empty model_group, which
-# is what /global/activity/model groups the spend logs by.
-#
-# The provider is faked at the HTTP boundary, so the retrieve call and the usage
-# accounting that reads the output file both run for real.
-# ---------------------------------------------------------------------------
-
 _BATCH_GROUP = "gemini-batch-group"
 _BATCH_DEPLOYMENT_MODEL = "openai/gpt-4o-mini"
 _BATCH_API_BASE = "http://localhost:4001/v1"

From 33d89c9814641a34cb66d357e2cc3a403677a06d Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 20:55:11 -0700
Subject: [PATCH 05/14] style: drop redundant comments per repo comment policy

---
 litellm/router.py                 | 3 +--
 tests/test_litellm/test_router.py | 3 ---
 2 files changed, 1 insertion(+), 5 deletions(-)

diff --git a/litellm/router.py b/litellm/router.py
index 20c9018abee..e85ca1bd7a8 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -6175,8 +6175,7 @@ class Router:
                         kwargs=new_kwargs,
                         function_name="aretrieve_batch",
                     )
-                    # A batch is retrieved by id, so only the deployment that answered knows the
-                    # group, and batch token usage is logged on this retrieve call.
+                    # Batch token usage is logged on this retrieve call, not on create.
                     model_group: Final = requested_model_group or model_name["model_name"]
                     new_kwargs[metadata_variable_name].setdefault("model_group", model_group)
                     new_kwargs.pop("custom_llm_provider", None)
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 258d1c973d0..dc51339cf55 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -1009,8 +1009,6 @@ _BATCH_OUTPUT_JSONL = "\n".join(
 
 
 class _BatchPayloadCollector(CustomLogger):
-    """Captures the StandardLoggingPayload the spend log row is built from."""
-
     def __init__(self):
         super().__init__()
         self.payloads = []
@@ -1043,7 +1041,6 @@ def _batch_model_group_router():
 
 
 def _mock_batch_provider(respx_mock):
-    """The completed batch, plus the output file the usage accounting reads."""
     respx_mock.get(f"{_BATCH_API_BASE}/batches/{_BATCH_ID}").mock(
         return_value=httpx.Response(200, json=_BATCH_COMPLETED)
     )

From 01bdfb34aa5ed88320dbd1c9f231876df820ca29 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 22:24:20 -0700
Subject: [PATCH 06/14] chore: drop redundant comments in aretrieve_batch
 router tests

---
 tests/test_litellm/test_router.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index dc51339cf55..bfa9ea7e3d1 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -1017,7 +1017,7 @@ class _BatchPayloadCollector(CustomLogger):
         self.payloads.append(kwargs.get("standard_logging_object"))
 
     async def retrieve_batch_payload(self):
-        for _ in range(100):  # the success handler runs as a background task
+        for _ in range(100):
             for payload in self.payloads:
                 if payload and payload.get("call_type") == "aretrieve_batch":
                     return payload
@@ -1065,8 +1065,6 @@ async def test_arouter_aretrieve_batch_without_model_stamps_model_group(monkeypa
     with respx.mock(assert_all_called=True) as respx_mock:
         _mock_batch_provider(respx_mock)
         response = await router.aretrieve_batch(batch_id=_BATCH_ID)
-        # the usage accounting reads the output file from the success handler,
-        # so the provider has to stay faked until that payload lands
         payload = await collector.retrieve_batch_payload()
 
     assert response.id == _BATCH_ID

From 9acf09f60d4f917053e1bfb9d493dce3cdd2771a Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 23:07:55 -0700
Subject: [PATCH 07/14] fix(router): keep batch retrieves out of the per-minute
 tpm/rpm counters

Stamping model_group let both router deployment callbacks past their
`model_group is None` early return for batch retrieves. A batch reports the
whole job's token total on retrieve and reports it again on every poll of the
finished batch, so those tokens are not load in the current minute: three polls
of one completed 1,200 token batch pushed a tpm:1000 deployment to 3,600. The
fan-out also probed unrelated deployments, adding an rpm tick to each.
---
 litellm/router.py                   |  5 ++
 litellm/router_utils/batch_utils.py | 18 +++++++
 tests/test_litellm/test_router.py   | 76 +++++++++++++++++++++++++++++
 3 files changed, 99 insertions(+)

diff --git a/litellm/router.py b/litellm/router.py
index e85ca1bd7a8..9dd267d7560 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -124,6 +124,7 @@ from litellm.router_utils.auto_router_model_naming import (
 )
 from litellm.router_utils.batch_utils import (
     _get_router_metadata_variable_name,
+    is_batch_retrieve_call_type,
     replace_model_in_jsonl,
     should_replace_model_in_jsonl,
 )
@@ -7878,6 +7879,8 @@ class Router:
             # WS session wrappers fire with result=None; per-turn costs tracked by inner calls.
             if kwargs.get("call_type") in ("_aresponses_websocket", "_arealtime"):
                 return
+            if is_batch_retrieve_call_type(kwargs.get("call_type")):
+                return
             standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
             if standard_logging_object is None:
                 raise ValueError("standard_logging_object is None")
@@ -8117,6 +8120,8 @@ class Router:
         """
         Update RPM usage for a deployment
         """
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         deployment_name: Final = kwargs["litellm_params"]["metadata"].get(
             "deployment", None
         )  # handles wildcard routes - by giving the original name sent to `litellm.completion`
diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py
index ccb6ad95519..6e110b586fb 100644
--- a/litellm/router_utils/batch_utils.py
+++ b/litellm/router_utils/batch_utils.py
@@ -5,6 +5,7 @@ from typing import Final
 
 from litellm._logging import verbose_logger
 from litellm.types.llms.openai import FileTypes, OpenAIFilesPurpose
+from litellm.types.utils import CallTypes
 
 
 class InMemoryFile(io.BytesIO):
@@ -170,3 +171,20 @@ def _get_router_metadata_variable_name(function_name: str | None) -> str:
         return "litellm_metadata"
     else:
         return "metadata"
+
+
+BATCH_RETRIEVE_CALL_TYPES: Final = frozenset(
+    {
+        CallTypes.aretrieve_batch.value,
+        CallTypes.retrieve_batch.value,
+    }
+)
+
+
+def is_batch_retrieve_call_type(call_type: object) -> bool:
+    """
+    A batch retrieve reports the whole job's token usage, which the provider spent
+    asynchronously over the life of the batch, and reports it again on every poll of the
+    finished batch. Per-minute usage counters must not be fed from it.
+    """
+    return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index bfa9ea7e3d1..d8045722998 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -1090,6 +1090,82 @@ async def test_arouter_aretrieve_batch_with_model_stamps_requested_model_group(m
     assert payload["model_group"] == _BATCH_GROUP
 
 
+_UNRELATED_BATCH_GROUP = "unrelated-batch-group"
+_UNRELATED_BATCH_API_BASE = "http://localhost:4002/v1"
+
+_BATCH_NOT_FOUND = {
+    "error": {
+        "message": f"No batch found with id '{_BATCH_ID}'.",
+        "type": "invalid_request_error",
+        "code": "batch_not_found",
+    }
+}
+
+
+async def _router_usage_keys(router, timeout: float = 2.0) -> list[str]:
+    loop = asyncio.get_event_loop()
+    deadline = loop.time() + timeout
+    while loop.time() < deadline:
+        keys = sorted(k for k in router.cache.in_memory_cache.cache_dict if k.startswith("global_router:"))
+        if keys:
+            return keys
+        await asyncio.sleep(0.05)
+    return []
+
+
+@pytest.mark.asyncio
+async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(monkeypatch: pytest.MonkeyPatch):
+    """
+    A batch reports the whole job's tokens on retrieve, and reports them again on every
+    poll of the finished batch, so they are not a measure of load in the current minute.
+    The fan-out also probes deployments the caller never named. Neither may reach the
+    per-minute tpm/rpm counters that gate live traffic.
+    """
+    import respx
+
+    collector = _BatchPayloadCollector()
+    monkeypatch.setattr(litellm, "callbacks", [collector])
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": _BATCH_GROUP,
+                "litellm_params": {
+                    "model": _BATCH_DEPLOYMENT_MODEL,
+                    "api_base": _BATCH_API_BASE,
+                    "api_key": "sk-fake",
+                },
+                "model_info": {"id": "batch-dep"},
+                "tpm": 1000,
+                "rpm": 10,
+            },
+            {
+                "model_name": _UNRELATED_BATCH_GROUP,
+                "litellm_params": {
+                    "model": _BATCH_DEPLOYMENT_MODEL,
+                    "api_base": _UNRELATED_BATCH_API_BASE,
+                    "api_key": "sk-fake",
+                },
+                "model_info": {"id": "unrelated-dep"},
+                "tpm": 1000,
+                "rpm": 10,
+            },
+        ]
+    )
+
+    with respx.mock(assert_all_called=True) as respx_mock:
+        _mock_batch_provider(respx_mock)
+        respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock(
+            return_value=httpx.Response(404, json=_BATCH_NOT_FOUND)
+        )
+        response = await router.aretrieve_batch(batch_id=_BATCH_ID)
+        payload = await collector.retrieve_batch_payload()
+        usage_keys = await _router_usage_keys(router)
+
+    assert response.id == _BATCH_ID
+    assert payload["model_group"] == _BATCH_GROUP
+    assert usage_keys == []
+
+
 @pytest.mark.asyncio
 async def test_arouter_aretrieve_file_content():
     """

From 58c3d04733f2bebfbc15e8f1f6dd702a37c6e2f6 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 23:24:47 -0700
Subject: [PATCH 08/14] test: cover is_batch_retrieve_call_type in router batch
 utils

---
 .../router_unit_tests/test_router_batch_utils.py  | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py
index c9f19731372..e274ac61a01 100644
--- a/tests/router_unit_tests/test_router_batch_utils.py
+++ b/tests/router_unit_tests/test_router_batch_utils.py
@@ -317,3 +317,18 @@ def test_replace_model_in_jsonl_with_embedded_newlines():
         == "This is a message\nwith multiple\nlines"
     )
     assert result_json["custom_id"] == "test123"
+
+
+def test_is_batch_retrieve_call_type_matches_only_batch_retrieves():
+    from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
+    from litellm.types.utils import CallTypes
+
+    assert is_batch_retrieve_call_type(CallTypes.aretrieve_batch.value) is True
+    assert is_batch_retrieve_call_type(CallTypes.retrieve_batch.value) is True
+
+    for call_type in CallTypes:
+        if call_type in (CallTypes.aretrieve_batch, CallTypes.retrieve_batch):
+            continue
+        assert is_batch_retrieve_call_type(call_type.value) is False
+
+    assert is_batch_retrieve_call_type(None) is False

From ad2afe5e6568b69389c2258680274098b9191b6d Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 5 Sep 2026 23:41:07 -0700
Subject: [PATCH 09/14] style(router): drop the inline comment on the batch
 retrieve stamp

---
 litellm/router.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/litellm/router.py b/litellm/router.py
index 9dd267d7560..2397c6b2fc7 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -6176,7 +6176,6 @@ class Router:
                         kwargs=new_kwargs,
                         function_name="aretrieve_batch",
                     )
-                    # Batch token usage is logged on this retrieve call, not on create.
                     model_group: Final = requested_model_group or model_name["model_name"]
                     new_kwargs[metadata_variable_name].setdefault("model_group", model_group)
                     new_kwargs.pop("custom_llm_provider", None)

From 63ea19743373fa1dd87081a66d64e5f580212f77 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 6 Sep 2026 01:02:53 -0700
Subject: [PATCH 10/14] fix(router): keep batch retrieves out of routing
 strategy state

Stamping model_group on a batch retrieve routed the whole batch job's token usage
into the per-model-group counters that usage-based, latency-based, cost-based and
least-busy routing read, so polling a finished batch could exhaust a group's TPM or
RPM window and lock live chat traffic out with RouterRateLimitError. Polling also
drove the least-busy in-flight counts negative once per poll per deployment, which
pinned chat to whichever deployment had been polled most.

The strategy callbacks now skip batch retrieve call types, so a retrieve still lands
in spend logs under its model group while the numbers that pick a deployment for the
next chat request stay driven by live traffic only.
---
 litellm/router.py                         |  3 +-
 litellm/router_strategy/least_busy.py     | 11 ++++
 litellm/router_strategy/lowest_cost.py    |  5 ++
 litellm/router_strategy/lowest_latency.py |  7 ++
 litellm/router_strategy/lowest_tpm_rpm.py |  5 ++
 litellm/router_utils/batch_utils.py       |  3 +-
 tests/test_litellm/test_router.py         | 79 +++++++++++++++++++++++
 7 files changed, 111 insertions(+), 2 deletions(-)

diff --git a/litellm/router.py b/litellm/router.py
index 2397c6b2fc7..8f913d96463 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -6177,7 +6177,8 @@ class Router:
                         function_name="aretrieve_batch",
                     )
                     model_group: Final = requested_model_group or model_name["model_name"]
-                    new_kwargs[metadata_variable_name].setdefault("model_group", model_group)
+                    if not new_kwargs[metadata_variable_name].get("model_group"):
+                        new_kwargs[metadata_variable_name]["model_group"] = model_group
                     new_kwargs.pop("custom_llm_provider", None)
                     data.pop("custom_llm_provider", None)
                     return await litellm.aretrieve_batch(
diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py
index 1433e8ba4d4..e93288fd9fe 100644
--- a/litellm/router_strategy/least_busy.py
+++ b/litellm/router_strategy/least_busy.py
@@ -11,6 +11,7 @@ from typing import Final
 
 from litellm.caching.caching import DualCache
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
 
 
 class LeastBusyLoggingHandler(CustomLogger):
@@ -27,6 +28,8 @@ class LeastBusyLoggingHandler(CustomLogger):
 
         Caching based on model group.
         """
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             if kwargs["litellm_params"].get("metadata") is None:
                 pass
@@ -48,6 +51,8 @@ class LeastBusyLoggingHandler(CustomLogger):
             pass
 
     def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             if kwargs["litellm_params"].get("metadata") is None:
                 pass
@@ -76,6 +81,8 @@ class LeastBusyLoggingHandler(CustomLogger):
             pass
 
     def log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             if kwargs["litellm_params"].get("metadata") is None:
                 pass
@@ -103,6 +110,8 @@ class LeastBusyLoggingHandler(CustomLogger):
             pass
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             if kwargs["litellm_params"].get("metadata") is None:
                 pass
@@ -131,6 +140,8 @@ class LeastBusyLoggingHandler(CustomLogger):
             pass
 
     async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             if kwargs["litellm_params"].get("metadata") is None:
                 pass
diff --git a/litellm/router_strategy/lowest_cost.py b/litellm/router_strategy/lowest_cost.py
index b927df0c438..aaad6186484 100644
--- a/litellm/router_strategy/lowest_cost.py
+++ b/litellm/router_strategy/lowest_cost.py
@@ -8,6 +8,7 @@ from litellm import ModelResponse, token_counter, verbose_logger
 from litellm._logging import verbose_router_logger
 from litellm.caching.caching import DualCache
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
 
 
 class LowestCostLoggingHandler(CustomLogger):
@@ -19,6 +20,8 @@ class LowestCostLoggingHandler(CustomLogger):
         self.router_cache = router_cache
 
     def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             """
             Update usage on success
@@ -96,6 +99,8 @@ class LowestCostLoggingHandler(CustomLogger):
             )
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             """
             Update cost usage on success
diff --git a/litellm/router_strategy/lowest_latency.py b/litellm/router_strategy/lowest_latency.py
index a1b67eaeaf9..598ca1227ec 100644
--- a/litellm/router_strategy/lowest_latency.py
+++ b/litellm/router_strategy/lowest_latency.py
@@ -9,6 +9,7 @@ from litellm import ModelResponse, token_counter, verbose_logger
 from litellm.caching.caching import DualCache
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs, safe_divide_seconds
+from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
 from litellm.types.utils import LiteLLMPydanticObjectBase
 
 if TYPE_CHECKING:
@@ -35,6 +36,8 @@ class LowestLatencyLoggingHandler(CustomLogger):
         self.routing_args = RoutingArgs(**routing_args)
 
     def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             """
             Update latency usage on success
@@ -167,6 +170,8 @@ class LowestLatencyLoggingHandler(CustomLogger):
         """
         Check if Timeout Error, if timeout set deployment latency -> 100
         """
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             metadata_field: Final = self._select_metadata_field(kwargs)
             _exception: Final = kwargs.get("exception", None)
@@ -221,6 +226,8 @@ class LowestLatencyLoggingHandler(CustomLogger):
             )
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             """
             Update latency usage on success
diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py
index 31c4b1d7e3f..d4abf1f8f70 100644
--- a/litellm/router_strategy/lowest_tpm_rpm.py
+++ b/litellm/router_strategy/lowest_tpm_rpm.py
@@ -8,6 +8,7 @@ from litellm import token_counter
 from litellm._logging import verbose_router_logger
 from litellm.caching.caching import DualCache
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
 from litellm.types.utils import LiteLLMPydanticObjectBase
 from litellm.utils import print_verbose
 
@@ -27,6 +28,8 @@ class LowestTPMLoggingHandler(CustomLogger):
         self.routing_args = RoutingArgs(**routing_args)
 
     def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             """
             Update TPM/RPM usage on success
@@ -79,6 +82,8 @@ class LowestTPMLoggingHandler(CustomLogger):
             verbose_router_logger.debug(traceback.format_exc())
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             """
             Update TPM/RPM usage on success
diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py
index 6e110b586fb..be20c358202 100644
--- a/litellm/router_utils/batch_utils.py
+++ b/litellm/router_utils/batch_utils.py
@@ -185,6 +185,7 @@ def is_batch_retrieve_call_type(call_type: object) -> bool:
     """
     A batch retrieve reports the whole job's token usage, which the provider spent
     asynchronously over the life of the batch, and reports it again on every poll of the
-    finished batch. Per-minute usage counters must not be fed from it.
+    finished batch. The counters that measure live traffic, per-minute rate limits and the
+    routing strategies' own state, must not be fed from it.
     """
     return isinstance(call_type, str) and call_type in BATCH_RETRIEVE_CALL_TYPES
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index d8045722998..836085c1f5d 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -1166,6 +1166,85 @@ async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(m
     assert usage_keys == []
 
 
+_ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:")
+
+
+async def _router_strategy_keys(router, timeout: float = 2.0) -> list[str]:
+    loop = asyncio.get_event_loop()
+    deadline = loop.time() + timeout
+    while loop.time() < deadline:
+        keys = sorted(
+            key
+            for key in router.cache.in_memory_cache.cache_dict
+            if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS)
+        )
+        if keys:
+            return keys
+        await asyncio.sleep(0.05)
+    return []
+
+
+def _batch_fan_out_router(routing_strategy: str):
+    return litellm.Router(
+        routing_strategy=routing_strategy,
+        model_list=[
+            {
+                "model_name": _BATCH_GROUP,
+                "litellm_params": {
+                    "model": _BATCH_DEPLOYMENT_MODEL,
+                    "api_base": _BATCH_API_BASE,
+                    "api_key": "sk-fake",
+                },
+                "model_info": {"id": "batch-dep"},
+            },
+            {
+                "model_name": _UNRELATED_BATCH_GROUP,
+                "litellm_params": {
+                    "model": _BATCH_DEPLOYMENT_MODEL,
+                    "api_base": _UNRELATED_BATCH_API_BASE,
+                    "api_key": "sk-fake",
+                },
+                "model_info": {"id": "unrelated-dep"},
+            },
+        ],
+    )
+
+
+@pytest.mark.parametrize(
+    "routing_strategy",
+    ["usage-based-routing", "latency-based-routing", "cost-based-routing", "least-busy"],
+)
+@pytest.mark.asyncio
+async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies(
+    monkeypatch: pytest.MonkeyPatch, routing_strategy: str
+):
+    """
+    Every routing strategy picks a deployment from what recent live traffic did.
+    A batch retrieve reports the whole job on every poll and probes deployments the
+    caller never named, so polling a finished batch must not move the numbers that
+    decide where the next chat request goes.
+    """
+    import respx
+
+    collector = _BatchPayloadCollector()
+    monkeypatch.setattr(litellm, "callbacks", [collector])
+    monkeypatch.setattr(litellm, "input_callback", [])
+    router = _batch_fan_out_router(routing_strategy)
+
+    with respx.mock(assert_all_called=True) as respx_mock:
+        _mock_batch_provider(respx_mock)
+        respx_mock.get(f"{_UNRELATED_BATCH_API_BASE}/batches/{_BATCH_ID}").mock(
+            return_value=httpx.Response(404, json=_BATCH_NOT_FOUND)
+        )
+        for _ in range(3):
+            response = await router.aretrieve_batch(batch_id=_BATCH_ID)
+        await collector.retrieve_batch_payload()
+        strategy_keys = await _router_strategy_keys(router)
+
+    assert response.id == _BATCH_ID
+    assert strategy_keys == []
+
+
 @pytest.mark.asyncio
 async def test_arouter_aretrieve_file_content():
     """

From 828a02f78f2c974f6458237c6fde1128b01b26bb Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sun, 6 Sep 2026 03:03:45 -0700
Subject: [PATCH 11/14] fix(batches): stamp the model group on the proxy's
 model-encoded retrieve path

The model-encoded batch id path calls the SDK directly, so the router never
labels it. Stamp the decoded group into the request's litellm_metadata, and
guard usage-based-routing-v2 the same way the other strategies already are.
---
 litellm/proxy/batches_endpoints/endpoints.py  | 21 +++++++---
 litellm/router_strategy/lowest_tpm_rpm_v2.py  |  5 +++
 .../proxy/batches_endpoints/test_endpoints.py | 40 ++++++++++++++++++-
 tests/test_litellm/test_router.py             | 26 +++++++-----
 4 files changed, 76 insertions(+), 16 deletions(-)

diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py
index 5c4bacd757c..b6a8b421e02 100644
--- a/litellm/proxy/batches_endpoints/endpoints.py
+++ b/litellm/proxy/batches_endpoints/endpoints.py
@@ -52,6 +52,20 @@ from litellm.types.llms.openai import LiteLLMBatchCreateRequest
 router: Final = APIRouter()
 
 
+def _litellm_metadata_of(data: dict) -> dict:
+    """The request's litellm_metadata mapping, created on the request when it carries none.
+
+    The success handler reads this mapping, so a flag or a model group set here has to live
+    inside it rather than beside it.
+    """
+    existing: Final = data.get("litellm_metadata")
+    if isinstance(existing, dict):
+        return existing
+    created: Final = {}  # mutable-ok: the logging layer copies and extends this mapping, so it cannot be a read-only view
+    data["litellm_metadata"] = created
+    return created
+
+
 def _raise_not_found_when_openai_fallback_unservable(
     requested_provider: "str | None",
     data: Mapping[str, object],
@@ -531,11 +545,7 @@ async def retrieve_batch(
 
         poller_owns_accounting: Final = bool(unified_batch_id) and batch_cost_poller_is_active()
         if poller_owns_accounting:
-            litellm_metadata = data.get("litellm_metadata")
-            if not isinstance(litellm_metadata, dict):
-                litellm_metadata = {}  # mutable-ok: the suppression flag must live inside litellm_metadata for the success handler to read it, and this request carried no mapping to extend
-                data["litellm_metadata"] = litellm_metadata
-            litellm_metadata["batch_ignore_default_logging"] = True
+            _litellm_metadata_of(data)["batch_ignore_default_logging"] = True
 
         # Retrieve from provider (for non-terminal states or if DB lookup failed)
         # SCENARIO 1: Batch ID is encoded with model info
@@ -558,6 +568,7 @@ async def retrieve_batch(
             # so litellm.aretrieve_batch can load BedrockBatchesConfig. Without
             # it the call falls into the legacy provider switch and 400s.
             data["model"] = model_from_id
+            _litellm_metadata_of(data).setdefault("model_group", model_from_id)
 
             # Retrieve batch using model credentials
             response = await litellm.aretrieve_batch(
diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py
index 665ff69ab47..a2acce5fcb5 100644
--- a/litellm/router_strategy/lowest_tpm_rpm_v2.py
+++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py
@@ -12,6 +12,7 @@ from litellm._logging import verbose_logger, verbose_router_logger
 from litellm.caching.caching import DualCache
 from litellm.integrations.custom_logger import CustomLogger
 from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs
+from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
 from litellm.types.router import RouterErrors
 from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload
 from litellm.utils import get_utc_datetime, print_verbose
@@ -210,6 +211,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger):
             return deployment  # don't fail calls if eg. redis fails to connect
 
     def log_success_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             """
             Update TPM/RPM usage on success
@@ -250,6 +253,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger):
             )
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return
         try:
             """
             Update TPM usage on success
diff --git a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
index a37c8ff2bb4..c6df8f2ffcf 100644
--- a/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
+++ b/tests/test_litellm/proxy/batches_endpoints/test_endpoints.py
@@ -1233,9 +1233,11 @@ async def call_retrieve(
     user: Optional[UserAPIKeyAuth] = None,
     headers: Optional[Dict[str, str]] = None,
     query: Optional[Dict[str, str]] = None,
+    enriched_data: Optional[Dict[str, Any]] = None,
 ):
-    # Mirror the real flow: data starts as RetrieveBatchRequest(batch_id=...).
-    harness.data["data"] = {"batch_id": batch_id}
+    # Mirror the real flow: data starts as RetrieveBatchRequest(batch_id=...),
+    # then pre-call enrichment adds key/team metadata to it.
+    harness.data["data"] = {"batch_id": batch_id, **(enriched_data or {})}
     return await endpoints.retrieve_batch(
         request=FakeRequest(headers=headers, query=query),
         fastapi_response=Response(),
@@ -1271,6 +1273,7 @@ async def test_retrieve__model_encoded_id(retrieve_harness):
         "api_key": "sk-azure",
         "api_base": "https://azure.test",
         "model": "azure/gpt-4o",
+        "litellm_metadata": {"model_group": "azure/gpt-4o"},
     }
 
     # 4. OUTPUT SHAPE - ids re-encoded with the model for the round-trip.
@@ -1293,6 +1296,39 @@ async def test_retrieve__model_encoded_id__forwards_decoded_model_not_deployment
     assert retrieve_harness.aretrieve_kwargs()["model"] == "azure/gpt-4o"
 
 
+@pytest.mark.asyncio
+async def test_retrieve__model_encoded_id__stamps_model_group(retrieve_harness):
+    """This path never goes through the router, so nothing else labels the call.
+    Without the stamp the spend log lands under a blank model group and the batch
+    disappears from per-model usage."""
+    await call_retrieve(retrieve_harness, AZURE_BATCH_ID)
+
+    litellm_metadata = retrieve_harness.aretrieve_kwargs()["litellm_metadata"]
+
+    assert litellm_metadata["model_group"] == "azure/gpt-4o"
+
+
+@pytest.mark.asyncio
+async def test_retrieve__model_encoded_id__stamps_model_group_beside_existing_metadata(
+    retrieve_harness,
+):
+    """The stamp joins the metadata pre-call enrichment already built. Replacing
+    that dict instead of adding to it drops the key and team labels the spend log
+    is attributed with."""
+    await call_retrieve(
+        retrieve_harness,
+        AZURE_BATCH_ID,
+        enriched_data={"litellm_metadata": {"user_api_key_alias": "team-a-key"}},
+    )
+
+    litellm_metadata = retrieve_harness.aretrieve_kwargs()["litellm_metadata"]
+
+    assert litellm_metadata == {
+        "user_api_key_alias": "team-a-key",
+        "model_group": "azure/gpt-4o",
+    }
+
+
 @pytest.mark.asyncio
 async def test_retrieve__model_encoded_id__encodes_output_and_error_ids(
     retrieve_harness,
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 836085c1f5d..e4e0dafd750 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -1169,17 +1169,19 @@ async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(m
 _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:")
 
 
-async def _router_strategy_keys(router, timeout: float = 2.0) -> list[str]:
+async def _moved_routing_counters(router, timeout: float = 2.0) -> list[str]:
     loop = asyncio.get_event_loop()
     deadline = loop.time() + timeout
     while loop.time() < deadline:
-        keys = sorted(
-            key
-            for key in router.cache.in_memory_cache.cache_dict
+        cache_dict = router.cache.in_memory_cache.cache_dict
+        moved = sorted(
+            f"{key}={cache_dict[key]}"
+            for key in cache_dict
             if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS)
+            and cache_dict[key]
         )
-        if keys:
-            return keys
+        if moved:
+            return moved
         await asyncio.sleep(0.05)
     return []
 
@@ -1212,7 +1214,13 @@ def _batch_fan_out_router(routing_strategy: str):
 
 @pytest.mark.parametrize(
     "routing_strategy",
-    ["usage-based-routing", "latency-based-routing", "cost-based-routing", "least-busy"],
+    [
+        "usage-based-routing",
+        "usage-based-routing-v2",
+        "latency-based-routing",
+        "cost-based-routing",
+        "least-busy",
+    ],
 )
 @pytest.mark.asyncio
 async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies(
@@ -1239,10 +1247,10 @@ async def test_arouter_aretrieve_batch_does_not_feed_routing_strategies(
         for _ in range(3):
             response = await router.aretrieve_batch(batch_id=_BATCH_ID)
         await collector.retrieve_batch_payload()
-        strategy_keys = await _router_strategy_keys(router)
+        moved_counters = await _moved_routing_counters(router)
 
     assert response.id == _BATCH_ID
-    assert strategy_keys == []
+    assert moved_counters == []
 
 
 @pytest.mark.asyncio

From 2d13ca06d0386d5b27daadb110d8c79c9d9fcf33 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 02:39:38 -0700
Subject: [PATCH 12/14] fix(router): keep batch retrieves out of the sync
 success counter and type the metadata helper

---
 litellm/proxy/batches_endpoints/endpoints.py | 10 +++---
 litellm/router.py                            |  2 ++
 tests/test_litellm/test_router.py            | 38 ++++++++++++++++++++
 3 files changed, 45 insertions(+), 5 deletions(-)

diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py
index 1c38748fc90..403669d7f97 100644
--- a/litellm/proxy/batches_endpoints/endpoints.py
+++ b/litellm/proxy/batches_endpoints/endpoints.py
@@ -6,7 +6,7 @@
 ######################################################################
 import asyncio
 import os
-from collections.abc import Mapping
+from collections.abc import Mapping, MutableMapping
 from types import MappingProxyType
 from typing import Any, Final, cast
 
@@ -57,17 +57,17 @@ from litellm.types.llms.openai import LiteLLMBatchCreateRequest
 router: Final = APIRouter()
 
 
-def _litellm_metadata_of(data: dict) -> dict:
+def _litellm_metadata_of(data: MutableMapping[str, object]) -> MutableMapping[str, object]:
     """The request's litellm_metadata mapping, created on the request when it carries none.
 
     The success handler reads this mapping, so a flag or a model group set here has to live
     inside it rather than beside it.
     """
     existing: Final = data.get("litellm_metadata")
-    if isinstance(existing, dict):
+    if isinstance(existing, MutableMapping):
         return existing
-    created: Final = {}  # mutable-ok: the logging layer copies and extends this mapping, so it cannot be a read-only view
-    data["litellm_metadata"] = created
+    created: Final[dict[str, object]] = {}  # mutable-ok: the logging layer copies and extends this mapping
+    data["litellm_metadata"] = created  # rebind-ok: the success handler reads the request's own mapping
     return created
 
 
diff --git a/litellm/router.py b/litellm/router.py
index cf6f637332c..4c6c736935f 100644
--- a/litellm/router.py
+++ b/litellm/router.py
@@ -8097,6 +8097,8 @@ class Router:
         - key: str - The key used to increment the cache
         - None: if no key is found
         """
+        if is_batch_retrieve_call_type(kwargs.get("call_type")):
+            return None
         id = None
         if kwargs["litellm_params"].get("metadata") is None:
             pass
diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py
index 193c007a2e1..79b6635ce6f 100644
--- a/tests/test_litellm/test_router.py
+++ b/tests/test_litellm/test_router.py
@@ -49,6 +49,7 @@ from litellm.router import (
 from litellm.router_strategy import simple_shuffle
 from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit
 from litellm.router_utils.cooldown_handlers import _async_get_cooldown_deployments
+from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute
 from litellm.types.llms.openai import ChatCompletionRequest
 from litellm.types.router import Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo, PreRoutingHookResponse, RetryPolicy
 
@@ -1217,6 +1218,43 @@ async def test_arouter_aretrieve_batch_does_not_consume_deployment_rate_limits(m
     assert usage_keys == []
 
 
+@pytest.mark.parametrize(
+    ("call_type", "expected_key", "expected_successes"),
+    [
+        ("aretrieve_batch", None, 0),
+        ("retrieve_batch", None, 0),
+        ("acompletion", "batch-dep:successes", 1),
+    ],
+)
+def test_sync_deployment_callback_on_success_skips_batch_retrieves(
+    call_type: str, expected_key: str | None, expected_successes: int
+):
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": _BATCH_GROUP,
+                "litellm_params": {"model": _BATCH_DEPLOYMENT_MODEL, "api_base": _BATCH_API_BASE, "api_key": "sk-fake"},
+                "model_info": {"id": "batch-dep"},
+            }
+        ]
+    )
+
+    key = router.sync_deployment_callback_on_success(
+        kwargs={
+            "call_type": call_type,
+            "litellm_params": {"metadata": {"model_group": _BATCH_GROUP}, "model_info": {"id": "batch-dep"}},
+        },
+        completion_response=None,
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+    )
+
+    assert key == expected_key
+    assert (
+        get_deployment_successes_for_current_minute(litellm_router_instance=router, deployment_id="batch-dep")
+        == expected_successes
+    )
+
 _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:")
 
 

From 8a33b37c39503419c50adad19d806a8f51fbe117 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 08:20:14 -0700
Subject: [PATCH 13/14] fix(proxy): charge a finished batch once against
 per-model budgets

A completed batch reports its whole cost on every retrieve, and the
per-model budget limiter added that cost to the key, user, team, and
end-user counters on each poll. Stamping model_group on plain-id
retrieves widened this from model-encoded batch ids to every poll, so
a key ran out of a budget it never spent. A marker per counter and
batch id now lets the first poll charge and later polls skip.
---
 .../proxy/hooks/model_max_budget_limiter.py   | 80 ++++++++++++++---
 .../hooks/test_model_max_budget_limiter.py    | 89 +++++++++++++++++++
 2 files changed, 155 insertions(+), 14 deletions(-)
 create mode 100644 tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py

diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py
index bbfc7325f40..67577a68f5c 100644
--- a/litellm/proxy/hooks/model_max_budget_limiter.py
+++ b/litellm/proxy/hooks/model_max_budget_limiter.py
@@ -5,6 +5,8 @@ from dataclasses import dataclass
 from types import MappingProxyType
 from typing import Final
 
+from openai.types import Batch
+
 import litellm
 from litellm._logging import verbose_proxy_logger
 from litellm.caching.caching import DualCache
@@ -13,6 +15,7 @@ from litellm.litellm_core_utils.duration_parser import duration_in_seconds
 from litellm.llms.bedrock.common_utils import get_bedrock_base_model
 from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth
 from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
+from litellm.router_utils.batch_utils import is_batch_retrieve_call_type
 from litellm.types.llms.openai import AllMessageValues
 from litellm.types.utils import BudgetConfig, StandardLoggingPayload
 
@@ -117,6 +120,17 @@ def model_budget_start_time_cache_key(
     return f"{_BUDGET_START_TIME_KEY_PREFIXES[entity_type]}:{entity_id}:{budget_model}:{budget_duration}"
 
 
+def batch_charged_once_marker_key(spend_key: str, batch_id: str) -> str:
+    return f"{spend_key}:batch:{batch_id}"
+
+
+def batch_id_to_charge_once(call_type: object, response_obj: object, response_cost: float) -> str | None:
+    """A finished batch reports its whole cost on every poll, so its id is charged once per counter."""
+    if response_cost <= 0 or not is_batch_retrieve_call_type(call_type):
+        return None
+    return response_obj.id if isinstance(response_obj, Batch) else None
+
+
 def resolve_model_budget(model: str, model_max_budget: Mapping[str, object]) -> ResolvedModelBudget | None:
     """Find the `model_max_budget` entry that governs `model`, or None."""
     for candidate in _budget_model_candidates(model):
@@ -537,22 +551,18 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
             )
             return
 
+        batch_id: Final = batch_id_to_charge_once(
+            call_type=kwargs.get("call_type"),
+            response_obj=response_obj,
+            response_cost=response_cost,
+        )
         for entity_type, entity_id, resolved in resolved_budgets:
-            await self._increment_spend_for_key(
-                budget_config=resolved.budget_config,
-                spend_key=model_budget_spend_cache_key(
-                    entity_type=entity_type,
-                    entity_id=entity_id,
-                    budget_model=resolved.budget_model,
-                    budget_duration=resolved.budget_config.budget_duration,
-                ),
-                start_time_key=model_budget_start_time_cache_key(
-                    entity_type=entity_type,
-                    entity_id=entity_id,
-                    budget_model=resolved.budget_model,
-                    budget_duration=resolved.budget_config.budget_duration,
-                ),
+            await self._charge_entity(
+                entity_type=entity_type,
+                entity_id=entity_id,
+                resolved=resolved,
                 response_cost=response_cost,
+                batch_id=batch_id,
             )
 
         if self.dual_cache.redis_cache is not None:
@@ -562,3 +572,45 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
             "current state of in memory cache %s",
             json.dumps(self.dual_cache.in_memory_cache.cache_dict, indent=4, default=str),
         )
+
+    async def _charge_entity(
+        self,
+        entity_type: Litellm_EntityType,
+        entity_id: str | None,
+        resolved: ResolvedModelBudget,
+        response_cost: float,
+        batch_id: str | None,
+    ) -> None:
+        budget_duration: Final = resolved.budget_config.budget_duration
+        if budget_duration is None:
+            return
+        spend_key: Final = model_budget_spend_cache_key(
+            entity_type=entity_type,
+            entity_id=entity_id,
+            budget_model=resolved.budget_model,
+            budget_duration=budget_duration,
+        )
+        if batch_id is not None and not await self._claim_batch_charge(
+            spend_key=spend_key,
+            batch_id=batch_id,
+            ttl_seconds=duration_in_seconds(budget_duration),
+        ):
+            return
+        await self._increment_spend_for_key(
+            budget_config=resolved.budget_config,
+            spend_key=spend_key,
+            start_time_key=model_budget_start_time_cache_key(
+                entity_type=entity_type,
+                entity_id=entity_id,
+                budget_model=resolved.budget_model,
+                budget_duration=budget_duration,
+            ),
+            response_cost=response_cost,
+        )
+
+    async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool:
+        marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id)
+        if await self.dual_cache.async_get_cache(key=marker_key) is not None:
+            return False
+        await self.dual_cache.async_set_cache(key=marker_key, value=1, ttl=ttl_seconds)
+        return True
diff --git a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py
new file mode 100644
index 00000000000..47b14438212
--- /dev/null
+++ b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py
@@ -0,0 +1,89 @@
+from typing import Final
+
+import pytest
+
+from litellm.caching.caching import DualCache
+from litellm.proxy.hooks.model_max_budget_limiter import (
+    _PROXY_VirtualKeyModelMaxBudgetLimiter,
+)
+from litellm.types.utils import LiteLLMBatch, Usage
+
+KEY_HASH: Final = "key-hash-batch"
+USER_ID: Final = "user-batch"
+MODEL_GROUP: Final = "batch-qa-primary"
+BATCH_COST: Final = 2.925e-05
+CHAT_COST: Final = 0.001
+KEY_SPEND_KEY: Final = f"virtual_key_spend:{KEY_HASH}:{MODEL_GROUP}:1d"
+USER_SPEND_KEY: Final = f"user_model_spend:{USER_ID}:{MODEL_GROUP}:1d"
+
+
+def _batch(batch_id: str, status: str) -> LiteLLMBatch:
+    return LiteLLMBatch(
+        id=batch_id,
+        completion_window="24h",
+        created_at=1,
+        endpoint="/v1/chat/completions",
+        input_file_id="file-batch",
+        object="batch",
+        status=status,
+        usage=Usage(prompt_tokens=20, completion_tokens=18, total_tokens=38),
+    )
+
+
+def _event(call_type: str, response_cost: float) -> dict:
+    return {
+        "call_type": call_type,
+        "standard_logging_object": {
+            "call_type": call_type,
+            "response_cost": response_cost,
+            "model": "openai/gpt-5.4-mini",
+            "model_group": MODEL_GROUP,
+            "metadata": {"user_api_key_hash": KEY_HASH, "user_api_key_user_id": USER_ID},
+        },
+        "litellm_params": {
+            "metadata": {
+                "user_api_key_model_max_budget": {MODEL_GROUP: {"budget_limit": 0.0001, "time_period": "1d"}},
+                "user_api_key_user_model_max_budget": {MODEL_GROUP: {"budget_limit": 0.0001, "time_period": "1d"}},
+            }
+        },
+    }
+
+
+async def _poll(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, batch: LiteLLMBatch, response_cost: float) -> None:
+    await limiter.async_log_success_event(
+        _event("aretrieve_batch", response_cost), response_obj=batch, start_time=None, end_time=None
+    )
+
+
+async def _chat(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter) -> None:
+    await limiter.async_log_success_event(_event("acompletion", CHAT_COST), response_obj=None, start_time=None, end_time=None)
+
+
+async def _spend(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, spend_key: str) -> float:
+    return await limiter.dual_cache.async_get_cache(key=spend_key) or 0.0
+
+
+@pytest.mark.asyncio
+async def test_polls_of_a_finished_batch_charge_each_per_model_budget_once():
+    limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
+    first: Final = _batch("batch_first", "completed")
+
+    await _poll(limiter, _batch("batch_first", "in_progress"), response_cost=0)
+    for _ in range(3):
+        await _poll(limiter, first, response_cost=BATCH_COST)
+
+    assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST)
+    assert await _spend(limiter, USER_SPEND_KEY) == pytest.approx(BATCH_COST)
+
+
+@pytest.mark.asyncio
+async def test_a_second_batch_and_chat_requests_still_charge_the_budget():
+    limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
+
+    await _poll(limiter, _batch("batch_first", "completed"), response_cost=BATCH_COST)
+    await _poll(limiter, _batch("batch_first", "completed"), response_cost=BATCH_COST)
+    await _poll(limiter, _batch("batch_second", "completed"), response_cost=BATCH_COST)
+    await _chat(limiter)
+    await _chat(limiter)
+
+    assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(2 * BATCH_COST + 2 * CHAT_COST)

From 867a4df347a2632f8d93af52a9154f1679545ca7 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Sat, 19 Sep 2026 17:41:19 -0700
Subject: [PATCH 14/14] fix(proxy): claim a finished batch's per-model budget
 charge atomically

A finished batch reports its whole cost on every poll. The charge-once
marker is now taken with one atomic increment on the shared cache, so two
workers polling the same batch at once cannot both charge it, and the
marker's TTL is refreshed on every poll so a batch polled within every
budget window is never charged again after the marker's first expiry.
---
 .../proxy/hooks/model_max_budget_limiter.py   |   8 +-
 .../hooks/test_model_max_budget_limiter.py    | 113 +++++++++++++++++-
 2 files changed, 115 insertions(+), 6 deletions(-)

diff --git a/litellm/proxy/hooks/model_max_budget_limiter.py b/litellm/proxy/hooks/model_max_budget_limiter.py
index 67577a68f5c..cfa54ae01a2 100644
--- a/litellm/proxy/hooks/model_max_budget_limiter.py
+++ b/litellm/proxy/hooks/model_max_budget_limiter.py
@@ -610,7 +610,7 @@ class _PROXY_VirtualKeyModelMaxBudgetLimiter(RouterBudgetLimiting):
 
     async def _claim_batch_charge(self, spend_key: str, batch_id: str, ttl_seconds: int) -> bool:
         marker_key: Final = batch_charged_once_marker_key(spend_key=spend_key, batch_id=batch_id)
-        if await self.dual_cache.async_get_cache(key=marker_key) is not None:
-            return False
-        await self.dual_cache.async_set_cache(key=marker_key, value=1, ttl=ttl_seconds)
-        return True
+        polls: Final = await self.dual_cache.async_increment_cache(
+            key=marker_key, value=1, ttl=ttl_seconds, refresh_ttl=True
+        )
+        return polls == 1
diff --git a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py
index 47b14438212..ffb60fb4651 100644
--- a/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py
+++ b/tests/test_litellm/proxy/hooks/test_model_max_budget_limiter.py
@@ -1,3 +1,7 @@
+import asyncio
+import time
+from collections.abc import Callable, Mapping
+from types import MappingProxyType
 from typing import Final
 
 import pytest
@@ -6,6 +10,7 @@ from litellm.caching.caching import DualCache
 from litellm.proxy.hooks.model_max_budget_limiter import (
     _PROXY_VirtualKeyModelMaxBudgetLimiter,
 )
+from litellm.types.caching import RedisPipelineIncrementOperation
 from litellm.types.utils import LiteLLMBatch, Usage
 
 KEY_HASH: Final = "key-hash-batch"
@@ -30,7 +35,7 @@ def _batch(batch_id: str, status: str) -> LiteLLMBatch:
     )
 
 
-def _event(call_type: str, response_cost: float) -> dict:
+def _event(call_type: str, response_cost: float) -> dict[str, object]:
     return {
         "call_type": call_type,
         "standard_logging_object": {
@@ -56,13 +61,88 @@ async def _poll(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, batch: LiteLLMB
 
 
 async def _chat(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter) -> None:
-    await limiter.async_log_success_event(_event("acompletion", CHAT_COST), response_obj=None, start_time=None, end_time=None)
+    await limiter.async_log_success_event(
+        _event("acompletion", CHAT_COST), response_obj=None, start_time=None, end_time=None
+    )
 
 
 async def _spend(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, spend_key: str) -> float:
     return await limiter.dual_cache.async_get_cache(key=spend_key) or 0.0
 
 
+def _local_spend(limiter: _PROXY_VirtualKeyModelMaxBudgetLimiter, spend_key: str) -> float:
+    return limiter.dual_cache.in_memory_cache.get_cache(key=spend_key) or 0.0
+
+
+class _Clock:
+    def __init__(self) -> None:
+        self.seconds = 0.0
+
+    def now(self) -> float:
+        return self.seconds
+
+    def advance(self, seconds: float) -> None:
+        self.seconds = self.seconds + seconds
+
+
+class _SharedRedisDouble:
+    def __init__(self, now: Callable[[], float] = time.time) -> None:
+        self.now = now
+        self.entries: Mapping[str, tuple[float, float | None]] = MappingProxyType({})
+
+    def _live(self, key: str) -> tuple[float, float | None] | None:
+        entry: Final = self.entries.get(key)
+        if entry is None:
+            return None
+        expires_at: Final = entry[1]
+        if expires_at is not None and expires_at <= self.now():
+            return None
+        return entry
+
+    def _store(self, key: str, value: float, expires_at: float | None) -> None:
+        self.entries = MappingProxyType({**self.entries, key: (value, expires_at)})
+
+    async def async_get_cache(self, key: str, **kwargs: object) -> float | None:
+        await asyncio.sleep(0)
+        entry: Final = self._live(key)
+        return None if entry is None else entry[0]
+
+    async def async_set_cache(self, key: str, value: float, ttl: int | None = None, **kwargs: object) -> None:
+        await asyncio.sleep(0)
+        self._store(key, value, None if ttl is None else self.now() + ttl)
+
+    async def async_increment(
+        self,
+        key: str,
+        value: float,
+        ttl: int | None = None,
+        parent_otel_span: object = None,
+        refresh_ttl: bool = False,
+    ) -> float:
+        await asyncio.sleep(0)
+        live: Final = self._live(key)
+        total: Final = value if live is None else live[0] + value
+        kept_expiry: Final = None if live is None else live[1]
+        expires_at: Final = (
+            kept_expiry if ttl is None or (kept_expiry is not None and not refresh_ttl) else self.now() + ttl
+        )
+        self._store(key, total, expires_at)
+        return total
+
+    async def async_increment_pipeline(self, increment_list: list[RedisPipelineIncrementOperation]) -> list[float]:
+        return [await self.async_increment(op["key"], op["increment_value"], ttl=op["ttl"]) for op in increment_list]
+
+
+def _worker(redis: _SharedRedisDouble) -> _PROXY_VirtualKeyModelMaxBudgetLimiter:
+    return _PROXY_VirtualKeyModelMaxBudgetLimiter(
+        dual_cache=DualCache(redis_cache=redis)  # pyright: ignore[reportArgumentType]  # duck-typed Redis double
+    )
+
+
+async def _drain_redis_pushes() -> None:
+    await asyncio.gather(*(task for task in asyncio.all_tasks() if task is not asyncio.current_task()))
+
+
 @pytest.mark.asyncio
 async def test_polls_of_a_finished_batch_charge_each_per_model_budget_once():
     limiter: Final = _PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
@@ -87,3 +167,32 @@ async def test_a_second_batch_and_chat_requests_still_charge_the_budget():
     await _chat(limiter)
 
     assert await _spend(limiter, KEY_SPEND_KEY) == pytest.approx(2 * BATCH_COST + 2 * CHAT_COST)
+
+
+@pytest.mark.asyncio
+async def test_two_workers_polling_the_same_finished_batch_at_once_charge_it_once():
+    redis: Final = _SharedRedisDouble()
+    worker_a: Final = _worker(redis)
+    worker_b: Final = _worker(redis)
+    finished: Final = _batch("batch_first", "completed")
+
+    await asyncio.gather(_poll(worker_a, finished, BATCH_COST), _poll(worker_b, finished, BATCH_COST))
+    await _drain_redis_pushes()
+
+    assert _local_spend(worker_a, KEY_SPEND_KEY) + _local_spend(worker_b, KEY_SPEND_KEY) == pytest.approx(BATCH_COST)
+    assert await redis.async_get_cache(KEY_SPEND_KEY) == pytest.approx(BATCH_COST)
+
+
+@pytest.mark.asyncio
+async def test_a_batch_polled_within_every_budget_window_is_never_charged_again():
+    clock: Final = _Clock()
+    limiter: Final = _worker(_SharedRedisDouble(now=clock.now))
+    finished: Final = _batch("batch_first", "completed")
+
+    await _poll(limiter, finished, BATCH_COST)
+    clock.advance(12 * 3600)
+    await _poll(limiter, finished, BATCH_COST)
+    clock.advance(18 * 3600)
+    await _poll(limiter, finished, BATCH_COST)
+
+    assert _local_spend(limiter, KEY_SPEND_KEY) == pytest.approx(BATCH_COST)