fix(batches): cost an unmanaged gs:// batch through the deployment that created it (#49301)

* fix(batches): cost an unmanaged gs:// batch through the deployment that created it

The create now stores the router deployment id as model_id inside the
managed object row's file_object JSON and CheckBatchCost reads it before
falling back to the first deployment matching the base model in the
gs:// path. Rows without the field keep the old lookup

* fix(batches): tolerate a non-string or malformed file_object when reading the stored deployment

* fix(batches): read the stored creating deployment from the writer, not a lagging read replica

* fix(batches): resolve the writer database once so the type gate count does not grow

* fix(batches): decode a dict-shaped file_object when reading the stored deployment, accumulate wire samples, annotate test locals

---------

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -10,6 +10,7 @@ from dataclasses import replace as dataclasses_replace
 from datetime import datetime, timedelta
 from types import MappingProxyType
 from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast
+from pydantic import TypeAdapter, ValidationError
 
 from litellm_enterprise.proxy.common_utils.managed_object_attribution import (
     ManagedObjectCreatorAttribution,
@@ -50,6 +51,7 @@ def _error_names_resource(error: Exception, resource_id: str) -> bool:
 CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost"
 RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48)
 GIVE_UP_POLLING_AFTER: Final = timedelta(days=2 * MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
+_FILE_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object])
 
 
 PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = (
@@ -550,6 +552,10 @@ class CheckBatchCost:
         llm_provider: str,
         bare_model_name: str,
     ) -> Optional[Tuple[str, str]]:
+        stored_model_id: Final = self._get_stored_model_id(job)
+        if stored_model_id is not None and self._batch_deployment_exists(stored_model_id):
+            return stored_model_id, job.unified_object_id
+
         deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider)
         if deployment_id is None:
             verbose_proxy_logger.info(
@@ -662,23 +668,36 @@ class CheckBatchCost:
         )
 
     @staticmethod
-    def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]:
-        from litellm.types.utils import LiteLLMBatch
-
+    def _parse_file_object(job: "ManagedObjectRow") -> Optional[dict[str, object]]:
         file_object: Final = job.file_object
         try:
-            parsed_file_object: Final[object] = (
+            return _FILE_OBJECT_ADAPTER.validate_python(
                 json.loads(file_object) if isinstance(file_object, str) else file_object
             )
-        except (json.JSONDecodeError, ValueError):
+        except (json.JSONDecodeError, ValueError, ValidationError):
             return None
-        if not isinstance(parsed_file_object, dict):
+
+    @classmethod
+    def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]:
+        from litellm.types.utils import LiteLLMBatch
+
+        parsed_file_object: Final = cls._parse_file_object(job)
+        if parsed_file_object is None:
             return None
         try:
             return LiteLLMBatch.model_validate(parsed_file_object).input_file_id
         except Exception:
             return None
 
