feat(proxy): track cost for unmanaged Vertex AI batch jobs (#31442)

* feat(proxy): track cost for unmanaged Vertex AI batch jobs

CheckBatchCost previously skipped Vertex batches created via the raw GCS
input_file_id path, since their unified_object_id is a raw provider job id
that fails the base64 managed-id check. Behind the opt-in general_settings
flag track_unmanaged_vertex_batch_cost, the poller now derives the model
from the gs:// input_file_id, maps it to a configured vertex_ai deployment,
polls the batch, computes cost, and marks batch_processed=True.

* Update tracking for failed", "expired", "cancelled"

* fix(proxy): apply ruff format to proxy_server.py

* address greptile review feedback (greploop iteration 1)

Filter unmanaged Vertex batch deployments by vertex_ai provider so a
shared model group name can't route to a wrong-provider deployment.
Move gs:// URI parsing into VertexAIBatchTransformation. Add test
coverage for the failed/expired/cancelled terminal-status DB update.

* fix: route unmanaged vertex batches to matching deployment

---------

Co-authored-by: Cursor Agent <cursoragent@cursor.com>

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -13,6 +13,8 @@ from litellm.constants import (
 )
 
 if TYPE_CHECKING:
+    from litellm.integrations.prometheus import PrometheusLogger
+    from litellm.proxy._types import LiteLLM_ManagedObjectTable
     from litellm.proxy.utils import PrismaClient, ProxyLogging
     from litellm.router import Router
 
@@ -26,6 +28,7 @@ class CheckBatchCost:
         proxy_logging_obj: "ProxyLogging",
         prisma_client: "PrismaClient",
         llm_router: "Router",
+        track_unmanaged_vertex_batch_cost: bool = False,
     ):
         from litellm.proxy.utils import PrismaClient, ProxyLogging
         from litellm.router import Router
@@ -33,6 +36,7 @@ class CheckBatchCost:
         self.proxy_logging_obj: ProxyLogging = proxy_logging_obj
         self.prisma_client: PrismaClient = prisma_client
         self.llm_router: Router = llm_router
+        self._track_unmanaged_vertex_batch_cost = track_unmanaged_vertex_batch_cost
         # Cached after the first poll cycle. Once we know the column is absent we skip
         # the guaranteed-failing primary query on every subsequent cycle.
         self._has_batch_processed_column: bool = True
@@ -97,6 +101,182 @@ class CheckBatchCost:
             order={"created_at": "asc"},
         )
 
