feat(batches): enrich batch cost rows with breakdown, identity, session, and org spend

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -112,35 +112,39 @@ class CheckBatchCost:
             verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
             return {}
 
-    async def _get_key_alias(self, batch_id: str, api_key: str | None) -> str | None:
-        """Resolve the creating virtual key's alias from its hashed token."""
+    async def _get_key_attribution(self, batch_id: str, api_key: str | None) -> tuple[str | None, str | None]:
+        """Resolve the creating virtual key's (alias, org_id) from its hashed token."""
         if not api_key:
-            return None
+            return None, None
         try:
             key_row: prisma_models.LiteLLM_VerificationToken | None = (
                 await self.prisma_client.db.litellm_verificationtoken.find_unique(
                     where={"token": api_key}
                 )
             )
-            return getattr(key_row, "key_alias", None) if key_row is not None else None
+            if key_row is None:
+                return None, None
+            return getattr(key_row, "key_alias", None), getattr(key_row, "org_id", None)
         except Exception as e:
             verbose_proxy_logger.error(f"CheckBatchCost: could not look up key alias for batch {batch_id}: {e}")
-            return None
+            return None, None
 
-    async def _get_team_alias(self, team_id: str | None) -> str | None:
-        """Resolve a team's alias from its id."""
+    async def _get_team_attribution(self, team_id: str | None) -> tuple[str | None, str | None]:
+        """Resolve a team's (alias, organization_id) from its id."""
         if not team_id:
-            return None
+            return None, None
         try:
             team_row: prisma_models.LiteLLM_TeamTable | None = (
                 await self.prisma_client.db.litellm_teamtable.find_unique(
                     where={"team_id": team_id}
                 )
             )
-            return getattr(team_row, "team_alias", None) if team_row is not None else None
+            if team_row is None:
+                return None, None
+            return getattr(team_row, "team_alias", None), getattr(team_row, "organization_id", None)
         except Exception as e:
             verbose_proxy_logger.error(f"CheckBatchCost: could not look up team alias for team {team_id}: {e}")