+    @classmethod
+    def _get_stored_model_id(cls, job: "ManagedObjectRow") -> Optional[str]:
+        """The deployment that created the batch, written into file_object by the managed
+        files hook at create time. Older rows have no such key and fall back to matching
+        the bare model name against the router."""
+        parsed_file_object: Final = cls._parse_file_object(job)
+        model_id: Final = parsed_file_object.get("model_id") if parsed_file_object else None
+        return model_id if isinstance(model_id, str) and model_id else None
+
     async def _track_completed_batch_cost(
         self,
         job: "ManagedObjectRow",

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -21,6 +21,7 @@ from typing import (
     TypedDict,
     Union,
     cast,
+    runtime_checkable,
 )
 from uuid import NAMESPACE_URL, uuid5
 
@@ -374,6 +375,12 @@ class _ManagedResourcePrismaClient(Protocol):
     def db(self) -> _ManagedResourceDatabase: ...
 
 
+@runtime_checkable
+class _WriterRoutedDatabase(Protocol):
+    @property
+    def writer(self) -> _ManagedResourceDatabase: ...
+
+
 class _SchedulerWithJobLookup(Protocol):
     def get_job(self, job_id: str) -> object: ...
 
@@ -479,6 +486,11 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag
     return prisma_client.db.litellm_managedobjecttable
 
 
+def _writer_database(db: _ManagedResourceDatabase) -> _ManagedResourceDatabase:
+    """Reads that feed a write must see the create that just happened; a lagging read replica may not"""
+    return db.writer if isinstance(db, _WriterRoutedDatabase) else db
+
+
 _JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
 _POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00"
 
@@ -512,6 +524,25 @@ def _without_model_output(
     )
 
 
+def _with_model_id(
+    file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse,
+    model_id: str | None,
+) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse:
+    return file_object if model_id is None else file_object.model_copy(update={"model_id": model_id})
+
+
+async def _stored_model_id(managed_object_table: _ManagedObjectTableActions, unified_object_id: str) -> str | None:
+    row: Final = await managed_object_table.find_first(where={"unified_object_id": unified_object_id})
+    if row is None:
+        return None
+    try:
+        parsed: Final = _JSON_VALUE_ADAPTER.validate_python(_decode_json_blob(row.file_object))
+    except (json.JSONDecodeError, ValidationError):
+        return None
+    model_id: Final = parsed.get("model_id") if isinstance(parsed, dict) else None
+    return model_id if isinstance(model_id, str) and model_id else None
+
+
 def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]:
     hidden_params: Final = cast(  # cast-ok: _hidden_params is an untyped attribute the upload path sets
         "Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {}
@@ -697,9 +728,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         persist_attribution: bool = False,
         create_if_missing: bool = True,
         batch_processed: bool = False,
+        creating_model_id: str | None = None,
     ) -> None:
         """Persist a managed object row, caching it and upserting it in the DB.
 
+        creating_model_id is the router deployment that created the object. It is kept
+        inside the stored file_object JSON as `model_id`, because the raw provider id an
+        unmanaged batch stores as unified_object_id says nothing about which deployment
+        made it, and the cost poller needs that deployment to read the result back. A
+        later observe of the row (create_if_missing=False) keeps the stored value rather
+        than overwriting it with whichever deployment served the retrieve.
+
         persist_attribution is set only by the batch and background response creates,
         the callers that can speak for the creator; it gates the api_key, request_tags, and,
         for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or
@@ -720,7 +759,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
         It is written only in the upsert create branch.
         """
         verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache")
-        stored_object: Final = _without_model_output(file_object)
+        managed_database: Final = _managed_database(self.prisma_client)
+        managed_object_table: Final = managed_database.litellm_managedobjecttable
+        preserved_model_id: Final = (
+            creating_model_id
+            if creating_model_id is not None or create_if_missing
+            else await _stored_model_id(_writer_database(managed_database).litellm_managedobjecttable, unified_object_id)
+        )
+        stored_object: Final = _with_model_id(_without_model_output(file_object), preserved_model_id)
         litellm_managed_object = LiteLLM_ManagedObjectTable(
             unified_object_id=unified_object_id,
             model_object_id=model_object_id,
@@ -752,7 +798,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             "status": file_object.status,
             "updated_by": user_api_key_dict.user_id,
         }