+    @staticmethod
+    def _record_error(
+        prom_logger: Optional["PrometheusLogger"], error_type: str
+    ) -> None:
+        if prom_logger is not None:
+            prom_logger.record_check_batch_cost_error(error_type)
+
+    def _resolve_job_routing(
+        self,
+        job: "LiteLLM_ManagedObjectTable",
+        prom_logger: Optional["PrometheusLogger"],
+    ) -> Optional[Tuple[str, str]]:
+        """
+        Resolve (model_id, batch_id) for a managed-object row, where model_id is a router
+        deployment id and batch_id is the raw provider batch id.
+
+        Managed batches encode both in a base64 unified id. Unmanaged Vertex batches, created with
+        a raw gs:// input_file_id, store the raw provider job id as unified_object_id; when
+        track_unmanaged_vertex_batch_cost is enabled the model is derived from the gs:// path and
+        mapped to a configured vertex_ai deployment. Returns None (recording a metric) when the row
+        can't be routed.
+        """
+        from litellm.proxy.openai_files_endpoints.common_utils import (
+            _is_base64_encoded_unified_file_id,
+            get_batch_id_from_unified_batch_id,
+            get_model_id_from_unified_batch_id,
+        )
+
+        unified_object_id = job.unified_object_id
+        decoded = _is_base64_encoded_unified_file_id(unified_object_id)
+        if decoded:
+            model_id = get_model_id_from_unified_batch_id(decoded)
+            if model_id is None:
+                verbose_proxy_logger.info(
+                    f"Skipping job {unified_object_id} because it is not a valid model id"
+                )
+                self._record_error(prom_logger, "invalid_model_id")
+                return None
+            return model_id, get_batch_id_from_unified_batch_id(decoded)
+
+        if self._track_unmanaged_vertex_batch_cost:
+            return self._resolve_unmanaged_vertex_routing(job, prom_logger)
+
+        verbose_proxy_logger.info(
+            f"Skipping job {unified_object_id} because it is not a valid unified object id"
+        )
+        self._record_error(prom_logger, "invalid_unified_id")
+        return None
+
+    def _resolve_unmanaged_vertex_routing(
+        self,
+        job: "LiteLLM_ManagedObjectTable",
+        prom_logger: Optional["PrometheusLogger"],
+    ) -> Optional[Tuple[str, str]]:
+        from litellm.llms.vertex_ai.batches.transformation import (
+            VertexAIBatchTransformation,
+        )
+
+        input_file_id = self._get_input_file_id(job)
+        if not VertexAIBatchTransformation.is_unmanaged_gcs_batch_input_file_id(
+            input_file_id
+        ):
+            verbose_proxy_logger.info(
+                f"Skipping job {job.unified_object_id}: not an unmanaged vertex batch "
+                "(no gs:// input_file_id with a publishers/ model path)"
+            )
+            self._record_error(prom_logger, "invalid_unified_id")
+            return None
+        assert input_file_id is not None  # narrowed by is_unmanaged_gcs_batch_input_file_id
+
+        bare_model_name = VertexAIBatchTransformation.get_bare_model_name_from_gcs_file(
+            input_file_id
+        )
+        deployment_id = self._get_vertex_ai_deployment_id_for_bare_model(
+            bare_model_name
+        )
+        if deployment_id is None:
+            verbose_proxy_logger.info(
+                f"Skipping unmanaged vertex batch {job.unified_object_id}: no vertex_ai "
+                f"deployment configured for model {bare_model_name}"
+            )
+            self._record_error(prom_logger, "unmanaged_no_matching_deployment")
+            return None
+
+        return deployment_id, job.unified_object_id
+
+    def _get_vertex_ai_deployment_id_for_bare_model(
+        self, bare_model_name: str
+    ) -> Optional[str]:
+        model_group = self.llm_router.resolve_model_name_from_model_id(bare_model_name)
+        deployment_id = (
+            self._get_vertex_ai_deployment_id(model_group) if model_group else None
+        )
+        if deployment_id is not None:
+            return deployment_id
+
+        return self._get_vertex_ai_deployment_id_from_matching_deployments(
+            bare_model_name
+        )
+
+    def _get_vertex_ai_deployment_id_from_matching_deployments(
+        self, bare_model_name: str
+    ) -> Optional[str]:
+        from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
+
+        for deployment in self.llm_router.get_model_list(model_name=None) or []:
+            litellm_params = deployment.get("litellm_params") or {}
+            actual_model = litellm_params.get("model")
+            if not isinstance(actual_model, str):
+                continue
+            if not self._is_bare_model_match(actual_model, bare_model_name):
+                continue
+            try:
+                _, llm_provider, _, _ = get_llm_provider(
+                    model=actual_model,
+                    custom_llm_provider=litellm_params.get("custom_llm_provider"),
+                )
+            except Exception:
+                continue
+            if llm_provider != "vertex_ai":
+                continue
+            model_info = deployment.get("model_info") or {}
+            deployment_id = model_info.get("id")
+            if isinstance(deployment_id, str):
+                return deployment_id
+        return None
+
+    @staticmethod
+    def _is_bare_model_match(actual_model: str, bare_model_name: str) -> bool:
+        return (
+            actual_model == bare_model_name
+            or actual_model.endswith(f"/{bare_model_name}")
+            or actual_model.endswith(f":{bare_model_name}")
+        )
+
+    def _get_vertex_ai_deployment_id(self, model_group: str) -> Optional[str]:
+        """
+        Returns the first deployment id for `model_group` whose provider is vertex_ai,
+        skipping deployments from other providers that happen to share the model group name.
+        """
+        from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
+
+        for deployment_id in self.llm_router.get_model_ids(model_name=model_group):
+            deployment_info = self.llm_router.get_deployment(model_id=deployment_id)
+            if deployment_info is None:
+                continue
+            try:
+                _, llm_provider, _, _ = get_llm_provider(
+                    model=deployment_info.litellm_params.model,
+                    custom_llm_provider=deployment_info.litellm_params.custom_llm_provider,
+                )
+            except Exception:
+                continue
+            if llm_provider == "vertex_ai":
+                return deployment_id
+        return None
+
+    @staticmethod
+    def _get_input_file_id(job: "LiteLLM_ManagedObjectTable") -> Optional[str]:
+        import json
+
+        from litellm.types.utils import LiteLLMBatch
+
+        file_object = job.file_object
+        if isinstance(file_object, str):
+            try:
+                file_object = json.loads(file_object)
+            except (json.JSONDecodeError, ValueError):
+                return None
+        if not isinstance(file_object, dict):
+            return None
+        try:
+            return LiteLLMBatch.model_validate(file_object).input_file_id
+        except Exception:
+            return None
+
     async def check_batch_cost(self):
         """
         Check if the batch JOB has been tracked.
@@ -114,8 +294,6 @@ class CheckBatchCost:
         from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
         from litellm.proxy.openai_files_endpoints.common_utils import (
             _is_base64_encoded_unified_file_id,
-            get_batch_id_from_unified_batch_id,
-            get_model_id_from_unified_batch_id,
         )
 
         try:
@@ -172,31 +350,10 @@ class CheckBatchCost:
         else:
             jobs = await self._fallback_find_jobs()
         for job in jobs:
-            # get the model from the job
-            unified_object_id = job.unified_object_id
-            decoded_unified_object_id = _is_base64_encoded_unified_file_id(
-                unified_object_id
-            )
-            if not decoded_unified_object_id:
-                verbose_proxy_logger.info(
-                    f"Skipping job {unified_object_id} because it is not a valid unified object id"
-                )
-                if prom_logger:
-                    prom_logger.record_check_batch_cost_error("invalid_unified_id")
-                continue
-            else:
-                unified_object_id = decoded_unified_object_id
-
-            model_id = get_model_id_from_unified_batch_id(unified_object_id)
-            batch_id = get_batch_id_from_unified_batch_id(unified_object_id)
-
-            if model_id is None:
-                verbose_proxy_logger.info(
-                    f"Skipping job {unified_object_id} because it is not a valid model id"
-                )
-                if prom_logger:
-                    prom_logger.record_check_batch_cost_error("invalid_model_id")
+            routing = self._resolve_job_routing(job, prom_logger)
+            if routing is None:
                 continue
+            model_id, batch_id = routing
 
             verbose_proxy_logger.info(
                 f"Querying model ID: {model_id} for cost and usage of batch ID: {batch_id}"
@@ -213,7 +370,7 @@ class CheckBatchCost:
                 )
             except Exception as e:
                 verbose_proxy_logger.info(
-                    f"Skipping job {unified_object_id} because of error querying model ID: {model_id} for cost and usage of batch ID: {batch_id}: {e}"
+                    f"Skipping job {job.unified_object_id} because of error querying model ID: {model_id} for cost and usage of batch ID: {batch_id}: {e}"
                 )
                 if prom_logger:
                     prom_logger.record_check_batch_cost_error("provider_retrieval_error")
@@ -287,7 +444,7 @@ class CheckBatchCost:
                 deployment_info = self.llm_router.get_deployment(model_id=model_id)
                 if deployment_info is None:
                     verbose_proxy_logger.info(
-                        f"Skipping job {unified_object_id} because it is not a valid deployment info"
+                        f"Skipping job {job.unified_object_id} because it is not a valid deployment info"
                     )
                     if prom_logger:
                         prom_logger.record_check_batch_cost_error("deployment_not_found")
@@ -413,6 +570,26 @@ class CheckBatchCost:
                         f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}"
                     )
 
+            elif response.status in ("failed", "expired", "cancelled"):
+                try:
+                    update_data = {
+                        "status": response.status,
+                        "file_object": response.model_dump_json(),
+                    }
+                    if self._has_batch_processed_column:
+                        update_data["batch_processed"] = True
+                    await self.prisma_client.db.litellm_managedobjecttable.update(
+                        where={"id": job.id},
+                        data=update_data,
+                    )
+                    verbose_proxy_logger.info(
+                        f"CheckBatchCost: marked job {job.id} as {response.status} in DB"
+                    )
+                except Exception as db_err:
+                    verbose_proxy_logger.error(
+                        f"CheckBatchCost: failed to mark job {job.id} as {response.status} in DB: {db_err}"
+                    )
+
         # Record polling run metrics (always, even if nothing was processed)
         if prom_logger:
             prom_logger.record_check_batch_cost_run(

litellm/llms/vertex_ai/batches/transformation.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 from litellm._uuid import uuid
 from litellm.llms.vertex_ai.common_utils import (
@@ -207,3 +207,19 @@ class VertexAIBatchTransformation:
         parts = model_path.split("/")
         model = f"publishers/{'/'.join(parts[:3])}"
         return model
+
+    @classmethod
+    def is_unmanaged_gcs_batch_input_file_id(cls, input_file_id: Optional[str]) -> bool:
+        """
+        Returns True if `input_file_id` is a raw gs:// Vertex batch input file (i.e. not a
+        LiteLLM-managed unified file id) with a `publishers/` model path that
+        `_get_model_from_gcs_file` can parse.
+        """
+        return input_file_id is not None and input_file_id.startswith("gs://") and "publishers/" in input_file_id
+
+    @classmethod
+    def get_bare_model_name_from_gcs_file(cls, gcs_file_uri: str) -> str:
+        """
+        Extracts the bare model name (e.g. "gemini-1.5-flash-001") from a gcs file uri.
+        """
+        return cls._get_model_from_gcs_file(gcs_file_uri).rsplit("/", 1)[-1]

litellm/proxy/dev_config.yaml

@@ -193,6 +193,9 @@ model_list:
 
 general_settings:
   master_key: sk-1234
+  # Opt-in: let CheckBatchCost track cost for unmanaged Vertex batches created with a raw gs:// input_file_id.
+  # Requires a vertex_ai deployment configured for the batched model. Defaults to false.
+  # track_unmanaged_vertex_batch_cost: true
 
 sandbox_tools:
   - sandbox_tool_name: e2b_sandbox

litellm/proxy/proxy_server.py

@@ -7627,6 +7627,7 @@ class ProxyStartupEvent:
                     proxy_logging_obj=proxy_logging_obj,
                     prisma_client=prisma_client,
                     llm_router=llm_router,
+                    track_unmanaged_vertex_batch_cost=general_settings.get("track_unmanaged_vertex_batch_cost", False),
                 )
                 scheduler.add_job(
                     check_batch_cost_job.check_batch_cost,

tests/proxy_unit_tests/test_check_batch_cost.py

@@ -1,13 +1,35 @@
 """
 Unit tests for CheckBatchCost class.
 Covers: stale-row cleanup (file_purpose scoping), paginated find_many,