-            return None
+            return None, None
 
     async def _build_creator_attribution_metadata(
         self, job: "LiteLLM_ManagedObjectTable", batch_id: str
@@ -153,6 +157,10 @@ class CheckBatchCost:
         user_api_key_alias; when it has no alias, or the key has since been rotated or
         deleted, the field keeps the creating user's alias that _get_user_info filled in,
         because a resolvable name is more useful on the spend row than a null.
+
+        user_api_key_org_id must be resolved here too: the spend update writer reads it
+        off this metadata to increment organization spend, so leaving it out silently
+        drops batch cost from org accounting for keys and teams that belong to one.
         """
         api_key = getattr(job, "api_key", None)
         team_id = getattr(job, "team_id", None)
@@ -165,12 +173,15 @@ class CheckBatchCost:
             **(await self._get_user_info(batch_id, job.created_by)),
         }
 
-        key_alias = await self._get_key_alias(batch_id, api_key)
+        key_alias, key_org_id = await self._get_key_attribution(batch_id, api_key)
         if key_alias is not None:
             metadata["user_api_key_alias"] = key_alias
-        team_alias = await self._get_team_alias(team_id)
+        team_alias, team_org_id = await self._get_team_attribution(team_id)
         if team_alias is not None:
             metadata["user_api_key_team_alias"] = team_alias
+        org_id: Final = key_org_id or team_org_id
+        if org_id is not None:
+            metadata["user_api_key_org_id"] = org_id
         if isinstance(request_tags, list) and request_tags:
             metadata["tags"] = [tag for tag in request_tags if isinstance(tag, str)]
 
@@ -804,6 +815,7 @@ class CheckBatchCost:
             function_id=str(uuid.uuid4()),
         )
 
+        deployment_api_base: Final = deployment_info.litellm_params.api_base
         logging_obj.update_environment_variables(
             litellm_params={
                 # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks
@@ -812,9 +824,17 @@ class CheckBatchCost:
                         "user-agent": CHECK_BATCH_COST_USER_AGENT,
                     }
                 },
-                "metadata": await self._build_creator_attribution_metadata(job, batch_id),
+                **({"api_base": deployment_api_base} if deployment_api_base else {}),
+                "metadata": {
+                    **(await self._build_creator_attribution_metadata(job, batch_id)),
+                    # spend logs read the deployment identity off these metadata keys, so
+                    # without them the batch cost row carries no model_id or model_group
+                    "model_info": {"id": model_id},
+                    "model_group": deployment_info.model_name,
+                },
             },
             optional_params={},
+            custom_llm_provider=str(llm_provider) if llm_provider else None,
         )
 
         if not await self._claim_job_for_costing(job):
@@ -832,6 +852,8 @@ class CheckBatchCost:
                 batch_models=batch_result.models,
                 batch_successful_requests=batch_result.successful_requests,
                 batch_failed_requests=batch_result.failed_requests,
+                batch_prompt_cost=batch_result.prompt_cost,
+                batch_completion_cost=batch_result.completion_cost,
             )
         except Exception:
             await self._release_job_claim(job)

litellm/batches/batch_utils.py

@@ -10,7 +10,7 @@ from litellm._logging import verbose_logger
 from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
 from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
 from litellm.types.llms.openai import Batch
-from litellm.types.utils import CallTypes, ModelInfo, Usage
+from litellm.types.utils import ModelInfo, Usage
 from litellm.utils import token_counter
 
 
@@ -23,6 +23,8 @@ class BatchCostUsageResult:
     models: list[str]
     successful_requests: int
     failed_requests: int
+    prompt_cost: float = 0.0
+    completion_cost: float = 0.0
 
 
 async def calculate_batch_cost_and_usage(
@@ -130,7 +132,8 @@ class _LineOutcome(Enum):
 
 @dataclass(frozen=True, slots=True)
 class _BatchOutputLineStats:
-    cost: float
+    prompt_cost: float
+    completion_cost: float
     prompt_tokens: int
     completion_tokens: int
     total_tokens: int
@@ -193,15 +196,16 @@ def _compute_output_line_stats(
     raw_model: Final = response_body.get("model")
     response_model: Final = raw_model if isinstance(raw_model, str) and raw_model else None
     completion_details: Final = usage.completion_tokens_details
+    line_prompt_cost, line_completion_cost = _output_line_cost(
+        usage=usage,
+        custom_llm_provider=custom_llm_provider,
+        model_name=model_name,
+        response_model=response_model,
+        model_info=model_info,
+    )
     return _BatchOutputLineStats(
-        cost=_output_line_cost(
-            response_body=response_body,
-            usage=usage,
-            custom_llm_provider=custom_llm_provider,
-            model_name=model_name,
-            response_model=response_model,
-            model_info=model_info,
-        ),
+        prompt_cost=line_prompt_cost,
+        completion_cost=line_completion_cost,
         prompt_tokens=usage.prompt_tokens,
         completion_tokens=usage.completion_tokens,
         total_tokens=usage.total_tokens,
@@ -213,31 +217,24 @@ def _compute_output_line_stats(
 
 
 def _output_line_cost(
-    response_body: Mapping[str, Any],
     usage: Usage,
     custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
     model_name: str | None,
     response_model: str | None,
     model_info: ModelInfo | None,
-) -> float:
+) -> tuple[float, float]:
+    """(prompt_cost, completion_cost) for one output line, priced at batch rates."""
     from litellm.cost_calculator import batch_cost_calculator
 
-    if model_info is None and custom_llm_provider not in ("anthropic", "bedrock"):
-        return litellm.completion_cost(
-            completion_response=response_body,
-            custom_llm_provider=custom_llm_provider,
-            call_type=CallTypes.aretrieve_batch.value,
-        )
     cost_model: Final = (
         model_name if custom_llm_provider == "bedrock" and model_name else response_model or model_name or ""
     )
-    prompt_cost, completion_cost = batch_cost_calculator(
+    return batch_cost_calculator(
         usage=usage,
         model=cost_model,
         custom_llm_provider=custom_llm_provider,
         model_info=model_info,
     )
-    return prompt_cost + completion_cost
 
 
 def _aggregate_batch_cost_usage_models(
@@ -270,7 +267,9 @@ def _aggregate_batch_cost_usage_models(
         **cache_token_params,
     )
     batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model]
-    total_cost: Final = sum((stats.cost for stats in line_stats), 0.0)
+    total_prompt_cost: Final = sum((stats.prompt_cost for stats in line_stats), 0.0)
+    total_completion_cost: Final = sum((stats.completion_cost for stats in line_stats), 0.0)
+    total_cost: Final = total_prompt_cost + total_completion_cost
     verbose_logger.debug(
         "batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d",
         total_cost,
@@ -285,6 +284,8 @@ def _aggregate_batch_cost_usage_models(
         models=batch_models,
         successful_requests=successful_requests,
         failed_requests=failed_requests,
+        prompt_cost=total_prompt_cost,
+        completion_cost=total_completion_cost,
     )
 
 
@@ -309,7 +310,8 @@ def calculate_vertex_ai_batch_cost_and_usage(
     """
     from litellm.cost_calculator import batch_cost_calculator
 
-    total_cost = 0.0
+    total_prompt_cost = 0.0  # rebind-ok: loop accumulator, matches total_tokens below
+    total_completion_cost = 0.0  # rebind-ok: loop accumulator, matches total_tokens below
     total_tokens = 0
     prompt_tokens = 0
     completion_tokens = 0
@@ -341,7 +343,8 @@ def calculate_vertex_ai_batch_cost_and_usage(
                 model=actual_model_name,
                 custom_llm_provider="vertex_ai",
             )
-            total_cost += p_cost + c_cost
+            total_prompt_cost += p_cost
+            total_completion_cost += c_cost
         except Exception as e:
             verbose_logger.debug("vertex_ai batch cost calculation error for line: %s", str(e))
 
@@ -349,6 +352,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
         completion_tokens += _completion
         total_tokens += _total
 
+    total_cost: Final = total_prompt_cost + total_completion_cost
     verbose_logger.info(
         "vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d",
         total_cost,
@@ -369,6 +373,8 @@ def calculate_vertex_ai_batch_cost_and_usage(
         models=[actual_model_name],
         successful_requests=successful_requests,
         failed_requests=failed_requests,
+        prompt_cost=total_prompt_cost,
+        completion_cost=total_completion_cost,
     )
 
 

litellm/litellm_core_utils/litellm_logging.py

@@ -2904,6 +2904,15 @@ class Logging(LiteLLMLoggingBaseClass):
                 result._hidden_params["batch_successful_requests"] = batch_successful_requests  # pyright: ignore[reportPrivateUsage]  # rebind-ok: same result._hidden_params pattern as response_cost/batch_models above
                 result._hidden_params["batch_failed_requests"] = batch_failed_requests  # pyright: ignore[reportPrivateUsage]  # rebind-ok: same pattern as above
                 result.usage = batch_usage
+                batch_prompt_cost: Final = kwargs.get("batch_prompt_cost", None)
+                batch_completion_cost: Final = kwargs.get("batch_completion_cost", None)
+                if batch_prompt_cost is not None and batch_completion_cost is not None:
+                    self.set_cost_breakdown(
+                        input_cost=batch_prompt_cost,
+                        output_cost=batch_completion_cost,
+                        total_cost=batch_cost,
+                        cost_for_built_in_tools_cost_usd_dollar=0.0,
+                    )
 
             elif should_compute_batch_data:
                 batch_result: Final = await _handle_completed_batch(
@@ -2919,6 +2928,12 @@ class Logging(LiteLLMLoggingBaseClass):
                 result._hidden_params["batch_successful_requests"] = batch_result.successful_requests  # pyright: ignore[reportPrivateUsage]  # rebind-ok: same pattern as above
                 result._hidden_params["batch_failed_requests"] = batch_result.failed_requests  # pyright: ignore[reportPrivateUsage]  # rebind-ok: same pattern as above
                 result.usage = batch_result.usage
+                self.set_cost_breakdown(
+                    input_cost=batch_result.prompt_cost,
+                    output_cost=batch_result.completion_cost,
+                    total_cost=batch_result.cost,
+                    cost_for_built_in_tools_cost_usd_dollar=0.0,
+                )
 
         start_time, end_time, result = self._success_handler_helper_fn(
             start_time=start_time,

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -582,6 +582,7 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
                 metadata=metadata,
                 standard_logging_payload=standard_logging_payload,
                 omit_when_missing=_omits_session_id_when_missing(metadata),
+                batch_trace_session_id=_get_batch_trace_session_id(call_type=call_type, request_id=id),
             ),
             request_duration_ms=_get_request_duration_ms(start_time, end_time),
             status=_get_status_for_spend_log(
@@ -620,20 +621,44 @@ def _omits_session_id_when_missing(metadata: Mapping[str, object] | None) -> boo
     return general_settings.get("missing_session_id") == "omit"
 
 
+_BATCH_TRACE_CALL_TYPES: Final = frozenset(
+    {
+        CallTypes.create_batch.value,
+        CallTypes.acreate_batch.value,
+        CallTypes.retrieve_batch.value,
+        CallTypes.aretrieve_batch.value,
+    }
+)
+
+
+def _get_batch_trace_session_id(call_type: str | None, request_id: str | None) -> str | None:
+    """A batch's create row and its poller-written cost row both derive their request id
+    from the same batch id (the cost row appends BATCH_COST_REQUEST_ID_SUFFIX), so using
+    that id as the session groups the batch lifecycle into one trace on the logs UI. The
+    poller builds its own logging context, so per-request trace ids can never link them."""
+    if call_type not in _BATCH_TRACE_CALL_TYPES or not request_id:
+        return None
+    return request_id.removesuffix(BATCH_COST_REQUEST_ID_SUFFIX)
+
+
 def _get_session_id_for_spend_log(
     kwargs: Mapping[str, object],
     metadata: Mapping[str, object] | None,
     standard_logging_payload: StandardLoggingPayload | None,
     omit_when_missing: bool,
+    batch_trace_session_id: str | None = None,
 ) -> str | None:
     """Under `omit` only `metadata.session_id`, the key Langfuse reads, counts as a session; `litellm_session_id` may
-    be a copied trace id."""
+    be a copied trace id. Batch call types carry a deterministic session derived from the batch id, which outranks
+    the per-request trace ids because those differ between the create call and the cost poller's row."""
     if omit_when_missing:
         session_id: Final = metadata.get("session_id") if metadata else None
         return str(session_id) if session_id else None
 
     from litellm._uuid import uuid
 
+    if batch_trace_session_id is not None:
+        return batch_trace_session_id
     if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None:
         return str(standard_logging_payload.get("trace_id"))
     if kwargs.get("litellm_trace_id") is not None:

tests/batches_tests/test_batches_logging_unit_tests.py

@@ -144,18 +144,20 @@ def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict):
 @pytest.mark.asyncio
 async def test_batch_cost_calculator(sample_file_content_dict):
     """
-    mock litellm.completion_cost to return 0.5
+    mock batch_cost_calculator to return (0.3, 0.2) per line
 
     we know sample_file_content_dict has 2 successful responses
 
-    so we expect the cost to be 0.5 * 2 = 1.0
+    so we expect the cost to be (0.3 + 0.2) * 2 = 1.0, split 0.6 / 0.4
     """
-    with patch("litellm.completion_cost", return_value=0.5):
+    with patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.3, 0.2)):
         result = _aggregate_batch_cost_usage_models(
             entries=sample_file_content_dict,
             custom_llm_provider="openai",
         )
-        assert result.cost == 1.0  # 0.5 * 2 successful responses
+        assert result.cost == pytest.approx(1.0)  # (0.3 + 0.2) * 2 successful responses
+        assert result.prompt_cost == pytest.approx(0.6)
+        assert result.completion_cost == pytest.approx(0.4)
 
 
 def test_get_response_from_batch_job_output_file(sample_file_content_dict):
@@ -402,6 +404,56 @@ async def test_batch_retrieve_cost_tracking_with_explicit_cost_data():
         assert mock_batch.usage == explicit_usage
 
 
+@pytest.mark.asyncio
+async def test_batch_retrieve_explicit_cost_split_sets_cost_breakdown():
+    """The poller passes the batch's prompt/completion cost split so the spend row's
+    cost_breakdown carries real input/output costs; without it the UI's Cost Breakdown
+    card renders blank for every batch. Regression for the split being dropped."""
+    from litellm.litellm_core_utils.litellm_logging import Logging
+    from litellm.types.utils import CallTypes, LiteLLMBatch
+
+    mock_batch = LiteLLMBatch(
+        id="batch-breakdown-1",
+        object="batch",
+        endpoint="/v1/chat/completions",
+        errors=None,
+        input_file_id="file-input-1",
+        completion_window="24h",
+        status="completed",
+        output_file_id="file-output-1",
+        created_at=1234567890,
+    )
+    mock_batch._hidden_params = {}
+
+    logging_obj = Logging(
+        model="gpt-5-mini",
+        messages=[{"role": "user", "content": "test"}],
+        stream=False,
+        call_type=CallTypes.aretrieve_batch.value,
+        litellm_call_id="test-call-breakdown",
+        function_id="test-function",
+        start_time=time.time(),
+        dynamic_success_callbacks=[],
+    )
+    logging_obj.custom_llm_provider = "openai"
+
+    await logging_obj.async_success_handler(
+        result=mock_batch,
+        start_time=time.time(),
+        end_time=time.time() + 1,
+        batch_cost=0.10,
+        batch_usage=litellm.Usage(prompt_tokens=200, completion_tokens=100, total_tokens=300),
+        batch_models=["gpt-5-mini"],
+        batch_prompt_cost=0.06,
+        batch_completion_cost=0.04,
+    )
+
+    assert logging_obj.cost_breakdown is not None
+    assert logging_obj.cost_breakdown["input_cost"] == 0.06
+    assert logging_obj.cost_breakdown["output_cost"] == 0.04
+    assert logging_obj.cost_breakdown["total_cost"] == 0.10
+
+
 @pytest.mark.asyncio
 async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batch():
     """

tests/proxy_unit_tests/test_check_batch_cost.py

@@ -2553,6 +2553,51 @@ class TestBatchCostAttribution:
 
         assert metadata["user_api_key_alias"] == "prod-key"
 
+    @pytest.mark.asyncio
+    async def test_org_id_comes_from_the_creating_key(self):
+        """The spend update writer increments organization spend from user_api_key_org_id,
+        so an org-scoped key's batch cost must carry the key's org id."""
+        from types import SimpleNamespace
+
+        instance = self._instance(
+            key_row=SimpleNamespace(key_alias="prod-key", org_id="org-42"),
+            team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
+        )
+
+        metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
+
+        assert metadata["user_api_key_org_id"] == "org-42"
+
+    @pytest.mark.asyncio
+    async def test_org_id_falls_back_to_the_team_organization(self):
+        """A key with no org of its own still books batch spend against its team's
+        organization, matching how the request path resolves org attribution."""
+        from types import SimpleNamespace
+
+        instance = self._instance(
+            key_row=SimpleNamespace(key_alias="prod-key", org_id=None),
+            team_row=SimpleNamespace(team_alias="Team Alpha", organization_id="org-team"),
+        )
+
+        metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
+
+        assert metadata["user_api_key_org_id"] == "org-team"
+
+    @pytest.mark.asyncio
+    async def test_no_org_leaves_the_key_unset(self):
+        """Without any org the key is absent entirely, so the spend writer's org update
+        stays skipped instead of matching an empty-string organization."""
+        from types import SimpleNamespace
+
+        instance = self._instance(
+            key_row=SimpleNamespace(key_alias="prod-key", org_id=None),
+            team_row=SimpleNamespace(team_alias="Team Alpha", organization_id=None),
+        )
+
+        metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
+
+        assert "user_api_key_org_id" not in metadata
+
 
 class TestPollPageStarvation:
     """LIT-5462 regression: a row that can never be costed used to keep its slot in the

tests/test_litellm/batches/test_batch_utils.py

@@ -489,7 +489,9 @@ def test_aggregate_counts_successful_and_failed_requests(monkeypatch):
 
 
 def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch):