-        managed_object_table: Final = _managed_object_table(self.prisma_client)
         if not create_if_missing:
             await managed_object_table.update_many(
                 where={"unified_object_id": unified_object_id},
@@ -2170,6 +2215,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
                 request_tags=request_tags_from_metadata(request_metadata if isinstance(request_metadata, dict) else {}),
                 persist_attribution=is_batch_create,
                 create_if_missing=is_batch_create,
+                creating_model_id=model_id if is_batch_create else None,
             )
             if not is_batch_create:
                 await self.store_batch_output_file_ownership(

tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py

@@ -0,0 +1,181 @@
+"""Two Vertex deployments share a base model but live in different projects with different
+buckets. An unmanaged gs:// batch created on the second one must be polled through that
+deployment (LIT-10326). Before the fix the cost poller matched the base model name in the
+gs:// path to the first deployment in the list and asked the wrong project for the job."""
+
+import json
+import re
+from collections.abc import Callable
+from contextlib import ExitStack
+from pathlib import Path
+from typing import Final
+from uuid import uuid4
+
+import pytest
+import yaml
+from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value
+from integration._support.process import owned_proxy_process
+from integration._support.vertex import service_account_json
+from integration._support.wire import Reply, Request, wire_server
+from pydantic import JsonValue
+
+LOCATION: Final = "us-central1"
+BASE_MODEL: Final = "gemini-2.5-flash"
+TEAMS: Final = ("a", "b")
+_JOB_PATH: Final = re.compile(
+    r"^/v1/projects/(?P<project>[^/]+)/locations/(?P<location>[^/]+)/batchPredictionJobs(?:/(?P<job>\d+))?$"
+)
+
+
+def _project(team: str) -> str:
+    return f"scripted-project-{team}"
+
+
+def _deployment(team: str) -> str:
+    return f"{BASE_MODEL}-team-{team}"
+
+
+def _input_file(team: str) -> str:
+    return f"gs://scripted-bucket-{team}/litellm-vertex-files/publishers/google/models/{BASE_MODEL}/{uuid4().hex}/input.jsonl"
+
+
+def _job_number(team: str) -> str:
+    return f"{TEAMS.index(team) + 1}{uuid4().int % 10**17:017d}"
+
+
+def _vertex_peer(jobs: dict[str, tuple[str, str]]) -> Callable[[Request], Reply]:
+    """A Vertex double that remembers which project each job was created in. A retrieve
+    from any other project answers the 404 real Vertex gives (the job does not exist
+    there), a retrieve from the right one reports the job still running."""
+
+    def respond(request: Request) -> Reply:
+        path: Final = _JOB_PATH.match(request.target.split("?", 1)[0])
+        assert path is not None, request.target
+        assert request.headers.get("authorization") == "Bearer scripted-token", request.headers
+        project, location, job = path.group("project"), path.group("location"), path.group("job")
+        if request.method == "POST":
+            team: Final = project.rsplit("-", 1)[-1]
+            uri: Final = re.search(r'"(gs://[^"]+)"', request.body.decode())
+            assert uri is not None and uri.group(1).startswith(f"gs://scripted-bucket-{team}/"), (project, request.body)
+            created, created_uri = jobs.setdefault(project, (_job_number(team), uri.group(1)))
+            return Reply(body=json.dumps(_job(project, location, created, created_uri, "JOB_STATE_PENDING")).encode())
+        assert request.method == "GET" and job is not None, (request.method, request.target)
+        known: Final = jobs.get(project)
+        if known is None or known[0] != job:
+            return Reply(
+                status=404,
+                body=json.dumps(
+                    {
+                        "error": {
+                            "code": 404,
+                            "message": f"The BatchPredictionJob does not exist: {job}",
+                            "status": "NOT_FOUND",
+                        }
+                    }
+                ).encode(),
+            )
+        return Reply(body=json.dumps(_job(project, location, job, known[1], "JOB_STATE_RUNNING")).encode())
+
+    return respond
+
+
+def _job(project: str, location: str, job: str, input_uri: str, state: str) -> dict[str, JsonValue]:
+    return {
+        "name": f"projects/{project}/locations/{location}/batchPredictionJobs/{job}",
+        "displayName": "litellm-vertex-batch-scripted",
+        "model": f"publishers/google/models/{BASE_MODEL}",
+        "inputConfig": {"gcsSource": {"uris": [input_uri]}, "instancesFormat": "jsonl"},
+        "outputConfig": {"predictionsFormat": "jsonl", "gcsDestination": {"outputUriPrefix": input_uri.rsplit("/", 1)[0]}},
+        "state": state,
+        "createTime": "2026-10-30T20:00:00.000000Z",
+        "updateTime": "2026-10-30T20:00:00.000000Z",
+    }
+
+
+def _owned_config(path: Path, gateway: Gateway, wire_url: str) -> Path:
+    config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()))
+    path.write_text(
+        yaml.safe_dump(
+            {
+                **config,
+                "model_list": [
+                    {
+                        "model_name": _deployment(team),
+                        "litellm_params": {
+                            "model": f"vertex_ai/{BASE_MODEL}",
+                            "api_base": wire_url,
+                            "vertex_project": _project(team),
+                            "vertex_location": LOCATION,
+                            "vertex_credentials": service_account_json(_project(team), gateway.upstream_url.rstrip("/")),
+                            "gcs_bucket_name": f"scripted-bucket-{team}",
+                        },
+                    }
+                    for team in TEAMS
+                ],
+                "general_settings": {
+                    **object_value(config.get("general_settings") or {}),
+                    "track_unmanaged_batch_cost": True,
+                },
+                "litellm_settings": {
+                    **object_value(config.get("litellm_settings") or {}),
+                    "enable_loadbalancing_on_batch_endpoints": True,
+                    "user_url_allowed_hosts": ["127.0.0.1"],
+                },
+            }
+        )
+    )
+    return path
+
+
+def _create_batch(gateway: Gateway, key: str, team: str) -> str:
+    response: Final = gateway.request(
+        "POST",
+        "/v1/batches",
+        {
+            "input_file_id": _input_file(team),
+            "endpoint": "/v1/chat/completions",
+            "completion_window": "24h",
+            "model": _deployment(team),
+        },
+        key=key,
+    )
+    assert response.status_code == 200, response.text
+    return string_value(JSON_OBJECT.validate_json(response.content)["id"])
+
+
+def _retrieve_targets(requests: tuple[Request, ...], job: str) -> tuple[str, ...]:
+    return tuple(request.target for request in requests if request.method == "GET" and request.target.endswith(f"/{job}"))
+
+
+@pytest.mark.timeout(240)
+def test_unmanaged_batch_on_the_second_deployment_sharing_a_base_model_is_polled_through_its_own_project(
+    gateway: Gateway, tmp_path: Path
+) -> None:
+    jobs: Final[dict[str, tuple[str, str]]] = {}  # mutable-ok: the wire double records which project created each job
+    with ExitStack() as stack:
+        wire: Final = stack.enter_context(wire_server(_vertex_peer(jobs)))
+        config: Final = _owned_config(tmp_path / "two-vertex-teams.yaml", gateway, wire.url)
+        owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config))
+        scenario: Final = stack.enter_context(owned.gateway.scenario())
+        key: Final = scenario.key(models=[_deployment(team) for team in TEAMS])
+
+        batch_a: Final = _create_batch(owned.gateway, key, "a")
+        batch_b: Final = _create_batch(owned.gateway, key, "b")
+        assert (batch_a, batch_b) == (jobs[_project("a")][0], jobs[_project("b")][0]), (batch_a, batch_b, jobs)
+
+        seen: Final[list[Request]] = []  # mutable-ok: drain empties the wire queue, so samples accumulate across retries
+
+        def polled_so_far() -> tuple[Request, ...]:
+            seen.extend(wire.drain())
+            return tuple(seen)
+
+        polled: Final = eventually(
+            polled_so_far,
+            lambda requests: bool(_retrieve_targets(requests, batch_b)) and bool(_retrieve_targets(requests, batch_a)),
+            seconds=90,
+        )
+        for team, batch in (("a", batch_a), ("b", batch_b)):
+            assert all(
+                target.startswith(f"/v1/projects/{_project(team)}/") for target in _retrieve_targets(polled, batch)
+            ), (team, _retrieve_targets(polled, batch))
+