-and the batch_processed-column fallback query.
+the batch_processed-column fallback query, and routing of unmanaged
+Vertex batches (raw gs:// input_file_id, no managed unified id).
 """
 
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
+_IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id"
+
+
+def _unmanaged_vertex_file_object(
+    input_file_id="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash/abc.jsonl",
+    status="validating",
+):
+    """A LiteLLMBatch JSON blob shaped like what the managed-files hook stores for an
+    unmanaged Vertex batch (raw gs:// input_file_id)."""
+    from litellm.types.utils import LiteLLMBatch
+
+    return LiteLLMBatch(
+        id="8823717160934178816",
+        completion_window="24h",
+        created_at=1,
+        endpoint="/v1/chat/completions",
+        input_file_id=input_file_id,
+        object="batch",
+        status=status,
+    ).model_dump_json()
+
 
 class TestCheckBatchCost:
     """Test suite for CheckBatchCost class"""
@@ -375,6 +397,76 @@ class TestCheckBatchCost:
         ), "update() must include batch_processed=True when column is present"
         assert update_data["status"] == "complete"
 
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize("terminal_status", ["failed", "expired", "cancelled"])
+    async def test_terminal_status_marks_job_processed(
+        self,
+        check_batch_cost_instance,
+        mock_prisma_client,
+        mock_llm_router,
+        terminal_status,
+    ):
+        """When the provider reports a terminal status (failed/expired/cancelled), the row
+        must be written back with that status and batch_processed=True so it stops being
+        polled forever.
+        """
+        from unittest.mock import patch
+
+        mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock(
+            return_value=0
+        )
+        mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock()
+        mock_prisma_client.db.litellm_usertable.find_unique = AsyncMock(
+            return_value=None
+        )
+
+        mock_job = MagicMock()
+        mock_job.id = "job-terminal-1"
+        mock_job.unified_object_id = "dW5pZmllZF9iYXRjaF9pZA=="
+        mock_job.created_by = "user-1"
+
+        assert check_batch_cost_instance._has_batch_processed_column is True
+        mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock(
+            return_value=[mock_job]
+        )
+
+        mock_response = MagicMock()
+        mock_response.status = terminal_status
+        mock_response.model_dump_json.return_value = (
+            f'{{"id":"batch-1","status":"{terminal_status}"}}'
+        )
+
+        mock_llm_router.aretrieve_batch = AsyncMock(return_value=mock_response)
+
+        decoded_id = "llm_model_id,model-123;llm_batch_id,batch-456;"
+
+        with (
+            patch(
+                "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id",
+                side_effect=[decoded_id, None],
+            ),
+            patch(
+                "litellm.proxy.openai_files_endpoints.common_utils.get_model_id_from_unified_batch_id",
+                return_value="model-123",
+            ),
+            patch(
+                "litellm.proxy.openai_files_endpoints.common_utils.get_batch_id_from_unified_batch_id",
+                return_value="batch-456",
+            ),
+        ):
+            await check_batch_cost_instance.check_batch_cost()
+
+        assert (
+            mock_prisma_client.db.litellm_managedobjecttable.update.call_count == 1
+        ), f"Expected update() to be called exactly once for a {terminal_status} job"
+        update_data = mock_prisma_client.db.litellm_managedobjecttable.update.call_args[
+            1
+        ]["data"]
+        assert update_data["status"] == terminal_status
+        assert (
+            update_data["batch_processed"] is True
+        ), "terminal-status update() must set batch_processed=True so polling stops"
+
     @pytest.mark.asyncio
     async def test_raw_output_file_id_converted_to_managed_id(
         self, check_batch_cost_instance, mock_prisma_client, mock_llm_router
@@ -512,3 +604,249 @@ class TestCheckBatchCost:
         }
         assert mock_response.output_file_id == fake_managed_output_id
         assert mock_response.error_file_id == fake_managed_error_id
+
+
+class TestUnmanagedVertexRouting:
+    """Routing of unmanaged Vertex batches whose unified_object_id is a raw provider job id."""
+
+    def _instance(self, track_unmanaged, router):
+        from litellm_enterprise.proxy.common_utils.check_batch_cost import (
+            CheckBatchCost,
+        )
+
+        return CheckBatchCost(
+            proxy_logging_obj=MagicMock(),
+            prisma_client=MagicMock(),
+            llm_router=router,
+            track_unmanaged_vertex_batch_cost=track_unmanaged,
+        )
+
+    def _job(self, file_object=None):
+        job = MagicMock()
+        job.unified_object_id = "8823717160934178816"
+        job.file_object = (
+            file_object if file_object is not None else _unmanaged_vertex_file_object()
+        )
+        return job
+
+    def test_flag_off_skips_unmanaged_id_unchanged(self):
+        """Default (flag off): a raw numeric unified_object_id is skipped exactly as before;
+        no model derivation or router lookup happens."""
+        router = MagicMock()
+        instance = self._instance(track_unmanaged=False, router=router)
+        prom = MagicMock()
+
+        with patch(_IS_B64, return_value=False):
+            result = instance._resolve_job_routing(self._job(), prom)
+
+        assert result is None
+        prom.record_check_batch_cost_error.assert_called_once_with("invalid_unified_id")
+        router.resolve_model_name_from_model_id.assert_not_called()
+        router.get_model_ids.assert_not_called()
+
+    def _vertex_deployment(self):
+        deployment = MagicMock()
+        deployment.litellm_params.custom_llm_provider = "vertex_ai"
+        deployment.litellm_params.model = "vertex_ai/gemini-2.5-flash"
+        return deployment
+
+    def test_flag_on_routes_to_vertex_deployment(self):
+        """Flag on: derive the bare model from the gs:// path, resolve it to a deployment id,
+        and use the raw unified_object_id as the provider batch id."""
+        router = MagicMock()
+        router.resolve_model_name_from_model_id.return_value = "gemini-2.5-flash"
+        router.get_model_ids.return_value = ["deploy-1"]
+        router.get_deployment = MagicMock(return_value=self._vertex_deployment())
+        instance = self._instance(track_unmanaged=True, router=router)
+
+        with patch(_IS_B64, return_value=False):
+            result = instance._resolve_job_routing(self._job(), MagicMock())
+
+        assert result == ("deploy-1", "8823717160934178816")
+        # bare model name (trailing GCS segment), not the full publishers/.. path
+        router.resolve_model_name_from_model_id.assert_called_once_with(
+            "gemini-2.5-flash"
+        )
+        router.get_model_ids.assert_called_once_with(model_name="gemini-2.5-flash")
+
+    def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self):
+        """Flag on, but the only deployment for the model group is a non-vertex_ai
+        provider: must not be selected, even though the model group name matches."""
+        router = MagicMock()
+        router.resolve_model_name_from_model_id.return_value = "gemini-2.5-flash"
+        router.get_model_ids.return_value = ["deploy-openai"]
+        non_vertex_deployment = MagicMock()
+        non_vertex_deployment.litellm_params.custom_llm_provider = "openai"
+        non_vertex_deployment.litellm_params.model = "gpt-4o"
+        router.get_deployment = MagicMock(return_value=non_vertex_deployment)
+        instance = self._instance(track_unmanaged=True, router=router)
+        prom = MagicMock()
+
+        with patch(_IS_B64, return_value=False):
+            result = instance._resolve_job_routing(self._job(), prom)
+
+        assert result is None
+        prom.record_check_batch_cost_error.assert_called_once_with(
+            "unmanaged_no_matching_deployment"
+        )
+
+    def test_flag_on_uses_later_vertex_deployment_with_matching_suffix(self):
+        router = MagicMock()
+        router.resolve_model_name_from_model_id.return_value = "azure-gemini"
+        router.get_model_ids.return_value = ["deploy-azure"]
+        non_vertex_deployment = MagicMock()
+        non_vertex_deployment.litellm_params.custom_llm_provider = "azure"
+        non_vertex_deployment.litellm_params.model = "azure/gemini-2.5-flash"
+        router.get_deployment = MagicMock(return_value=non_vertex_deployment)
+        router.get_model_list.return_value = [
+            {
+                "model_name": "azure-gemini",
+                "litellm_params": {
+                    "model": "azure/gemini-2.5-flash",
+                    "custom_llm_provider": "azure",
+                },
+                "model_info": {"id": "deploy-azure"},
+            },
+            {
+                "model_name": "vertex-gemini",
+                "litellm_params": {
+                    "model": "vertex_ai/gemini-2.5-flash",
+                    "custom_llm_provider": "vertex_ai",
+                },
+                "model_info": {"id": "deploy-vertex"},
+            },
+        ]
+        instance = self._instance(track_unmanaged=True, router=router)
+
+        with patch(_IS_B64, return_value=False):
+            result = instance._resolve_job_routing(self._job(), MagicMock())
+
+        assert result == ("deploy-vertex", "8823717160934178816")
+        router.get_model_ids.assert_called_once_with(model_name="azure-gemini")
+
+    def test_flag_on_no_matching_deployment_records_metric(self):
+        """Flag on but no vertex_ai deployment for the model: skip with a distinct metric."""
+        router = MagicMock()
+        router.resolve_model_name_from_model_id.return_value = None
+        router.get_model_ids.return_value = []
+        instance = self._instance(track_unmanaged=True, router=router)
+        prom = MagicMock()
+
+        with patch(_IS_B64, return_value=False):
+            result = instance._resolve_job_routing(self._job(), prom)
+
+        assert result is None
+        prom.record_check_batch_cost_error.assert_called_once_with(
+            "unmanaged_no_matching_deployment"
+        )
+
+    def test_flag_on_non_gcs_input_is_not_unmanaged_vertex(self):
+        """Flag on, but input_file_id is not a gs:// publishers path: treat as unroutable,
+        do not attempt model derivation."""
+        router = MagicMock()
+        instance = self._instance(track_unmanaged=True, router=router)
+        prom = MagicMock()
+        job = self._job(
+            file_object=_unmanaged_vertex_file_object(input_file_id="file-abc-123")
+        )
+
+        with patch(_IS_B64, return_value=False):
+            result = instance._resolve_job_routing(job, prom)
+
+        assert result is None
+        prom.record_check_batch_cost_error.assert_called_once_with("invalid_unified_id")
+        router.resolve_model_name_from_model_id.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_end_to_end_costs_unmanaged_batch(self):
+        """Flag on, completed unmanaged batch: the poller polls Vertex with the raw job id,
+        computes cost, and marks batch_processed=True. Fails before this change (the row is
+        skipped at the unified-id gate)."""
+        router = MagicMock()
+        router.resolve_model_name_from_model_id.return_value = "gemini-2.5-flash"
+        router.get_model_ids.return_value = ["deploy-1"]
+
+        mock_response = MagicMock()
+        mock_response.status = "completed"
+        mock_response.output_file_id = "gs://bucket/out/predictions.jsonl"
+        mock_response.error_file_id = None
+        mock_response.completed_at = None
+        mock_response.created_at = None
+        mock_response.model_dump_json.return_value = (
+            '{"id":"8823717160934178816","status":"completed"}'
+        )
+        router.aretrieve_batch = AsyncMock(return_value=mock_response)
+        router.get_deployment_credentials_with_provider = MagicMock(
+            return_value={"vertex_project": "p", "vertex_location": "us-central1"}
+        )
+
+        deployment = MagicMock()
+        deployment.litellm_params.custom_llm_provider = "vertex_ai"
+        deployment.litellm_params.model = "vertex_ai/gemini-2.5-flash"
+        deployment.model_name = "gemini-2.5-flash"
+        deployment.model_info.model_dump.return_value = {}
+        router.get_deployment = MagicMock(return_value=deployment)
+
+        instance = self._instance(track_unmanaged=True, router=router)
+        instance.proxy_logging_obj.get_proxy_hook.return_value = None
+        instance._has_batch_processed_column = True
+
+        prisma = instance.prisma_client
+        prisma.db = MagicMock()
+        prisma.db.litellm_managedobjecttable = MagicMock()
+        prisma.db.litellm_managedobjecttable.update_many = AsyncMock(return_value=0)
+        prisma.db.litellm_managedobjecttable.update = AsyncMock()
+        prisma.db.litellm_managedobjecttable.find_many = AsyncMock(
+            return_value=[self._job()]
+        )
+        prisma.db.litellm_usertable = MagicMock()
+        prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=None)
+
+        mock_file_content = MagicMock()
+        mock_file_content.content = b'{"id":"req-1"}'
+
+        with (
+            patch(_IS_B64, side_effect=[False, None]),
+            patch(
+                "litellm.files.main.afile_content",
+                new_callable=AsyncMock,
+                return_value=mock_file_content,
+            ),
+            patch(
+                "litellm.batches.batch_utils._get_file_content_as_dictionary",
+                return_value=[{"id": "req-1"}],
+            ),
+            patch(
+                "litellm.batches.batch_utils.calculate_batch_cost_and_usage",
+                new_callable=AsyncMock,
+                return_value=(
+                    0.01,
+                    {"prompt_tokens": 10, "completion_tokens": 5},
+                    ["gemini-2.5-flash"],
+                ),
+            ),
+            patch(
+                "litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
+                return_value=("gemini-2.5-flash", "vertex_ai", None, None),
+            ),
+            patch(
+                "litellm.litellm_core_utils.litellm_logging.Logging"
+            ) as mock_logging_cls,
+        ):
+            mock_logging_obj = MagicMock()
+            mock_logging_obj.async_success_handler = AsyncMock()
+            mock_logging_cls.return_value = mock_logging_obj
+
+            await instance.check_batch_cost()
+
+        router.aretrieve_batch.assert_awaited_once()
+        assert router.aretrieve_batch.call_args[1]["model"] == "deploy-1"
+        assert router.aretrieve_batch.call_args[1]["batch_id"] == "8823717160934178816"
+
+        mock_logging_obj.async_success_handler.assert_awaited_once()
+        assert mock_logging_obj.async_success_handler.call_args[1]["batch_cost"] == 0.01
+
+        assert prisma.db.litellm_managedobjecttable.update.call_count == 1
+        update_data = prisma.db.litellm_managedobjecttable.update.call_args[1]["data"]
+        assert update_data["batch_processed"] is True
+        assert update_data["status"] == "complete"