-    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 1.0)
+    import litellm.cost_calculator as cc
+
+    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.4, 0.6))
     result = bu._aggregate_batch_cost_usage_models(
         entries=[_success_row(usage=_usage(10, 5))], custom_llm_provider="openai"
     )
@@ -500,6 +502,7 @@ def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch):
         1,
         0,
     )
+    assert (result.prompt_cost, result.completion_cost) == (0.4, 0.6)
 
 
 # =========================================================================== #
@@ -507,15 +510,17 @@ def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch):
 # =========================================================================== #
 
 
-def test_cost_from_content_completion_cost_path(monkeypatch):
-    # model_info is None -> litellm.completion_cost per successful row.
+def test_cost_without_model_info_prices_each_row_by_its_response_model(monkeypatch):
+    # model_info is None -> batch_cost_calculator per successful row, model from the response body.
+    import litellm.cost_calculator as cc
+
     calls = []
 
-    def _completion_cost(**kw):
+    def _batch_cost(**kw):
         calls.append(kw)
-        return 0.5
+        return (0.3, 0.2)
 
-    monkeypatch.setattr(litellm, "completion_cost", _completion_cost)
+    monkeypatch.setattr(cc, "batch_cost_calculator", _batch_cost)
     rows = [
         _success_row(usage=_usage(10, 5)),
         _failed_row(),  # excluded -> not costed
@@ -524,8 +529,10 @@ def test_cost_from_content_completion_cost_path(monkeypatch):
 
     result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
 
-    assert result.cost == 1.0  # 2 successful * 0.5
+    assert result.cost == pytest.approx(1.0)  # 2 successful * (0.3 + 0.2)
+    assert (result.prompt_cost, result.completion_cost) == (pytest.approx(0.6), pytest.approx(0.4))
     assert len(calls) == 2  # failed row not costed
+    assert all(call["model"] == "gpt-4o" and call["model_info"] is None for call in calls)
     assert result.successful_requests == 2
     assert result.failed_requests == 1
 
@@ -578,7 +585,9 @@ def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch):
     """A one-shot generator: any implementation that iterates the entries twice
     (e.g. separate cost and usage passes) sees nothing on the second pass and
     returns wrong totals for at least one of cost/usage/models."""