tests/unit/enterprise/proxy/test_managed_files_hook.py

@@ -1437,6 +1437,8 @@ def _make_object_store_instance():
     mock_prisma = MagicMock()
     mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock()
     mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock()
+    mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None)
+    mock_prisma.db.writer = mock_prisma.db
 
     return (
         PROXY_LiteLLMManagedFiles(
@@ -2161,3 +2163,108 @@ async def test_store_unified_file_id_caches_the_storage_location_the_db_row_gets
     assert (cached.storage_backend, cached.storage_url) == ("litellm_db", "litellm_db://content-row-1")
     create_data = file_table.upsert.await_args.kwargs["data"]["create"]
     assert (create_data["storage_backend"], create_data["storage_url"]) == ("litellm_db", "litellm_db://content-row-1")
+
+
+@pytest.mark.asyncio
+async def test_create_stores_the_creating_deployment_in_file_object():
+    """Regression (LIT-10326): an unmanaged gs:// batch stores the raw Vertex job id, which
+    says nothing about which of two deployments sharing a base model created it. The create
+    must keep the router's deployment id inside the stored file_object so the cost poller
+    reads the result through that deployment instead of the first base-model match."""
+    managed_files, mock_prisma = _make_object_store_instance()
+    creator: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None)
+
+    await managed_files.store_unified_object_id(
+        unified_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_object=_make_batch_response(status="validating"),
+        litellm_parent_otel_span=None,
+        model_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_purpose="batch",
+        user_api_key_dict=creator,
+        persist_attribution=True,
+        create_if_missing=True,
+        creating_model_id="deploy-team-b",
+    )
+
+    created: Final = mock_prisma.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"]
+    assert json.loads(created["file_object"])["model_id"] == "deploy-team-b"
+
+
+@pytest.mark.asyncio
+async def test_poll_refresh_keeps_the_stored_creating_deployment():
+    """A later observe of the row (a GET through whichever deployment the client named)
+    rewrites file_object; it must carry the create-time model_id forward, not drop it or
+    replace it with the deployment that served the retrieve."""
+    managed_files, mock_prisma = _make_object_store_instance()
+    stored_row: Final = MagicMock()
+    stored_row.file_object = json.dumps({**json.loads(_make_batch_response(status="validating").model_dump_json()), "model_id": "deploy-team-b"})
+    mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row)
+    observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None)
+
+    await managed_files.store_unified_object_id(
+        unified_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_object=_make_batch_response(status="completed"),
+        litellm_parent_otel_span=None,
+        model_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_purpose="batch",
+        user_api_key_dict=observer,
+        persist_attribution=False,
+        create_if_missing=False,
+    )
+
+    written: Final = json.loads(mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"])
+    assert written["model_id"] == "deploy-team-b"
+    assert written["status"] == "completed"
+
+
+@pytest.mark.asyncio
+async def test_poll_refresh_reads_the_stored_creating_deployment_from_the_writer():
+    """With a read replica configured, a retrieve that lands before the replica has the create
+    must still find the create-time model_id, so the lookup goes to the writer. Prisma may hand
+    the Json column back already decoded, so the row here is a dict rather than JSON text."""
+    managed_files, mock_prisma = _make_object_store_instance()
+    stored_row: Final = MagicMock()
+    stored_row.file_object = {**json.loads(_make_batch_response(status="validating").model_dump_json()), "model_id": "deploy-team-b"}
+    mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None)
+    mock_prisma.db.writer = MagicMock()
+    mock_prisma.db.writer.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row)
+    observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None)
+
+    await managed_files.store_unified_object_id(
+        unified_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_object=_make_batch_response(status="completed"),
+        litellm_parent_otel_span=None,
+        model_object_id="projects/123/locations/us/batchPredictionJobs/456",
+        file_purpose="batch",
+        user_api_key_dict=observer,
+        persist_attribution=False,
+        create_if_missing=False,
+    )
+
+    written: Final = json.loads(mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"])
+    assert written["model_id"] == "deploy-team-b"
+    mock_prisma.db.litellm_managedobjecttable.find_first.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_batch_create_hook_passes_the_router_deployment_as_creating_model_id():
+    """The post-call hook on a batch create hands the router's model_id to the object store;
+    a retrieve of the same batch hands None so it cannot overwrite the creator."""
+    from litellm_enterprise.proxy.hooks.managed_files import BATCH_CREATE_HIDDEN_PARAM
+
+    managed_files: Final = _make_managed_files_instance()
+    created: Final = _make_batch_response(model_id="deploy-team-b", output_file_id=None, status="validating")
+    created._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True
+    retrieved: Final = _make_batch_response(model_id="deploy-team-a", output_file_id=None, status="in_progress")
+
+    with patch("litellm.proxy.proxy_server.llm_router", MagicMock()):
+        await managed_files.async_post_call_success_hook(
+            data={}, user_api_key_dict=_make_user_api_key_dict(), response=created
+        )
+        await managed_files.async_post_call_success_hook(
+            data={}, user_api_key_dict=_make_user_api_key_dict(), response=retrieved
+        )
+
+    first, second = managed_files.store_unified_object_id.await_args_list
+    assert first.kwargs["creating_model_id"] == "deploy-team-b"
+    assert second.kwargs["creating_model_id"] is None

tests/unit/proxy/common_utils/test_check_batch_cost.py

@@ -1771,6 +1771,53 @@ class TestUnmanagedVertexRouting:
 
         assert result == ("deploy-ft", "8823717160934178816")
 
+    def _two_team_deployments(self):
+        """Two Vertex deployments sharing the base model, team A listed first."""
+        return [
+            {
+                "model_name": f"gemini-2.5-flash-team-{team}",
+                "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "custom_llm_provider": "vertex_ai"},
+                "model_info": {"id": f"deploy-team-{team}"},
+            }
+            for team in ("a", "b")
+        ]
+
+    def test_flag_on_prefers_the_deployment_stored_on_the_row_over_the_first_base_model_match(self):
+        """Regression (LIT-10326): two deployments share gemini-2.5-flash with different
+        buckets. A batch created on team B's deployment stores deploy-team-b in file_object,
+        so the poller must read it back through team B, not the first match (team A)."""
+        router: Final = MagicMock()
+        router.registry.return_value.resolve.return_value = None
+        router.registry.return_value.deployments.return_value = self._two_team_deployments()
+        router.registry.return_value.deployment = MagicMock(return_value=self._vertex_deployment())
+        instance: Final = self._instance(track_unmanaged=True, router=router)
+        job: Final = self._job(
+            file_object=json.dumps({**json.loads(_unmanaged_vertex_file_object()), "model_id": "deploy-team-b"})
+        )
+
+        with patch(_IS_B64, return_value=False):
+            result: Final = instance._resolve_job_routing(job, MagicMock())
+
+        assert result == ("deploy-team-b", "8823717160934178816")
+        router.registry.return_value.deployment.assert_called_once_with("deploy-team-b")
+
+    def test_flag_on_falls_back_to_base_model_match_when_the_stored_deployment_is_gone(self):
+        """The stored deployment was removed from the config after the batch was created:
+        the row still routes through today's base-model lookup instead of being skipped."""
+        router: Final = MagicMock()
+        router.registry.return_value.resolve.return_value = None
+        router.registry.return_value.deployments.return_value = self._two_team_deployments()
+        router.registry.return_value.deployment = MagicMock(return_value=None)
+        instance: Final = self._instance(track_unmanaged=True, router=router)
+        job: Final = self._job(
+            file_object=json.dumps({**json.loads(_unmanaged_vertex_file_object()), "model_id": "deploy-removed"})
+        )
+
+        with patch(_IS_B64, return_value=False):
+            result: Final = instance._resolve_job_routing(job, MagicMock())
+
+        assert result == ("deploy-team-a", "8823717160934178816")
+
     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."""