-    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.5)
+    import litellm.cost_calculator as cc
+
+    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.25, 0.25))
     one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))])
 
     result = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai")
@@ -753,12 +762,15 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch):
 
 @pytest.mark.asyncio
 async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch):
+    import litellm.cost_calculator as cc
+
     rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
-    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5)
+    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (1.5, 1.0))
 
     result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai")
 
     assert result.cost == 2.5
+    assert (result.prompt_cost, result.completion_cost) == (1.5, 1.0)
     assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15)
     assert result.models == ["gpt-4o"]
 
@@ -1107,8 +1119,10 @@ async def test_handle_completed_batch_orchestration(monkeypatch):
     async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
         return _vertex_jsonl(rows)
 
+    import litellm.cost_calculator as cc
+
     monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
-    monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.3)
+    monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3))
 
     result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai")
 

tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py

@@ -128,6 +128,68 @@ def test_legacy_policy_keeps_trace_id_fallback():
     assert len(str(generated)) == 36
 
 
+def test_batch_lifecycle_rows_derive_the_same_session_from_the_batch_id():
+    """The create call's request id IS the batch id and the poller's cost row appends
+    _batch_cost to it, so deriving the session from the request id lands both rows in one
+    trace on the logs UI even though the poller builds a fresh logging context per cycle."""
+    from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id
+
+    create_session: Final = _get_batch_trace_session_id(call_type="acreate_batch", request_id="batch-uid-1")
+    cost_session: Final = _get_batch_trace_session_id(
+        call_type="aretrieve_batch", request_id="batch-uid-1_batch_cost"
+    )
+    assert create_session == cost_session == "batch-uid-1"
+
+
+def test_non_batch_call_types_derive_no_batch_session():
+    from litellm.proxy.spend_tracking.spend_tracking_utils import _get_batch_trace_session_id
+
+    assert _get_batch_trace_session_id(call_type="acompletion", request_id="chatcmpl-1") is None
+
+
+def test_batch_session_outranks_the_per_request_trace_id():
+    """Each batch lifecycle call carries its own auto-generated trace id, so letting the
+    trace id win would scatter the rows across sessions again."""
+    session_id: Final = _get_session_id_for_spend_log(
+        kwargs={"litellm_trace_id": "trace-abc"},
+        metadata={"trace_id": "trace-abc"},
+        standard_logging_payload=_TRACE_ONLY_STANDARD_LOGGING,
+        omit_when_missing=False,
+        batch_trace_session_id="batch-uid-1",
+    )
+    assert session_id == "batch-uid-1"
+
+
+def test_omit_policy_still_suppresses_batch_sessions():
+    session_id: Final = _get_session_id_for_spend_log(
+        kwargs={},
+        metadata=None,
+        standard_logging_payload=None,
+        omit_when_missing=True,
+        batch_trace_session_id="batch-uid-1",
+    )
+    assert session_id is None
+
+
+def test_get_logging_payload_groups_batch_create_and_cost_rows_in_one_session():
+    def _payload(call_type: str) -> SpendLogsPayload:
+        return get_logging_payload(
+            kwargs={
+                "call_type": call_type,
+                "model": "gpt-4o-mini",
+                "litellm_params": {"metadata": {"user_api_key": "test-key"}},
+            },
+            response_obj=litellm.ModelResponse(id="batch-uid-1", choices=[], usage=litellm.Usage()),
+            start_time=datetime.datetime.now(timezone.utc),
+            end_time=datetime.datetime.now(timezone.utc),
+        )
+
+    create_payload: Final = _payload("acreate_batch")
+    cost_payload: Final = _payload("aretrieve_batch")
+    assert cost_payload["request_id"] == "batch-uid-1_batch_cost"
+    assert create_payload["session_id"] == cost_payload["session_id"] == "batch-uid-1"
+
+
 @pytest.mark.parametrize(
     ("request_metadata", "expected"),
     [

ui/litellm-dashboard/src/components/view_logs/RequestLogsTableColumns.tsx

@@ -64,10 +64,12 @@ export const getRequestLogsTableColumns = ({
       const sessionAgentCount = log.session_agent_count ?? (isAgent ? sessionCount : 0);
       const sessionMcpCount = log.mcp_tool_call_count ?? (isMcp ? sessionCount : 0);
 
+      if (isBatchCallType(log.call_type)) {
+        return <BatchBadge count={sessionCount > 1 ? sessionCount : undefined} />;
+      }
       if (sessionCount <= 1) {
         if (isMcp) return <McpBadge />;
         if (isAgent) return <AgentBadge />;
-        if (isBatchCallType(log.call_type)) return <BatchBadge />;
         return <LlmBadge />;
       }