litellm/tests/unit/proxy/batches_endpoints/test_endpoints.py

"""
Routing-contract tests for litellm/proxy/batches_endpoints/endpoints.py

These are not happy-path smoke tests. Each row of the matrix locks the full
contract of a single routing branch so that *any* behavior change in this layer
fails loudly:

  1. DISPATCH      - exactly which downstream seam fired (litellm.acreate_batch
                     vs llm_router.acreate_batch), and every sibling seam is
                     asserted NOT called. A reordered/negated branch flips this.
  2. CREDENTIALS   - the credential resolver receives the model derived from the
                     request, not a hardcoded value. The router's
                     get_deployment_credentials_with_provider is input-locked.
  3. SEAM PAYLOAD  - the *entire* kwargs dict forwarded to the provider call is
                     exact-matched. Because LiteLLMBatchCreateRequest is a
                     TypedDict (zero runtime filtering), nothing else stops a
                     newly-added param from silently reaching every provider.
                     This exact-match is that missing guard: a new key fails the
                     test and forces a reviewer to ask "does this work for all
                     providers, or just openai".
  4. OUTPUT SHAPE  - the id encode/decode round-trip clients depend on.

Only true I/O boundaries are mocked (provider call, router, proxy logging,
request parsing, pre-call enrichment). The pure encode/decode/credential-merge
helpers run for real so the payload assertions reflect production exactly.

The object mocks are spec'd to their real classes, so a brand-new method call
added to this layer raises instead of silently passing - the inventory of seams
cannot drift without a test failure.
"""

import base64
import json
import logging
from contextlib import ExitStack
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any, Dict, Optional
from unittest.mock import AsyncMock, MagicMock, patch

import httpx
import pytest
import respx
from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles

import litellm
import litellm.proxy.batches_endpoints.endpoints as endpoints
import litellm.proxy.proxy_server as proxy_server
from litellm.proxy._types import ProxyException, UserAPIKeyAuth
from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing
from litellm.proxy.openai_files_endpoints.common_utils import (
    encode_file_id_with_model,
)
from litellm.proxy.utils import ProxyLogging
from litellm.router import Router
from litellm.types.llms.openai import BatchJobStatus
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
from litellm.types.utils import CredentialItem, LiteLLMBatch, SpecialEnums

from fastapi import HTTPException, Request, Response

# --------------------------------------------------------------------------- #
# Fixtures: distinguishable credentials per model so a wrong/hardcoded model_id
# produces wrong creds (or KeyError) and is impossible to hide.
# --------------------------------------------------------------------------- #

CREDS: Dict[str, Dict[str, str]] = {
    "azure/gpt-4o": {
        "custom_llm_provider": "azure",
        "api_key": "sk-azure",
        "api_base": "https://azure.test",
        "model": "azure/gpt-4o-deployment",
    },
    "vertex-model": {
        "custom_llm_provider": "vertex_ai",
        "api_key": "sk-vertex",
        "api_base": "https://vertex.test",
        "model": "vertex_ai/gemini-2.0",
    },
    "my-vllm": {
        "custom_llm_provider": "hosted_vllm",
        "api_key": "sk-vllm",
        "api_base": "http://vllm.test/v1",
        "model": "hosted_vllm/qwen",
    },
}

# A real model-encoded file id: decodes to "azure/gpt-4o", strips to "file-original123".
AZURE_FILE_ID = encode_file_id_with_model("file-original123", "azure/gpt-4o", id_type="file")


def make_batch(
    *,
    id: str = "batch-provider-id",
    output_file_id: Optional[str] = None,
    error_file_id: Optional[str] = None,
    input_file_id: Optional[str] = None,
    status: BatchJobStatus = "validating",
) -> LiteLLMBatch:
    batch = LiteLLMBatch(
        id=id,
        completion_window="24h",
        created_at=1234567890,
        endpoint="/v1/chat/completions",
        input_file_id=input_file_id or "file-provider-input",
        object="batch",
        status=status,
    )
    if output_file_id is not None:
        batch.output_file_id = output_file_id
    if error_file_id is not None:
        batch.error_file_id = error_file_id
    batch._hidden_params = {}
    return batch


class FakeRequest:
    """Minimal stand-in. The request is only read via .headers/.query_params on
    the model-param fallback path; everything else that touches it is mocked."""

    def __init__(
        self,
        headers: Optional[Dict[str, str]] = None,
        query: Optional[Dict[str, str]] = None,
    ):
        self.headers = headers or {}
        self.query_params = query or {}


@pytest.fixture
def openai_env_creds(monkeypatch):
    """Deterministic env creds so the implicit-openai fallback forwards instead
    of tripping the no-creds 404 gate, regardless of the host environment."""
    monkeypatch.setenv("OPENAI_API_KEY", "sk-env-openai")


@pytest.fixture
def no_openai_creds(monkeypatch):
    """Neutralize every credential source the 404 gate checks."""
    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
    monkeypatch.setattr(litellm, "api_key", None)
    monkeypatch.setattr(litellm, "openai_key", None)


@dataclass
class Harness:
    """Holds every mocked seam so a test can configure inputs and assert calls."""

    body: Dict[str, Any]
    read_body: AsyncMock
    pre_call: AsyncMock
    get_headers: MagicMock
    provider_from_headers: MagicMock
    is_known_model: MagicMock
    litellm_acreate: AsyncMock
    router: MagicMock
    logging: MagicMock
    creds_resolver: MagicMock
    upstream_files_route: respx.Route

    @property
    def router_acreate(self) -> AsyncMock:
        return self.router.acreate_batch

    def acreate_kwargs(self) -> Dict[str, Any]:
        """Exact kwargs forwarded to litellm.acreate_batch."""
        assert self.litellm_acreate.call_count == 1
        return dict(self.litellm_acreate.call_args.kwargs)

    def router_kwargs(self) -> Dict[str, Any]:
        assert self.router_acreate.call_count == 1
        return dict(self.router_acreate.call_args.kwargs)


def _creds_lookup(*, model_id: str, team_id: str | None = None) -> dict[str, str] | None:
    # An unknown/hardcoded model_id resolves to None exactly like the real router,
    # which the endpoint turns into a 400 and a missing dispatch - the bug cannot hide.
    return dict(CREDS[model_id]) if model_id in CREDS else None


@pytest.fixture
def harness(monkeypatch: pytest.MonkeyPatch):
    """Seam harness. Patches only true I/O boundaries; pure encode/decode/merge
    helpers run for real. Object mocks are spec'd so unknown method calls raise."""
    body_holder: Dict[str, Any] = {}
    logging = MagicMock(spec=ProxyLogging)
    logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"])
    logging.post_call_failure_hook = AsyncMock()
    logging.update_request_status = AsyncMock()
    logging.get_proxy_hook = MagicMock(return_value=None)

    router = MagicMock(spec=Router)
    router.model_group_alias = {}
    router.get_model_access_groups = MagicMock(return_value={})
    router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id)
    router.model_list = []
    router.acreate_batch = AsyncMock(return_value=make_batch())
    router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)

    read_body = AsyncMock(side_effect=lambda request: body_holder["body"])
    pre_call = AsyncMock(side_effect=lambda **kw: (body_holder["body"], MagicMock()))
    get_headers = MagicMock(return_value={})
    provider_from_headers = MagicMock(return_value=None)
    is_known_model = MagicMock(return_value=False)
    litellm_acreate = AsyncMock(return_value=make_batch())
    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)

    with ExitStack() as stack:
        stack.enter_context(patch.object(endpoints, "read_request_body", read_body))
        stack.enter_context(
            patch.object(
                ProxyBaseLLMRequestProcessing,
                "common_processing_pre_call_logic",
                pre_call,
            )
        )
        stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
        stack.enter_context(
            patch.object(
                endpoints,
                "get_custom_llm_provider_from_request_headers",
                provider_from_headers,
            )
        )
        stack.enter_context(patch.object(endpoints, "is_known_model", is_known_model))
        stack.enter_context(patch.object(litellm, "acreate_batch", litellm_acreate))
        stack.enter_context(patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False))
        upstream = stack.enter_context(respx.mock(assert_all_called=False))
        upstream_files_route = upstream.get(f"{CREDS['my-vllm']['api_base']}/files").mock(
            return_value=httpx.Response(404, json={"detail": "Not Found"})
        )
        stack.enter_context(patch.object(proxy_server, "llm_router", router))
        stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging))
        stack.enter_context(patch.object(proxy_server, "general_settings", {}))
        stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock()))
        stack.enter_context(patch.object(proxy_server, "version", "test-version"))
        stack.enter_context(patch.object(proxy_server, "prisma_client", None))

        h = Harness(
            body=body_holder,
            read_body=read_body,
            pre_call=pre_call,
            get_headers=get_headers,
            provider_from_headers=provider_from_headers,
            is_known_model=is_known_model,
            litellm_acreate=litellm_acreate,
            router=router,
            logging=logging,
            creds_resolver=router.get_deployment_credentials_with_provider,
            upstream_files_route=upstream_files_route,
        )
        yield h


def set_body(harness: Harness, body: Dict[str, Any]) -> None:
    harness.body["body"] = body


async def call_create(
    harness: Harness,
    *,
    provider: Optional[str] = None,
    user: Optional[UserAPIKeyAuth] = None,
    headers: Optional[Dict[str, str]] = None,
    query: Optional[Dict[str, str]] = None,
):
    return await endpoints.create_batch(
        request=FakeRequest(headers=headers, query=query),
        fastapi_response=Response(),
        provider=provider,
        user_api_key_dict=user or UserAPIKeyAuth(api_key="sk-test"),
    )


@pytest.fixture
def executed_runner():
    runner = MagicMock(spec=endpoints.LiteLLMExecutedBatchRunner)
    runner.create = AsyncMock(return_value=make_batch(id="litellm-executed-batch"))
    runner.cancel = AsyncMock(return_value=make_batch(id="litellm-executed-batch", status="cancelling"))
    factory = MagicMock(return_value=runner)
    with patch.object(  # test-quality-ok: the route builds its runner from proxy_server globals; the factory is the only seam
        endpoints, "_litellm_executed_batch_runner", factory
    ):
        yield runner, factory


def _managed_input_file_id(model: str) -> str:
    unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
        "application/jsonl", "managed-id", model, "file-id", "file-model-id"
    )
    return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")


# =========================================================================== #
# SCENARIO 1 - input_file_id encoded with model. The full showcase: every
# assertion type from the design lives here.
# =========================================================================== #


@pytest.mark.asyncio
async def test_create__model_encoded_file_id(harness):
    set_body(
        harness,
        {
            "input_file_id": AZURE_FILE_ID,
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    resp = await call_create(harness)

    # 1. DISPATCH - model-credential path fired via litellm, router did not.
    assert harness.litellm_acreate.call_count == 1
    harness.router_acreate.assert_not_called()

    # 2. CREDENTIALS - resolved for the model decoded FROM the file id.
    harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")

    # 3. SEAM PAYLOAD - exact, whole dict. A new forwarded key breaks this.
    assert harness.acreate_kwargs() == {
        "custom_llm_provider": "azure",
        "input_file_id": "file-original123",  # encoding stripped by this layer
        "endpoint": "/v1/chat/completions",
        "completion_window": "24h",
        "metadata": None,  # sanitize_openai_provider_metadata(None)
        "api_key": "sk-azure",
        "api_base": "https://azure.test",
        "model": "azure/gpt-4o-deployment",
    }

    # 4. OUTPUT SHAPE - ids re-encoded with the model; input_file_id restored.
    assert resp.id == encode_file_id_with_model("batch-provider-id", "azure/gpt-4o", id_type="batch")
    assert resp.input_file_id == AZURE_FILE_ID


@pytest.mark.asyncio
async def test_create__model_encoded_file_id__encodes_output_and_error_ids(harness):
    set_body(
        harness,
        {
            "input_file_id": AZURE_FILE_ID,
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    harness.litellm_acreate.return_value = make_batch(
        id="batch-xyz",
        output_file_id="file-out-raw",
        error_file_id="file-err-raw",
    )

    resp = await call_create(harness)

    assert resp.output_file_id == encode_file_id_with_model("file-out-raw", "azure/gpt-4o")
    assert resp.error_file_id == encode_file_id_with_model("file-err-raw", "azure/gpt-4o")


@pytest.mark.asyncio
async def test_create__model_encoded_file_id__resolver_gets_decoded_model(harness):
    """Regression guard: model_id for credential resolution must be derived from
    the file id. A hardcode would call the resolver with the wrong model."""
    set_body(
        harness,
        {
            "input_file_id": AZURE_FILE_ID,
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    await call_create(harness)

    harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")


# =========================================================================== #
# SCENARIO 2 - model from body / header / query. Locks source precedence.
# =========================================================================== #


@pytest.mark.asyncio
async def test_create__model_from_body(harness):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
            "model": "vertex-model",
        },
    )

    resp = await call_create(harness)

    assert harness.litellm_acreate.call_count == 1
    harness.router_acreate.assert_not_called()
    harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
    payload = harness.acreate_kwargs()
    assert payload["custom_llm_provider"] == "vertex_ai"
    assert payload["input_file_id"] == "file-plain"
    assert resp.id == encode_file_id_with_model("batch-provider-id", "vertex-model", id_type="batch")


@pytest.mark.asyncio
async def test_create__model_from_header(harness):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    await call_create(harness, headers={"x-litellm-model": "vertex-model"})

    harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
    harness.router_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_create__model_from_query(harness):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    await call_create(harness, query={"model": "vertex-model"})

    harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
    harness.router_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_create__body_model_beats_header_and_query(harness):
    """Precedence row: body > header > query (data.get('model') first)."""
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
            "model": "azure/gpt-4o",
        },
    )

    await call_create(
        harness,
        headers={"x-litellm-model": "vertex-model"},
        query={"model": "vertex-model"},
    )

    harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")


# =========================================================================== #
# SCENARIO 3 - fallback to custom_llm_provider (env-var creds). MUST NOT touch
# the credential resolver.
# =========================================================================== #


@pytest.mark.asyncio
async def test_create__fallback_default_openai(harness, openai_env_creds):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    await call_create(harness)

    assert harness.litellm_acreate.call_count == 1
    harness.router_acreate.assert_not_called()
    harness.creds_resolver.assert_not_called()  # inverse-bug guard
    assert harness.acreate_kwargs()["custom_llm_provider"] == "openai"


@pytest.mark.asyncio
async def test_create__fallback_no_creds_404(harness, no_openai_creds):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    with pytest.raises(ProxyException) as exc:
        await call_create(harness)

    assert exc.value.code == "404"
    assert exc.value.type == "invalid_request_error"
    assert exc.value.param is None
    assert exc.value.message == "No such File object: file-plain"
    harness.litellm_acreate.assert_not_called()
    harness.router_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_create__fallback_explicit_provider_bypasses_not_found_gate(harness, no_openai_creds):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    await call_create(harness, provider="anthropic")

    assert harness.acreate_kwargs()["custom_llm_provider"] == "anthropic"


@pytest.mark.asyncio
async def test_create__fallback_workload_identity_alone_forwards(
    harness: Harness, no_openai_creds: None, monkeypatch: pytest.MonkeyPatch, tmp_path: Path
):
    token_file = tmp_path / "subject_token.jwt"
    token_file.write_text("subject-token-from-file")
    monkeypatch.setenv("OPENAI_IDENTITY_PROVIDER_ID", "idp_test123")
    monkeypatch.setenv("OPENAI_SERVICE_ACCOUNT_ID", "user-test456")
    monkeypatch.setenv("OPENAI_IDENTITY_TOKEN_FILE", str(token_file))
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    await call_create(harness)

    assert harness.acreate_kwargs()["custom_llm_provider"] == "openai"


@pytest.mark.asyncio
async def test_create__fallback_env_key_alone_forwards(harness, monkeypatch):
    monkeypatch.setattr(litellm, "api_key", None)
    monkeypatch.setattr(litellm, "openai_key", None)
    monkeypatch.setenv("OPENAI_API_KEY", "sk-env-openai")
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    await call_create(harness)

    assert harness.acreate_kwargs()["custom_llm_provider"] == "openai"


@pytest.mark.asyncio
async def test_create__fallback_provider_path_param(harness):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    await call_create(harness, provider="anthropic")

    harness.creds_resolver.assert_not_called()
    assert harness.acreate_kwargs()["custom_llm_provider"] == "anthropic"


@pytest.mark.asyncio
async def test_create__fallback_body_custom_llm_provider(harness):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
            "custom_llm_provider": "bedrock",
        },
    )

    await call_create(harness)

    payload = harness.acreate_kwargs()
    assert payload["custom_llm_provider"] == "bedrock"


# =========================================================================== #
# Unified file id routing (-> llm_router). Helpers mocked only here because a
# real unified id is opaque base64; the routing contract is what we lock.
# =========================================================================== #


@pytest.mark.asyncio
async def test_create__unified_file_id_single_model_disables_cross_model_fallbacks(harness):
    set_body(
        harness,
        {
            "input_file_id": "litellm_proxy_unified_id",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
            "disable_fallbacks": False,
        },
    )
    with (
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"),
        patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gpt-4o-mini"]),
    ):
        resp = await call_create(harness)

    # DISPATCH - router fired, direct litellm did not.
    assert harness.router_acreate.call_count == 1
    harness.litellm_acreate.assert_not_called()
    # model injected from the unified id, input_file_id restored, hidden param set
    assert harness.router_kwargs()["model"] == "gpt-4o-mini"
    assert harness.router_kwargs()["disable_fallbacks"] is True
    assert resp.input_file_id == "litellm_proxy_unified_id"
    assert resp._hidden_params["unified_file_id"] == "unified-xyz"


@pytest.mark.asyncio
@pytest.mark.parametrize("models", [[], ["m1", "m2"]])
async def test_create__unified_file_id_not_exactly_one_model_400(harness, models):
    set_body(
        harness,
        {
            "input_file_id": "litellm_proxy_unified_id",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    with (
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"),
        patch.object(endpoints, "get_models_from_unified_file_id", return_value=models),
    ):
        with pytest.raises(ProxyException) as exc:
            await call_create(harness)

    assert exc.value.code == "400"
    harness.router_acreate.assert_not_called()
    harness.litellm_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_create__unified_file_id_resolves_real_storage_url(harness):
    """A base64 unified_file_id is a LiteLLM-internal token, not a real
    provider-side file reference (e.g. Vertex AI's batch transformation parses
    a `publishers/` segment out of the file URI and crashes on the opaque
    base64 string). The real backend location (`storage_url`) must be looked
    up from LiteLLM_ManagedFileTable and substituted before dispatch.

    Regression lock on the lookup key: LiteLLM_ManagedFileTable.unified_file_id
    stores the raw base64 file id (see schema.prisma and the enterprise
    managed-files hook, which queries with the raw id), NOT the decoded
    litellm_proxy:... string. Querying with the decoded string never matches
    and silently falls back."""
    set_body(
        harness,
        {
            "input_file_id": "litellm_proxy_unified_id",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    fake_db_file = MagicMock(
        storage_url="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.0/abc",
    )
    find_first = AsyncMock(return_value=fake_db_file)
    fake_repo_instance = MagicMock()
    fake_repo_instance.table.find_first = find_first
    fake_repo_cls = MagicMock(return_value=fake_repo_instance)

    with (
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"),
        patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]),
        patch.object(proxy_server, "prisma_client", MagicMock()),
        patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
    ):
        resp = await call_create(harness)

    assert harness.router_kwargs()["input_file_id"] == fake_db_file.storage_url
    find_first.assert_awaited_once_with(where={"unified_file_id": "litellm_proxy_unified_id"})
    assert resp.input_file_id == "litellm_proxy_unified_id"
    assert resp._hidden_params["unified_file_id"] == "unified-xyz"


@pytest.mark.asyncio
async def test_create__unified_file_id_db_error_falls_back_to_raw_id(harness):
    """Resolution is additive and best-effort: a lookup error leaves the id
    unresolved and dispatch falls back to the original id, exactly as before
    this change (the managed-files deployment hook still maps it). No new
    failure mode is introduced."""
    set_body(
        harness,
        {
            "input_file_id": "litellm_proxy_unified_id",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    find_first = AsyncMock(side_effect=Exception("db unavailable"))
    fake_repo_instance = MagicMock()
    fake_repo_instance.table.find_first = find_first
    fake_repo_cls = MagicMock(return_value=fake_repo_instance)

    with (
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"),
        patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]),
        patch.object(proxy_server, "prisma_client", MagicMock()),
        patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
    ):
        await call_create(harness)

    assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"


@pytest.mark.asyncio
async def test_create__multi_model_unified_file_with_loadbalancing_keeps_router_branch(harness):
    """Regression guard: a multi-model managed file dispatched with an explicit
    router model under load balancing must keep taking the load-balanced router
    branch, exactly as on the base revision, where the managed-files deployment
    hook remaps the unified id per model. Routing it into the unified branch
    instead would trip that branch's "exactly one model" 400 and break a path
    that works today, so the unified-file resolution must not steal the
    load-balanced branch."""
    set_body(
        harness,
        {
            "input_file_id": "litellm_proxy_unified_id",
            "model": "vertex-model",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    harness.is_known_model.return_value = True

    with (
        patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True),
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"),
        patch.object(endpoints, "get_models_from_unified_file_id", return_value=["model-a", "model-b"]),
    ):
        await call_create(harness)

    assert harness.router_acreate.call_count == 1
    assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"
    harness.litellm_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_create__unified_file_id_missing_row_falls_back_to_raw_id(harness):
    """Resolution is additive: when no managed-file row exists there is nothing
    to substitute, so dispatch falls back to the original id exactly as before
    this change (the managed-files deployment hook still maps it). No new
    failure mode is introduced for this case."""
    set_body(
        harness,
        {
            "input_file_id": "litellm_proxy_unified_id",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    find_first = AsyncMock(return_value=None)
    fake_repo_instance = MagicMock()
    fake_repo_instance.table.find_first = find_first
    fake_repo_cls = MagicMock(return_value=fake_repo_instance)

    with (
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"),
        patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]),
        patch.object(proxy_server, "prisma_client", MagicMock()),
        patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
    ):
        await call_create(harness)

    assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"


@pytest.mark.asyncio
async def test_create__unified_file_id_legacy_row_without_storage_url_dispatches_raw(
    harness,
):
    """A managed file whose row predates the storage_url column still dispatches
    the original id (the managed-files deployment hook maps it); the row exists,
    so this is not the missing-row fail-closed case."""
    set_body(
        harness,
        {
            "input_file_id": "litellm_proxy_unified_id",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    fake_db_file = MagicMock(storage_url=None)
    find_first = AsyncMock(return_value=fake_db_file)
    fake_repo_instance = MagicMock()
    fake_repo_instance.table.find_first = find_first
    fake_repo_cls = MagicMock(return_value=fake_repo_instance)

    with (
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"),
        patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gemini-2.0"]),
        patch.object(proxy_server, "prisma_client", MagicMock()),
        patch.object(endpoints, "ManagedFileRepository", fake_repo_cls),
    ):
        await call_create(harness)

    assert harness.router_kwargs()["input_file_id"] == "litellm_proxy_unified_id"


# --------------------------------------------------------------------------- #
# LiteLLM-executed batches: a unified file targeting a provider whose API has
# no /v1/batches (hosted_vllm) runs inside LiteLLM instead of being forwarded.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_create__unified_executed_provider_runs_inside_litellm(harness, executed_runner):
    runner, factory = executed_runner
    caller = UserAPIKeyAuth(api_key="sk-test", team_id="team-vllm")
    input_file_id = _managed_input_file_id("my-vllm")
    set_body(
        harness,
        {
            "input_file_id": input_file_id,
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
            "litellm_metadata": {"tags": ["batch-tag"]},
        },
    )
    resp = await call_create(harness, user=caller)

    harness.router_acreate.assert_not_called()
    harness.litellm_acreate.assert_not_called()
    harness.creds_resolver.assert_called_once_with(model_id="my-vllm", team_id="team-vllm")
    factory.assert_called_once_with(harness.router, harness.logging)
    runner.create.assert_awaited_once()
    create_kwargs = runner.create.call_args.kwargs
    assert create_kwargs["unified_input_file_id"] == input_file_id
    assert create_kwargs["model"] == "my-vllm"
    assert create_kwargs["provider"] == "hosted_vllm"
    assert create_kwargs["request_tags"] == ("batch-tag",)
    assert create_kwargs["user_api_key_dict"] is caller
    assert create_kwargs["create_request"]["model"] == "my-vllm"
    assert resp.id == "litellm-executed-batch"


@pytest.mark.asyncio
async def test_create__unified_executed_provider_without_database_400(harness):
    set_body(
        harness,
        {
            "input_file_id": _managed_input_file_id("my-vllm"),
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    with pytest.raises(ProxyException) as exc:
        await call_create(harness)

    assert exc.value.code == "400"
    assert "need a database" in exc.value.message
    harness.router_acreate.assert_not_called()
    harness.litellm_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_create__unified_executed_provider_with_its_own_files_api_goes_to_the_provider(harness, executed_runner):
    runner, factory = executed_runner
    harness.upstream_files_route.mock(return_value=httpx.Response(200, json={"object": "list", "data": []}))
    set_body(
        harness,
        {
            "input_file_id": _managed_input_file_id("my-vllm"),
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    await call_create(harness)

    factory.assert_not_called()
    runner.create.assert_not_called()
    assert harness.router_kwargs()["model"] == "my-vllm"


@pytest.mark.asyncio
async def test_create__unified_provider_model_never_touches_executed_runner(harness, executed_runner):
    runner, factory = executed_runner
    set_body(
        harness,
        {
            "input_file_id": _managed_input_file_id("azure/gpt-4o"),
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    await call_create(harness)

    factory.assert_not_called()
    runner.create.assert_not_called()
    harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o", team_id=None)
    assert harness.router_kwargs()["model"] == "azure/gpt-4o"


@pytest.mark.asyncio
@pytest.mark.parametrize("via", ["body", "header"])
async def test_create__raw_file_with_executed_model_400_with_upload_guidance(harness, via):
    body = {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"}
    set_body(harness, {**body, "model": "my-vllm"} if via == "body" else body)
    headers = {"x-litellm-model": "my-vllm"} if via == "header" else None

    with pytest.raises(ProxyException) as exc:
        await call_create(harness, headers=headers)

    assert exc.value.code == "400"
    assert "POST /v1/files" in exc.value.message
    assert "x-litellm-model" in exc.value.message
    harness.litellm_acreate.assert_not_called()
    harness.router_acreate.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "upstream_answer",
    [httpx.Response(200, json={"object": "list", "data": []}), httpx.Response(405), httpx.ConnectError("refused")],
    ids=["lists files", "files route without list", "unreachable"],
)
async def test_create__raw_file_with_executed_model_is_forwarded_unless_the_server_lacks_a_files_api(
    harness, upstream_answer
):
    harness.upstream_files_route.mock(side_effect=[upstream_answer])
    set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"})

    await call_create(harness, headers={"x-litellm-model": "my-vllm"})

    forwarded = harness.acreate_kwargs()
    assert forwarded["input_file_id"] == "file-plain"
    assert forwarded["custom_llm_provider"] == "hosted_vllm"
    assert forwarded["api_base"] == CREDS["my-vllm"]["api_base"]


@pytest.mark.asyncio
async def test_create__model_encoded_beats_unified(harness):
    """Precedence row: a file id that is BOTH model-encoded and (pretend) unified
    must take the model-encoded branch (checked first)."""
    set_body(
        harness,
        {
            "input_file_id": AZURE_FILE_ID,
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    with (
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"),
        patch.object(endpoints, "get_models_from_unified_file_id", return_value=["something-else"]),
    ):
        await call_create(harness)

    assert harness.litellm_acreate.call_count == 1
    harness.router_acreate.assert_not_called()
    harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")


# =========================================================================== #
# Loadbalancing branch (-> llm_router) and its precedence vs model-encoded.
# =========================================================================== #


@pytest.mark.asyncio
async def test_create__loadbalancing_routes_to_router(harness):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
            "model": "lb-model",
        },
    )
    harness.is_known_model.return_value = True
    with patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True):
        await call_create(harness)

    harness.is_known_model.assert_called_once_with(model="lb-model", llm_router=harness.router)
    assert harness.router_acreate.call_count == 1
    harness.litellm_acreate.assert_not_called()
    harness.creds_resolver.assert_not_called()


@pytest.mark.asyncio
async def test_create__model_encoded_beats_loadbalancing(harness):
    set_body(
        harness,
        {
            "input_file_id": AZURE_FILE_ID,
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
            "model": "lb-model",
        },
    )
    harness.is_known_model.return_value = True
    with patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True):
        await call_create(harness)

    assert harness.litellm_acreate.call_count == 1
    harness.router_acreate.assert_not_called()
    harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "body, missing_param",
    [
        ({"endpoint": "/v1/chat/completions", "completion_window": "24h"}, "input_file_id"),
        ({"input_file_id": "file-abc", "completion_window": "24h"}, "endpoint"),
        ({"input_file_id": "file-abc", "endpoint": "/v1/chat/completions"}, "completion_window"),
        ({}, "input_file_id"),
    ],
)
async def test_create__missing_required_param_is_400(harness, body, missing_param):
    set_body(harness, body)

    with pytest.raises(ProxyException) as exc_info:
        await call_create(harness)

    assert exc_info.value.code == "400"
    assert exc_info.value.type == "invalid_request_error"
    assert exc_info.value.param == missing_param
    assert exc_info.value.message == f"/batches: Missing required parameter: '{missing_param}'."
    harness.litellm_acreate.assert_not_called()
    harness.router_acreate.assert_not_called()


def _raw_batches_request(body: Dict[str, Any]) -> MagicMock:
    """A request that reaches the real pre-call logic, which the `harness` fixture
    mocks out. Metadata validation lives there, so it cannot be seen through the seam."""
    request = MagicMock(spec=Request)
    request.url = MagicMock()
    request.url.__str__.return_value = "http://localhost/v1/batches"
    request.url.path = "/v1/batches"
    request.method = "POST"
    request.scope = {"type": "http", "method": "POST", "path": "/v1/batches"}
    request.query_params = {}
    request.headers = {"Content-Type": "application/json"}
    request.client = MagicMock()
    request.client.host = "127.0.0.1"
    request.body = AsyncMock(return_value=json.dumps(body).encode())
    return request


@pytest.mark.asyncio
@pytest.mark.parametrize("field", ["metadata", "litellm_metadata"])
async def test_create__non_object_metadata_is_400(field):
    """A non-object metadata field is rejected with a 400 naming it (#37147), rather
    than being coerced to None and silently dropped, or crashing behind a 500 with
    "'str' object has no attribute 'update'"."""
    body = {
        "input_file_id": "file-abc",
        "endpoint": "/v1/chat/completions",
        "completion_window": "24h",
        field: "abc",
    }

    with pytest.raises(ProxyException) as exc_info:
        await endpoints.create_batch(
            request=_raw_batches_request(body),
            fastapi_response=Response(),
            provider=None,
            user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"),
        )

    assert exc_info.value.code == "400"
    assert exc_info.value.param == field
    assert "has no attribute 'update'" not in exc_info.value.message


# =========================================================================== #
# Team-level batch expiry enforcement (independent of routing).
# =========================================================================== #


def _user_with_expiry(expiry: Any) -> UserAPIKeyAuth:
    return UserAPIKeyAuth(
        api_key="sk-test",
        team_metadata={"enforced_batch_output_expires_after": expiry},
    )


@pytest.mark.asyncio
async def test_create__team_expiry_injected(harness, openai_env_creds):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    await call_create(harness, user=_user_with_expiry({"anchor": "created_at", "seconds": 3600}))

    assert harness.acreate_kwargs()["output_expires_after"] == {
        "anchor": "created_at",
        "seconds": 3600,
    }


@pytest.mark.asyncio
async def test_create__no_team_expiry_not_injected(harness, openai_env_creds):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    await call_create(harness, user=UserAPIKeyAuth(api_key="sk-test"))

    assert "output_expires_after" not in harness.acreate_kwargs()


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "expiry",
    [
        {"seconds": 3600},  # missing anchor
        {"anchor": "created_at"},  # missing seconds
        {"anchor": "completed_at", "seconds": 3600},  # wrong anchor
    ],
)
async def test_create__team_expiry_malformed_500(harness, expiry):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    with pytest.raises(ProxyException) as exc:
        await call_create(harness, user=_user_with_expiry(expiry))

    assert exc.value.code == "500"
    harness.litellm_acreate.assert_not_called()
    harness.router_acreate.assert_not_called()


# =========================================================================== #
# Cross-cutting: enrichment route_type, metadata sanitization, failure hook.
# =========================================================================== #


@pytest.mark.asyncio
async def test_create__uses_acreate_batch_route_type(harness, openai_env_creds):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    await call_create(harness)

    assert harness.pre_call.call_args.kwargs["route_type"] == "acreate_batch"


@pytest.mark.asyncio
@pytest.mark.parametrize("loadbalancing", [True, False], ids=["load_balanced", "unified_branch"])
async def test_create__pre_call_ownership_denial_stops_every_dispatch_branch(harness, loadbalancing: bool):
    set_body(
        harness,
        {
            "input_file_id": "litellm_proxy_unified_id",
            "model": "vertex-model",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    harness.is_known_model.return_value = True
    harness.pre_call.side_effect = HTTPException(status_code=403, detail="does not have access to the file")

    with (
        patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", loadbalancing),
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"),
        patch.object(endpoints, "get_models_from_unified_file_id", return_value=["vertex-model"]),
        pytest.raises(ProxyException) as exc,
    ):
        await call_create(harness)

    assert exc.value.code == "403", exc.value
    harness.router_acreate.assert_not_called()
    harness.litellm_acreate.assert_not_called()


def install_managed_files_hook(harness: Harness) -> AsyncMock:
    prisma_client = AsyncMock()
    managed_files = PROXY_LiteLLMManagedFiles(MagicMock(async_set_cache=AsyncMock()), prisma_client=prisma_client)
    harness.logging.post_call_success_hook = AsyncMock(side_effect=managed_files.async_post_call_success_hook)
    harness.router.model_list = []
    return prisma_client


TEAM_A_KEY = UserAPIKeyAuth(api_key="sk-team-a", user_id="user_a", team_id="team_a")


def assert_ownership_registered_for_team_a(prisma_client: AsyncMock, batch_id: str) -> None:
    upsert = prisma_client.db.litellm_managedobjecttable.upsert
    upsert.assert_awaited_once()
    assert upsert.await_args.kwargs["where"] == {"unified_object_id": batch_id}
    created = upsert.await_args.kwargs["data"]["create"]
    assert created["created_by"] == "user_a"
    assert created["team_id"] == "team_a"
    prisma_client.db.litellm_managedobjecttable.update_many.assert_not_awaited()


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "body",
    [
        {"input_file_id": AZURE_FILE_ID},
        {"input_file_id": "file-plain", "model": "vertex-model"},
        {"input_file_id": "file-plain"},
    ],
    ids=["model_encoded_file_id", "model_param", "provider_fallback"],
)
async def test_create__registers_ownership_for_creator(harness, openai_env_creds, body):
    set_body(harness, {**body, "endpoint": "/v1/chat/completions", "completion_window": "24h"})
    prisma_client = install_managed_files_hook(harness)

    resp = await call_create(harness, user=TEAM_A_KEY)

    assert_ownership_registered_for_team_a(prisma_client, resp.id)


@pytest.mark.asyncio
async def test_create__unified_file_id_registers_ownership_for_creator(harness):
    unified_input_file_id = base64.urlsafe_b64encode(
        b"litellm_proxy:application/octet-stream;unified_id,input-uuid;target_model_names,gpt-4o-mini"
    ).decode()
    set_body(
        harness,
        {
            "input_file_id": unified_input_file_id,
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    prisma_client = install_managed_files_hook(harness)

    resp = await call_create(harness, user=TEAM_A_KEY)

    assert harness.router_acreate.call_count == 1
    assert_ownership_registered_for_team_a(prisma_client, resp.id)


@pytest.mark.asyncio
async def test_create__metadata_sanitized_before_forwarding(harness, openai_env_creds):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
            "metadata": {"user_key": "user_val", "spend_logs_metadata": {"x": 1}},
        },
    )

    await call_create(harness)

    # provider-internal-only key dropped, string key kept (real sanitize runs)
    assert harness.acreate_kwargs()["metadata"] == {"user_key": "user_val"}


@pytest.mark.asyncio
async def test_create__exception_calls_failure_hook(harness, openai_env_creds):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    harness.litellm_acreate.side_effect = ValueError("provider boom")

    with pytest.raises(ProxyException):
        await call_create(harness)

    harness.logging.post_call_failure_hook.assert_called_once()
    assert harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom"


async def test_create__exception_carries_the_litellm_call_id(harness, openai_env_creds, caplog):
    call_id = "lit7836-batch-call-id"
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
            "litellm_call_id": call_id,
        },
    )
    harness.litellm_acreate.side_effect = ValueError("provider boom")

    with caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), pytest.raises(ProxyException) as raised:
        await call_create(harness)

    assert raised.value.headers["x-litellm-call-id"] == call_id
    record = next(r for r in caplog.records if "Exception occured" in r.getMessage())
    assert record.litellm_call_id == call_id
    assert call_id in record.getMessage()


# =========================================================================== #
#                                                                             #
#   GET /v1/batches/{batch_id}  -  retrieve_batch routing-contract tests      #
#                                                                             #
#   Same discipline as create_batch above. retrieve_batch has more seams: it  #
#   first consults the ManagedObjectTable (get_batch_from_database) and may    #
#   short-circuit on a terminal-status row WITHOUT ever calling a provider,    #
#   then on a miss/non-terminal row routes to one of three downstream seams    #
#   (litellm.aretrieve_batch via model creds, llm_router.aretrieve_batch, or   #
#   litellm.aretrieve_batch via env-var provider) and writes the fresh state   #
#   back via update_batch_in_database. Every test below locks exactly which    #
#   of those seams fired and asserts the siblings did NOT, so a reordered or   #
#   negated branch - or a dropped DB short-circuit / write-back - fails loud.  #
#                                                                             #
#   The DB seams (get_batch_from_database, update_batch_in_database,           #
#   resolve_*_to_unified) are true prisma I/O boundaries and are mocked. The   #
#   encode/decode/credential-merge helpers run for real, so payload and id     #
#   round-trip assertions reflect production exactly.                          #
# =========================================================================== #


# A real model-encoded BATCH id: decodes to "azure/gpt-4o", strips to
# "batch_orig123". Distinct from AZURE_FILE_ID so retrieve tests can't pass by
# accidentally reusing the create fixture's value.
AZURE_BATCH_ID = encode_file_id_with_model("batch_orig123", "azure/gpt-4o", id_type="batch")

# A realistic decoded unified batch id (what _is_base64_encoded_unified_file_id
# returns). model_id / llm_batch_id are parsed out of this by the real helpers.
UNIFIED_BATCH_ID = "litellm_proxy;model_id:gpt-4o-mini;llm_batch_id:batch-raw-xyz"

# A decoded unified id of a batch LiteLLM runs itself: the llm_batch_id carries
# the litellm_batch_ prefix, so no provider holds a batch to sync with.
EXECUTED_BATCH_ID = "litellm_proxy;model_id:my-vllm;llm_batch_id:litellm_batch_abc"
EXECUTED_BATCH_B64 = base64.urlsafe_b64encode(EXECUTED_BATCH_ID.encode()).decode().rstrip("=")


@dataclass
class RetrieveHarness:
    """Seams for retrieve_batch. `data['data']` is the dict pre-call enrichment
    returns; the routing branches mutate it, so it is reset per call."""

    data: Dict[str, Any]
    pre_call: AsyncMock
    get_headers: MagicMock
    provider_from_headers: MagicMock
    provider_from_query: MagicMock
    litellm_aretrieve: AsyncMock
    router: MagicMock
    logging: MagicMock
    creds_resolver: MagicMock
    get_batch_from_db: AsyncMock
    update_batch_in_db: AsyncMock
    ensure_managed_files: AsyncMock

    @property
    def router_aretrieve(self) -> AsyncMock:
        return self.router.aretrieve_batch

    def aretrieve_kwargs(self) -> Dict[str, Any]:
        """Exact kwargs forwarded to litellm.aretrieve_batch."""
        assert self.litellm_aretrieve.call_count == 1
        return dict(self.litellm_aretrieve.call_args.kwargs)

    def router_kwargs(self) -> Dict[str, Any]:
        assert self.router_aretrieve.call_count == 1
        return dict(self.router_aretrieve.call_args.kwargs)


@pytest.fixture
def retrieve_harness():
    data_holder: Dict[str, Any] = {"data": {}}
    logging = MagicMock(spec=ProxyLogging)
    logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"])
    logging.post_call_failure_hook = AsyncMock()
    logging.update_request_status = AsyncMock()
    logging.get_proxy_hook = MagicMock(return_value=None)

    router = MagicMock(spec=Router)
    router.model_group_alias = {}
    router.get_model_access_groups = MagicMock(return_value={})
    router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id)
    router.model_list = []
    router.aretrieve_batch = AsyncMock(return_value=make_batch())
    router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)
    router.get_credential_deployment = MagicMock(return_value=None)

    pre_call = AsyncMock(side_effect=lambda **kw: (data_holder["data"], MagicMock()))
    get_headers = MagicMock(return_value={})
    provider_from_headers = MagicMock(return_value=None)
    provider_from_query = MagicMock(return_value=None)
    litellm_aretrieve = AsyncMock(return_value=make_batch())
    # Default: DB miss -> always fall through to provider routing.
    get_batch_from_db = AsyncMock(return_value=(None, None))
    update_batch_in_db = AsyncMock(return_value=None)
    ensure_managed_files = AsyncMock(return_value=None)

    with ExitStack() as stack:
        stack.enter_context(
            patch.object(
                ProxyBaseLLMRequestProcessing,
                "common_processing_pre_call_logic",
                pre_call,
            )
        )
        stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
        stack.enter_context(
            patch.object(
                endpoints,
                "get_custom_llm_provider_from_request_headers",
                provider_from_headers,
            )
        )
        stack.enter_context(
            patch.object(
                endpoints,
                "get_custom_llm_provider_from_request_query",
                provider_from_query,
            )
        )
        stack.enter_context(patch.object(endpoints, "get_batch_from_database", get_batch_from_db))
        stack.enter_context(patch.object(endpoints, "update_batch_in_database", update_batch_in_db))
        stack.enter_context(patch.object(endpoints, "ensure_batch_response_managed_file_ids", ensure_managed_files))
        stack.enter_context(patch.object(litellm, "aretrieve_batch", litellm_aretrieve))
        stack.enter_context(patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False))
        stack.enter_context(patch.object(proxy_server, "llm_router", router))
        stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging))
        stack.enter_context(patch.object(proxy_server, "general_settings", {}))
        stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock()))
        stack.enter_context(patch.object(proxy_server, "version", "test-version"))
        stack.enter_context(patch.object(proxy_server, "prisma_client", MagicMock()))

        yield RetrieveHarness(
            data=data_holder,
            pre_call=pre_call,
            get_headers=get_headers,
            provider_from_headers=provider_from_headers,
            provider_from_query=provider_from_query,
            litellm_aretrieve=litellm_aretrieve,
            router=router,
            logging=logging,
            creds_resolver=router.get_deployment_credentials_with_provider,
            get_batch_from_db=get_batch_from_db,
            update_batch_in_db=update_batch_in_db,
            ensure_managed_files=ensure_managed_files,
        )


async def call_retrieve(
    harness: RetrieveHarness,
    batch_id: str,
    *,
    provider: Optional[str] = None,
    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=...),
    # 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(),
        user_api_key_dict=user or UserAPIKeyAuth(api_key="sk-test"),
        provider=provider,
        batch_id=batch_id,
    )


# --------------------------------------------------------------------------- #
# SCENARIO 1 - batch id encoded with model. litellm.aretrieve_batch via the
# model's resolved credentials; response ids re-encoded for the client.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_retrieve__model_encoded_id(retrieve_harness):
    resp = await call_retrieve(retrieve_harness, AZURE_BATCH_ID)

    # 1. DISPATCH - model-credential path fired via litellm, router did not.
    assert retrieve_harness.litellm_aretrieve.call_count == 1
    retrieve_harness.router_aretrieve.assert_not_called()

    # 2. CREDENTIALS - resolved for the model decoded FROM the batch id.
    retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")

    # 3. SEAM PAYLOAD - exact, whole dict forwarded to the provider call.
    #    Note `model` is the DECODED model, not the deployment from creds: the
    #    endpoint overrides it (provider-config providers like bedrock need it).
    assert retrieve_harness.aretrieve_kwargs() == {
        "custom_llm_provider": "azure",
        "batch_id": "batch_orig123",  # encoding stripped by this layer
        "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.
    assert resp.id == encode_file_id_with_model("batch-provider-id", "azure/gpt-4o", id_type="batch")

    # write-back to the managed-object table happened, tagged as a retrieve.
    assert retrieve_harness.update_batch_in_db.call_count == 1
    assert retrieve_harness.update_batch_in_db.call_args.kwargs["operation"] == "retrieve"


@pytest.mark.asyncio
async def test_retrieve__model_encoded_id__stamps_deployment_model_info_for_cost(retrieve_harness):
    """Regression: this path calls litellm.aretrieve_batch directly, so nothing stamped the
    deployment's model_info the way the router does for routed calls. Cost tracking then never
    saw the deployment id, and a completed batch on a deployment with its own per-page pricing
    was billed at the published rate with an empty model_id on the spend row."""
    retrieve_harness.router.get_credential_deployment.return_value = Deployment(
        model_name="azure-gpt",
        litellm_params=LiteLLM_Params(model="azure/gpt-4o"),
        model_info=ModelInfo(id="dep-123"),
    )
    retrieve_harness.pre_call.side_effect = lambda **kw: (
        {**retrieve_harness.data["data"], "litellm_metadata": {"user_api_key_alias": "qa-key"}},
        MagicMock(),
    )

    await call_retrieve(retrieve_harness, AZURE_BATCH_ID)

    retrieve_harness.router.get_credential_deployment.assert_called_once_with(model_id="azure/gpt-4o")
    litellm_metadata = retrieve_harness.aretrieve_kwargs()["litellm_metadata"]
    assert litellm_metadata["model_info"]["id"] == "dep-123"
    assert litellm_metadata["user_api_key_alias"] == "qa-key"


@pytest.mark.asyncio
async def test_retrieve__model_encoded_id__forwards_decoded_model_not_deployment(
    retrieve_harness,
):
    """Regression guard for the line-483 override: the model forwarded to the
    provider must be the decoded model id, never the deployment name that the
    credential merge pulled in. Dropping the override silently 400s bedrock."""
    await call_retrieve(retrieve_harness, AZURE_BATCH_ID)

    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,
):
    retrieve_harness.litellm_aretrieve.return_value = make_batch(
        id="batch-xyz",
        output_file_id="file-out-raw",
        error_file_id="file-err-raw",
    )

    resp = await call_retrieve(retrieve_harness, AZURE_BATCH_ID)

    assert resp.output_file_id == encode_file_id_with_model("file-out-raw", "azure/gpt-4o")
    assert resp.error_file_id == encode_file_id_with_model("file-err-raw", "azure/gpt-4o")


@pytest.mark.asyncio
async def test_retrieve__model_encoded_beats_loadbalancing(retrieve_harness):
    """Precedence: model-encoded id is checked before the loadbalancing/unified
    elif, so it wins even with loadbalancing enabled."""
    with patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True):
        await call_retrieve(retrieve_harness, AZURE_BATCH_ID)

    assert retrieve_harness.litellm_aretrieve.call_count == 1
    retrieve_harness.router_aretrieve.assert_not_called()
    retrieve_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")


# --------------------------------------------------------------------------- #
# Unified managed batch id -> llm_router.aretrieve_batch. model_id is parsed
# out of the unified id and stamped onto hidden params; raw file ids on the
# response are resolved back to unified ids.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_retrieve__unified_batch_id_routes_to_router(retrieve_harness):
    with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
        resp = await call_retrieve(retrieve_harness, "batch-unified-blob")

    # DISPATCH - router fired, direct litellm did not.
    assert retrieve_harness.router_aretrieve.call_count == 1
    retrieve_harness.litellm_aretrieve.assert_not_called()

    # Credentials are resolved for the deployment behind the unified id so the batch's
    # output file can be read for cost accounting. This id resolves to nothing here, and
    # the retrieve must still serve the batch rather than fail on the lookup.
    retrieve_harness.creds_resolver.assert_called_once_with(model_id="gpt-4o-mini")

    # router receives the (still-encoded) batch id verbatim - this layer does
    # not decode it for the unified path.
    assert retrieve_harness.router_kwargs() == {"batch_id": "batch-unified-blob"}

    # hidden params: unified id passed through, model_id parsed from it.
    assert resp._hidden_params["unified_batch_id"] == UNIFIED_BATCH_ID
    assert resp._hidden_params["model_id"] == "gpt-4o-mini"

    # raw provider file ids on the response are normalized to managed unified ids.
    retrieve_harness.ensure_managed_files.assert_called_once()


@pytest.mark.asyncio
async def test_retrieve__loadbalancing_raw_id_routes_to_router(retrieve_harness):
    """Loadbalancing on + a plain (non-encoded, non-unified) batch id routes to
    the router. Locks the current dispatch contract of the shared elif."""
    with patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True):
        resp = await call_retrieve(retrieve_harness, "batch-raw-xyz")

    assert retrieve_harness.router_aretrieve.call_count == 1
    retrieve_harness.litellm_aretrieve.assert_not_called()
    assert retrieve_harness.router_kwargs() == {"batch_id": "batch-raw-xyz"}
    # not a unified id -> hidden param reflects that, no model_id stamped.
    assert resp._hidden_params["unified_batch_id"] is False
    assert "model_id" not in resp._hidden_params
    # not a unified id -> no managed-file normalization.
    retrieve_harness.ensure_managed_files.assert_not_called()


# --------------------------------------------------------------------------- #
# SCENARIO 3 - fallback to custom_llm_provider (env-var creds). MUST NOT touch
# the credential resolver or the router.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_retrieve__fallback_default_openai(retrieve_harness, openai_env_creds):
    await call_retrieve(retrieve_harness, "batch-raw-xyz")

    assert retrieve_harness.litellm_aretrieve.call_count == 1
    retrieve_harness.router_aretrieve.assert_not_called()
    retrieve_harness.creds_resolver.assert_not_called()  # inverse-bug guard
    assert retrieve_harness.aretrieve_kwargs() == {
        "custom_llm_provider": "openai",
        "batch_id": "batch-raw-xyz",
    }
    assert retrieve_harness.update_batch_in_db.call_count == 1


@pytest.mark.asyncio
async def test_retrieve__fallback_no_creds_404(retrieve_harness, no_openai_creds):
    with pytest.raises(ProxyException) as exc:
        await call_retrieve(retrieve_harness, "batch-raw-xyz")

    assert exc.value.code == "404"
    assert exc.value.type == "invalid_request_error"
    assert exc.value.param is None
    assert exc.value.message == "No batch found with id 'batch-raw-xyz'."
    retrieve_harness.litellm_aretrieve.assert_not_called()
    retrieve_harness.router_aretrieve.assert_not_called()


@pytest.mark.asyncio
async def test_retrieve__fallback_explicit_provider_bypasses_not_found_gate(retrieve_harness, no_openai_creds):
    await call_retrieve(retrieve_harness, "batch-raw-xyz", provider="anthropic")

    assert retrieve_harness.aretrieve_kwargs()["custom_llm_provider"] == "anthropic"


@pytest.mark.asyncio
async def test_retrieve__fallback_env_key_alone_forwards(retrieve_harness, monkeypatch):
    monkeypatch.setattr(litellm, "api_key", None)
    monkeypatch.setattr(litellm, "openai_key", None)
    monkeypatch.setenv("OPENAI_API_KEY", "sk-env-openai")

    await call_retrieve(retrieve_harness, "batch-raw-xyz")

    assert retrieve_harness.aretrieve_kwargs() == {
        "custom_llm_provider": "openai",
        "batch_id": "batch-raw-xyz",
    }


@pytest.mark.asyncio
async def test_retrieve__fallback_provider_path_param(retrieve_harness):
    await call_retrieve(retrieve_harness, "batch-raw-xyz", provider="anthropic")

    retrieve_harness.creds_resolver.assert_not_called()
    assert retrieve_harness.aretrieve_kwargs()["custom_llm_provider"] == "anthropic"


@pytest.mark.asyncio
async def test_retrieve__fallback_provider_from_header(retrieve_harness):
    retrieve_harness.provider_from_headers.return_value = "bedrock"

    await call_retrieve(retrieve_harness, "batch-raw-xyz")

    assert retrieve_harness.aretrieve_kwargs()["custom_llm_provider"] == "bedrock"


@pytest.mark.asyncio
async def test_retrieve__fallback_provider_from_query(retrieve_harness):
    retrieve_harness.provider_from_query.return_value = "vertex_ai"

    await call_retrieve(retrieve_harness, "batch-raw-xyz")

    assert retrieve_harness.aretrieve_kwargs()["custom_llm_provider"] == "vertex_ai"


@pytest.mark.asyncio
async def test_retrieve__fallback_provider_precedence_path_over_header(
    retrieve_harness,
):
    """provider path param beats the header-derived provider."""
    retrieve_harness.provider_from_headers.return_value = "bedrock"

    await call_retrieve(retrieve_harness, "batch-raw-xyz", provider="anthropic")

    assert retrieve_harness.aretrieve_kwargs()["custom_llm_provider"] == "anthropic"


# --------------------------------------------------------------------------- #
# ManagedObjectTable short-circuit. A terminal-status DB row is returned
# immediately - no provider call, no write-back. A non-terminal row falls
# through to a provider sync.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
@pytest.mark.parametrize("status", ["completed", "complete", "failed", "cancelled", "expired"])
async def test_retrieve__db_terminal_state_short_circuits(retrieve_harness, status):
    # "complete" is the DB-normalized alias of "completed"; it is not a valid
    # constructor literal but reaches the endpoint via a stored row, so set it
    # post-construction to exercise that exact branch.
    db_response = make_batch(id="batch-from-db", status="completed")
    db_response.status = status
    retrieve_harness.get_batch_from_db.return_value = (MagicMock(), db_response)

    resp = await call_retrieve(retrieve_harness, "batch-raw-xyz")

    # No provider seam fired, and no write-back (the row is already terminal).
    retrieve_harness.litellm_aretrieve.assert_not_called()
    retrieve_harness.router_aretrieve.assert_not_called()
    retrieve_harness.update_batch_in_db.assert_not_called()
    # The DB object is what the client gets back.
    assert resp is db_response


@pytest.mark.asyncio
async def test_retrieve__db_terminal_unified_resolves_file_ids(retrieve_harness):
    db_response = make_batch(id="batch-from-db", status="completed")
    db_batch_object = MagicMock()
    retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response)

    with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
        await call_retrieve(retrieve_harness, "batch-unified-blob")

    # Terminal short-circuit still registers/normalizes raw provider file ids.
    retrieve_harness.ensure_managed_files.assert_called_once()
    ensure_kwargs = retrieve_harness.ensure_managed_files.call_args.kwargs
    assert ensure_kwargs["response"] is db_response
    assert ensure_kwargs["db_batch_object"] is db_batch_object
    assert ensure_kwargs["unified_batch_id"] == UNIFIED_BATCH_ID
    retrieve_harness.litellm_aretrieve.assert_not_called()
    retrieve_harness.router_aretrieve.assert_not_called()


@pytest.mark.asyncio
async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harness, openai_env_creds):
    """A non-terminal DB row must NOT short-circuit; the endpoint syncs with the
    provider to refresh state."""
    db_response = make_batch(id="batch-from-db", status="validating")
    retrieve_harness.get_batch_from_db.return_value = (MagicMock(), db_response)

    await call_retrieve(retrieve_harness, "batch-raw-xyz")

    # Provider sync happened despite the DB hit.
    assert retrieve_harness.litellm_aretrieve.call_count == 1
    assert retrieve_harness.update_batch_in_db.call_count == 1


@pytest.mark.asyncio
@pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"])
async def test_retrieve__executed_batch_served_from_db_in_every_status(retrieve_harness, status):
    db_response = make_batch(id="litellm-executed-batch", status=status)
    db_batch_object = MagicMock()
    db_batch_object.updated_at = datetime.now(timezone.utc)
    retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response)

    resp = await call_retrieve(retrieve_harness, EXECUTED_BATCH_B64)

    assert resp is db_response
    retrieve_harness.litellm_aretrieve.assert_not_called()
    retrieve_harness.router_aretrieve.assert_not_called()
    retrieve_harness.update_batch_in_db.assert_not_called()
    retrieve_harness.ensure_managed_files.assert_called_once()
    assert retrieve_harness.ensure_managed_files.call_args.kwargs["unified_batch_id"] == EXECUTED_BATCH_ID


@pytest.mark.asyncio
async def test_retrieve__executed_batch_abandoned_by_its_runner_is_served_failed(retrieve_harness, executed_runner):
    runner, _ = executed_runner
    failed = make_batch(id="litellm-executed-batch", status="failed")
    runner.fail_abandoned = AsyncMock(return_value=failed)
    db_response = make_batch(id="litellm-executed-batch", status="in_progress")
    db_batch_object = MagicMock()
    db_batch_object.updated_at = datetime.now(timezone.utc) - timedelta(minutes=10)
    retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response)
    user = UserAPIKeyAuth(api_key="sk-test", user_id="user-1")

    resp = await call_retrieve(retrieve_harness, EXECUTED_BATCH_B64, user=user)

    assert resp is failed
    runner.fail_abandoned.assert_awaited_once_with(db_response, user)
    retrieve_harness.litellm_aretrieve.assert_not_called()
    retrieve_harness.router_aretrieve.assert_not_called()
    retrieve_harness.ensure_managed_files.assert_called_once()


@pytest.mark.asyncio
async def test_retrieve__executed_batch_with_a_fresh_heartbeat_is_left_running(retrieve_harness, executed_runner):
    runner, _ = executed_runner
    runner.fail_abandoned = AsyncMock()
    db_response = make_batch(id="litellm-executed-batch", status="in_progress")
    db_batch_object = MagicMock()
    db_batch_object.updated_at = datetime.now(timezone.utc) - timedelta(seconds=30)
    retrieve_harness.get_batch_from_db.return_value = (db_batch_object, db_response)

    resp = await call_retrieve(retrieve_harness, EXECUTED_BATCH_B64)

    assert resp is db_response
    runner.fail_abandoned.assert_not_awaited()


@pytest.mark.asyncio
async def test_retrieve__executed_batch_without_db_row_404(retrieve_harness):
    with pytest.raises(ProxyException) as exc:
        await call_retrieve(retrieve_harness, EXECUTED_BATCH_B64)

    assert exc.value.code == "404"
    retrieve_harness.litellm_aretrieve.assert_not_called()
    retrieve_harness.router_aretrieve.assert_not_called()


# --------------------------------------------------------------------------- #
# Cross-cutting: enrichment route_type and failure-hook on provider error.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_retrieve__uses_aretrieve_batch_route_type(retrieve_harness, openai_env_creds):
    await call_retrieve(retrieve_harness, "batch-raw-xyz")

    assert retrieve_harness.pre_call.call_args.kwargs["route_type"] == "aretrieve_batch"


@pytest.mark.asyncio
async def test_retrieve__exception_calls_failure_hook(retrieve_harness, openai_env_creds):
    retrieve_harness.litellm_aretrieve.side_effect = ValueError("provider boom")

    with pytest.raises(ProxyException):
        await call_retrieve(retrieve_harness, "batch-raw-xyz")

    retrieve_harness.logging.post_call_failure_hook.assert_called_once()
    assert (
        retrieve_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0]
        == "provider boom"
    )


# =========================================================================== #
#                                                                             #
#   GET /v1/batches  -  list_batches routing-contract tests                    #
#                                                                             #
#   Branch order (first match wins):                                          #
#     1. managed_files hook present  -> managed_files_obj.list_user_batches    #
#     2. model from body/query/header -> litellm.alist_batches + id encode     #
#     3. target_model_names (param or body) -> llm_router.alist_batches        #
#     4. fallback -> litellm.alist_batches via env-var custom_llm_provider     #
#                                                                             #
#   llm_router is required; absence is a 500 before any branch runs.          #
# =========================================================================== #


class FakeListPage:
    """Stand-in for the SyncCursorPage[Batch] that alist_batches returns. The
    endpoint only touches `.data` (to encode ids) and `._hidden_params`."""

    def __init__(self, data: Any):
        self.data = data
        self._hidden_params: Dict[str, Any] = {}


@dataclass
class ListHarness:
    body: Dict[str, Any]
    read_body: AsyncMock
    pre_call: AsyncMock
    get_headers: MagicMock
    provider_from_headers: MagicMock
    provider_from_query: MagicMock
    litellm_alist: AsyncMock
    router: MagicMock
    logging: MagicMock
    creds_resolver: MagicMock

    @property
    def router_alist(self) -> AsyncMock:
        return self.router.alist_batches

    def set_managed_files(self, page: Any) -> AsyncMock:
        """Install a managed_files hook exposing list_user_batches -> page."""
        hook = MagicMock()
        hook.list_user_batches = AsyncMock(return_value=page)
        self.logging.get_proxy_hook = MagicMock(return_value=hook)
        return hook.list_user_batches

    def alist_kwargs(self) -> Dict[str, Any]:
        assert self.litellm_alist.call_count == 1
        return dict(self.litellm_alist.call_args.kwargs)

    def router_kwargs(self) -> Dict[str, Any]:
        assert self.router_alist.call_count == 1
        return dict(self.router_alist.call_args.kwargs)


@pytest.fixture
def list_harness():
    body_holder: Dict[str, Any] = {"body": {}}
    logging = MagicMock(spec=ProxyLogging)
    logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"])
    logging.post_call_failure_hook = AsyncMock()
    logging.update_request_status = AsyncMock()
    # Default: no managed_files hook -> branches 2/3/4 are reachable.
    logging.get_proxy_hook = MagicMock(return_value=None)

    router = MagicMock(spec=Router)
    router.model_group_alias = {}
    router.get_model_access_groups = MagicMock(return_value={})
    router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id)
    router.model_list = []
    router.alist_batches = AsyncMock(return_value=FakeListPage([]))
    router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)

    read_body = AsyncMock(side_effect=lambda request: body_holder["body"])
    pre_call = AsyncMock(side_effect=lambda **kw: (body_holder["body"], MagicMock()))
    get_headers = MagicMock(return_value={})
    provider_from_headers = MagicMock(return_value=None)
    provider_from_query = MagicMock(return_value=None)
    litellm_alist = AsyncMock(return_value=FakeListPage([]))

    with ExitStack() as stack:
        stack.enter_context(patch.object(endpoints, "read_request_body", read_body))
        stack.enter_context(
            patch.object(
                ProxyBaseLLMRequestProcessing,
                "common_processing_pre_call_logic",
                pre_call,
            )
        )
        stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
        stack.enter_context(
            patch.object(
                endpoints,
                "get_custom_llm_provider_from_request_headers",
                provider_from_headers,
            )
        )
        stack.enter_context(
            patch.object(
                endpoints,
                "get_custom_llm_provider_from_request_query",
                provider_from_query,
            )
        )
        stack.enter_context(patch.object(litellm, "alist_batches", litellm_alist))
        stack.enter_context(patch.object(proxy_server, "llm_router", router))
        stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging))
        stack.enter_context(patch.object(proxy_server, "general_settings", {}))
        stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock()))
        stack.enter_context(patch.object(proxy_server, "version", "test-version"))

        yield ListHarness(
            body=body_holder,
            read_body=read_body,
            pre_call=pre_call,
            get_headers=get_headers,
            provider_from_headers=provider_from_headers,
            provider_from_query=provider_from_query,
            litellm_alist=litellm_alist,
            router=router,
            logging=logging,
            creds_resolver=router.get_deployment_credentials_with_provider,
        )


async def call_list(
    harness: ListHarness,
    *,
    provider: Optional[str] = None,
    limit: Optional[int] = None,
    after: Optional[str] = None,
    target_model_names: Optional[str] = None,
    user: Optional[UserAPIKeyAuth] = None,
    headers: Optional[Dict[str, str]] = None,
    query: Optional[Dict[str, str]] = None,
    body: Optional[Dict[str, Any]] = None,
):
    harness.body["body"] = body if body is not None else {}
    return await endpoints.list_batches(
        request=FakeRequest(headers=headers, query=query),
        fastapi_response=Response(),
        provider=provider,
        limit=limit,
        after=after,
        user_api_key_dict=user or UserAPIKeyAuth(api_key="sk-test"),
        target_model_names=target_model_names,
    )


# --------------------------------------------------------------------------- #
# Branch 1 - ManagedObjectTable listing. This is the default production path
# (the managed_files hook is registered) and wins over every other branch.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_list__managed_files_path(list_harness):
    page = FakeListPage([make_batch(id="batch-1")])
    list_user_batches = list_harness.set_managed_files(page)

    user = UserAPIKeyAuth(api_key="sk-test")
    resp = await call_list(
        list_harness,
        user=user,
        limit=7,
        after="batch-cursor",
        provider="openai",
        target_model_names="m1,m2",
    )

    # DISPATCH - managed-files seam fired, neither provider seam did.
    list_user_batches.assert_called_once_with(
        user_api_key_dict=user,
        limit=7,
        after="batch-cursor",
        provider="openai",
        target_model_names="m1,m2",
        llm_router=list_harness.router,
    )
    list_harness.litellm_alist.assert_not_called()
    list_harness.router_alist.assert_not_called()
    assert resp is page


@pytest.mark.asyncio
async def test_list__managed_files_beats_model_param(list_harness):
    """Branch 1 is checked before the model branch: a model in the body does not
    divert away from managed-files listing."""
    page = FakeListPage([])
    list_user_batches = list_harness.set_managed_files(page)

    await call_list(list_harness, body={"model": "azure/gpt-4o"})

    list_user_batches.assert_called_once()
    list_harness.litellm_alist.assert_not_called()
    list_harness.router_alist.assert_not_called()
    list_harness.creds_resolver.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "limit, expected_message, expected_openai_code",
    [
        (
            -1,
            "Invalid 'limit': integer below minimum value. Expected a value >= 0, but got -1 instead.",
            "integer_below_min_value",
        ),
        (
            101,
            "Invalid 'limit': integer above maximum value. Expected a value <= 100, but got 101 instead.",
            "integer_above_max_value",
        ),
        (
            1000,
            "Invalid 'limit': integer above maximum value. Expected a value <= 100, but got 1000 instead.",
            "integer_above_max_value",
        ),
    ],
)
async def test_list__out_of_range_limit_rejected_with_400(list_harness, limit, expected_message, expected_openai_code):
    """OpenAI parity: GET /v1/batches rejects limit < 0 and limit > 100 with an
    OpenAI-shaped 400 before any listing branch runs (issue #37149)."""
    list_user_batches = list_harness.set_managed_files(FakeListPage([]))

    with pytest.raises(ProxyException) as exc:
        await call_list(list_harness, limit=limit)

    assert exc.value.code == "400"
    assert exc.value.param == "limit"
    assert exc.value.type == "invalid_request_error"
    assert exc.value.openai_code == expected_openai_code
    assert exc.value.message == expected_message
    list_user_batches.assert_not_called()
    list_harness.litellm_alist.assert_not_called()
    list_harness.router_alist.assert_not_called()


@pytest.mark.asyncio
@pytest.mark.parametrize("limit", [None, 0, 1, 100])
async def test_list__in_range_limit_dispatches(list_harness, limit):
    """OpenAI parity: live OpenAI accepts limit=0 (empty page) and 1..100, so
    those values must keep flowing through to the listing branch untouched."""
    page = FakeListPage([])
    list_user_batches = list_harness.set_managed_files(page)

    resp = await call_list(list_harness, limit=limit)

    assert resp is page
    assert list_user_batches.call_args.kwargs["limit"] == limit


# --------------------------------------------------------------------------- #
# Branch 2 - model from body/query/header. The endpoint resolves credentials
# for the body model, forwards custom_llm_provider once (it pops it from data
# via prepare_data_with_credentials the way create/retrieve do), and encodes
# the response ids. Regression guard for the duplicate-kwarg
# "multiple values for keyword argument 'custom_llm_provider'" bug.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_list__model_from_body_routes_and_encodes(list_harness):
    list_harness.litellm_alist.return_value = FakeListPage([make_batch(id="batch-1"), make_batch(id="batch-2")])

    resp = await call_list(list_harness, body={"model": "azure/gpt-4o"})

    assert list_harness.litellm_alist.call_count == 1
    list_harness.router_alist.assert_not_called()
    list_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")
    assert resp.data[0].id == encode_file_id_with_model("batch-1", "azure/gpt-4o", id_type="batch")
    assert resp.data[1].id == encode_file_id_with_model("batch-2", "azure/gpt-4o", id_type="batch")


# --------------------------------------------------------------------------- #
# Branch 3 - target_model_names (function param or body) -> llm_router. Routes
# to the FIRST model in the comma list; `model` is stripped from data first.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_list__target_model_names_param_routes_to_router(list_harness):
    await call_list(list_harness, target_model_names="m1,m2", limit=3, after="cur")

    assert list_harness.router_alist.call_count == 1
    list_harness.litellm_alist.assert_not_called()
    list_harness.creds_resolver.assert_not_called()
    # first model only; after/limit forwarded; nothing else (param not in data).
    assert list_harness.router_kwargs() == {
        "model": "m1",
        "after": "cur",
        "limit": 3,
    }


@pytest.mark.asyncio
async def test_list__target_model_names_from_body(list_harness):
    await call_list(list_harness, body={"target_model_names": "m1,m2"})

    assert list_harness.router_alist.call_count == 1
    list_harness.litellm_alist.assert_not_called()
    kwargs = list_harness.router_kwargs()
    assert kwargs["model"] == "m1"
    # body-sourced target_model_names stays in the forwarded data.
    assert kwargs["target_model_names"] == "m1,m2"


@pytest.mark.asyncio
async def test_list__target_model_names_takes_first_only(list_harness):
    """Locks the current behavior: with multiple target models, only the first
    is routed to (silently, unlike create which 400s on >1). A change here -
    intentional or not - must update this test."""
    await call_list(list_harness, target_model_names="alpha,beta,gamma")

    assert list_harness.router_kwargs()["model"] == "alpha"


# --------------------------------------------------------------------------- #
# Branch 4 - fallback to custom_llm_provider (env-var creds). MUST NOT touch
# the credential resolver or the router.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_list__fallback_default_openai(list_harness, no_openai_creds):
    """list stays ungated by the no-creds 404 guard: it answers about a
    collection, not a specific id, so there is nothing to 404 about."""
    await call_list(list_harness)

    assert list_harness.litellm_alist.call_count == 1
    list_harness.router_alist.assert_not_called()
    list_harness.creds_resolver.assert_not_called()  # inverse-bug guard
    assert list_harness.alist_kwargs() == {
        "custom_llm_provider": "openai",
        "after": None,
        "limit": None,
    }


@pytest.mark.asyncio
async def test_list__fallback_provider_path_param(list_harness):
    await call_list(list_harness, provider="anthropic")

    list_harness.creds_resolver.assert_not_called()
    assert list_harness.alist_kwargs()["custom_llm_provider"] == "anthropic"


@pytest.mark.asyncio
async def test_list__fallback_provider_from_header(list_harness):
    list_harness.provider_from_headers.return_value = "bedrock"

    await call_list(list_harness)

    assert list_harness.alist_kwargs()["custom_llm_provider"] == "bedrock"


@pytest.mark.asyncio
async def test_list__fallback_provider_from_query(list_harness):
    list_harness.provider_from_query.return_value = "vertex_ai"

    await call_list(list_harness)

    assert list_harness.alist_kwargs()["custom_llm_provider"] == "vertex_ai"


@pytest.mark.asyncio
async def test_list__fallback_after_and_limit_forwarded(list_harness):
    await call_list(list_harness, after="cursor-9", limit=42)

    kwargs = list_harness.alist_kwargs()
    assert kwargs["after"] == "cursor-9"
    assert kwargs["limit"] == 42


# --------------------------------------------------------------------------- #
# Cross-cutting: router requirement, route_type, failure hook.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_list__no_router_raises_500(list_harness):
    with patch.object(proxy_server, "llm_router", None):
        with pytest.raises(ProxyException) as exc:
            await call_list(list_harness)

    assert exc.value.code == "500"
    list_harness.litellm_alist.assert_not_called()


@pytest.mark.asyncio
async def test_list__uses_alist_batches_route_type(list_harness):
    await call_list(list_harness)

    assert list_harness.pre_call.call_args.kwargs["route_type"] == "alist_batches"


@pytest.mark.asyncio
async def test_list__exception_calls_failure_hook(list_harness):
    list_harness.litellm_alist.side_effect = ValueError("provider boom")

    with pytest.raises(ProxyException):
        await call_list(list_harness)

    list_harness.logging.post_call_failure_hook.assert_called_once()
    assert list_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom"


@pytest.mark.asyncio
async def test_list__failure_hook_and_response_share_the_request_litellm_call_id(list_harness):
    call_id = "lit7836-list-batches-call-id"
    list_harness.pre_call.side_effect = lambda **kw: (
        {**list_harness.body["body"], "litellm_call_id": call_id},
        MagicMock(),
    )
    list_harness.litellm_alist.side_effect = ValueError("provider boom")

    with pytest.raises(ProxyException) as raised:
        await call_list(list_harness, after="batch-0", limit=5)

    failure_request_data = list_harness.logging.post_call_failure_hook.call_args.kwargs["request_data"]
    assert failure_request_data["litellm_call_id"] == call_id
    assert (failure_request_data["after"], failure_request_data["limit"]) == ("batch-0", 5)
    assert raised.value.headers["x-litellm-call-id"] == call_id


# =========================================================================== #
#                                                                             #
#   POST /v1/batches/{batch_id}/cancel  -  cancel_batch routing-contract tests #
#                                                                             #
#   Three branches, first match wins:                                         #
#     1. model-encoded batch id -> litellm.acancel_batch via model creds       #
#     2. unified batch id       -> llm_router.acancel_batch (model+batch_id    #
#                                  parsed out of the unified id)               #
#     3. fallback               -> litellm.acancel_batch via env-var provider  #
#   Every branch then writes state back via update_batch_in_database(          #
#   operation="cancel"). There is NO ManagedObjectTable read short-circuit     #
#   here (unlike retrieve).                                                    #
#                                                                             #
#   These tests pin CURRENT behavior so a refactor can't silently change it.  #
#   Two current-behavior quirks are locked deliberately and noted inline:     #
#     - SCENARIO 1 forwards the DEPLOYMENT model from creds, not the decoded   #
#       model (retrieve overrides it; cancel does not).                        #
#     - SCENARIO 3 rebuilds a CancelBatchRequest and forwards only            #
#       {custom_llm_provider, batch_id}, dropping enrichment keys.            #
# =========================================================================== #


@dataclass
class CancelHarness:
    data: Dict[str, Any]
    pre_call: AsyncMock
    add_data: AsyncMock
    get_headers: MagicMock
    provider_from_headers: MagicMock
    provider_from_query: MagicMock
    litellm_acancel: AsyncMock
    router: MagicMock
    logging: MagicMock
    creds_resolver: MagicMock
    update_batch_in_db: AsyncMock

    @property
    def router_acancel(self) -> AsyncMock:
        return self.router.acancel_batch

    def acancel_kwargs(self) -> Dict[str, Any]:
        assert self.litellm_acancel.call_count == 1
        return dict(self.litellm_acancel.call_args.kwargs)

    def router_kwargs(self) -> Dict[str, Any]:
        assert self.router_acancel.call_count == 1
        return dict(self.router_acancel.call_args.kwargs)


@pytest.fixture
def cancel_harness():
    data_holder: Dict[str, Any] = {"data": {}}
    logging = MagicMock(spec=ProxyLogging)
    logging.post_call_success_hook = AsyncMock(side_effect=lambda **kw: kw["response"])
    logging.post_call_failure_hook = AsyncMock()
    logging.update_request_status = AsyncMock()
    logging.get_proxy_hook = MagicMock(return_value=None)

    router = MagicMock(spec=Router)
    router.model_group_alias = {}
    router.get_model_access_groups = MagicMock(return_value={})
    router.resolve_model_name_from_model_id = MagicMock(side_effect=lambda model_id: model_id)
    router.model_list = []
    router.acancel_batch = AsyncMock(return_value=make_batch())
    router.get_deployment_credentials_with_provider = MagicMock(side_effect=_creds_lookup)

    pre_call = AsyncMock(side_effect=lambda **kw: (data_holder["data"], MagicMock()))
    # add_litellm_data_to_request is a passthrough that returns the data it got.
    add_data = AsyncMock(side_effect=lambda **kw: kw["data"])
    get_headers = MagicMock(return_value={})
    provider_from_headers = MagicMock(return_value=None)
    provider_from_query = MagicMock(return_value=None)
    litellm_acancel = AsyncMock(return_value=make_batch())
    update_batch_in_db = AsyncMock(return_value=None)

    with ExitStack() as stack:
        stack.enter_context(
            patch.object(
                ProxyBaseLLMRequestProcessing,
                "common_processing_pre_call_logic",
                pre_call,
            )
        )
        stack.enter_context(patch.object(ProxyBaseLLMRequestProcessing, "get_custom_headers", get_headers))
        stack.enter_context(
            patch.object(
                endpoints,
                "get_custom_llm_provider_from_request_headers",
                provider_from_headers,
            )
        )
        stack.enter_context(
            patch.object(
                endpoints,
                "get_custom_llm_provider_from_request_query",
                provider_from_query,
            )
        )
        stack.enter_context(patch.object(endpoints, "update_batch_in_database", update_batch_in_db))
        stack.enter_context(patch.object(litellm, "acancel_batch", litellm_acancel))
        stack.enter_context(patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", False))
        stack.enter_context(patch.object(proxy_server, "llm_router", router))
        stack.enter_context(patch.object(proxy_server, "proxy_logging_obj", logging))
        stack.enter_context(patch.object(proxy_server, "general_settings", {}))
        stack.enter_context(patch.object(proxy_server, "proxy_config", MagicMock()))
        stack.enter_context(patch.object(proxy_server, "version", "test-version"))
        stack.enter_context(patch.object(proxy_server, "prisma_client", MagicMock()))
        stack.enter_context(patch.object(proxy_server, "add_litellm_data_to_request", add_data))

        yield CancelHarness(
            data=data_holder,
            pre_call=pre_call,
            add_data=add_data,
            get_headers=get_headers,
            provider_from_headers=provider_from_headers,
            provider_from_query=provider_from_query,
            litellm_acancel=litellm_acancel,
            router=router,
            logging=logging,
            creds_resolver=router.get_deployment_credentials_with_provider,
            update_batch_in_db=update_batch_in_db,
        )


async def call_cancel(
    harness: CancelHarness,
    batch_id: str,
    *,
    provider: Optional[str] = None,
    user: Optional[UserAPIKeyAuth] = None,
    headers: Optional[Dict[str, str]] = None,
    query: Optional[Dict[str, str]] = None,
    data_extra: Optional[Dict[str, Any]] = None,
):
    harness.data["data"] = {"batch_id": batch_id, **(data_extra or {})}
    return await endpoints.cancel_batch(
        request=FakeRequest(headers=headers, query=query),
        batch_id=batch_id,
        fastapi_response=Response(),
        provider=provider,
        user_api_key_dict=user or UserAPIKeyAuth(api_key="sk-test"),
    )


# --------------------------------------------------------------------------- #
# SCENARIO 1 - model-encoded batch id -> litellm.acancel_batch via model creds.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_cancel__model_encoded_id(cancel_harness):
    resp = await call_cancel(cancel_harness, AZURE_BATCH_ID)

    # DISPATCH - model-credential path via litellm; router untouched.
    assert cancel_harness.litellm_acancel.call_count == 1
    cancel_harness.router_acancel.assert_not_called()

    # CREDENTIALS - resolved for the model decoded from the batch id.
    cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")

    # SEAM PAYLOAD - exact dict. NOTE current behavior: `model` is the
    # DEPLOYMENT name from creds, NOT the decoded model (cancel, unlike
    # retrieve, does not override it). Locking this guards the difference.
    assert cancel_harness.acancel_kwargs() == {
        "custom_llm_provider": "azure",
        "batch_id": "batch_orig123",  # decoded/stripped original id
        "api_key": "sk-azure",
        "api_base": "https://azure.test",
        "model": "azure/gpt-4o-deployment",
    }

    # OUTPUT SHAPE - response id re-encoded with the DECODED model.
    assert resp.id == encode_file_id_with_model("batch-provider-id", "azure/gpt-4o", id_type="batch")

    # write-back tagged as a cancel.
    assert cancel_harness.update_batch_in_db.call_count == 1
    assert cancel_harness.update_batch_in_db.call_args.kwargs["operation"] == "cancel"


@pytest.mark.asyncio
async def test_cancel__model_encoded_id_forwards_deployment_model(cancel_harness):
    """Pin the current contract: cancel forwards the creds' deployment model.
    If someone adds a decoded-model override (as retrieve has), this flips and
    must be reviewed."""
    await call_cancel(cancel_harness, AZURE_BATCH_ID)

    assert cancel_harness.acancel_kwargs()["model"] == "azure/gpt-4o-deployment"


@pytest.mark.asyncio
async def test_cancel__model_encoded_beats_unified(cancel_harness):
    with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
        await call_cancel(cancel_harness, AZURE_BATCH_ID)

    assert cancel_harness.litellm_acancel.call_count == 1
    cancel_harness.router_acancel.assert_not_called()
    cancel_harness.creds_resolver.assert_called_once_with(model_id="azure/gpt-4o")


# --------------------------------------------------------------------------- #
# SCENARIO 2 - unified batch id -> llm_router.acancel_batch. model and batch_id
# are parsed out of the unified id; hidden params stamped.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_cancel__unified_batch_id_routes_to_router(cancel_harness):
    with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
        resp = await call_cancel(cancel_harness, "batch-unified-blob")

    # DISPATCH - router fired, litellm did not, no creds lookup.
    assert cancel_harness.router_acancel.call_count == 1
    cancel_harness.litellm_acancel.assert_not_called()
    cancel_harness.creds_resolver.assert_not_called()

    # model + batch_id are extracted from the unified id and forwarded.
    assert cancel_harness.router_kwargs() == {
        "batch_id": "batch-raw-xyz",
        "model": "gpt-4o-mini",
    }

    # hidden params: unified id passed through, model_id stamped from data.
    assert resp._hidden_params["unified_batch_id"] == UNIFIED_BATCH_ID
    assert resp._hidden_params["model_id"] == "gpt-4o-mini"

    assert cancel_harness.update_batch_in_db.call_args.kwargs["operation"] == "cancel"


@pytest.mark.asyncio
async def test_cancel__db_write_receives_caller_auth(cancel_harness):
    """update_batch_in_database can only mint managed IDs for a cancelled batch's
    output files when it has an auth context, so cancel must forward the caller's."""
    caller = UserAPIKeyAuth(api_key="sk-test", user_id="user-cancel-1")
    with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID):
        await call_cancel(cancel_harness, "batch-unified-blob", user=caller)

    assert cancel_harness.update_batch_in_db.call_args.kwargs["user_api_key_dict"] is caller


@pytest.mark.asyncio
async def test_cancel__unified_missing_model_id_400(cancel_harness):
    # unified id with no model_id segment -> get_model_id returns None -> 400.
    with patch.object(
        endpoints,
        "is_base64_encoded_unified_file_id",
        return_value="litellm_proxy;llm_batch_id:batch-xyz",
    ):
        with pytest.raises(ProxyException) as exc:
            await call_cancel(cancel_harness, "batch-unified-blob")

    assert exc.value.code == "400"
    cancel_harness.router_acancel.assert_not_called()
    cancel_harness.litellm_acancel.assert_not_called()


@pytest.mark.asyncio
async def test_cancel__unified_no_router_500(cancel_harness):
    with (
        patch.object(proxy_server, "llm_router", None),
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID),
    ):
        with pytest.raises(ProxyException) as exc:
            await call_cancel(cancel_harness, "batch-unified-blob")

    assert exc.value.code == "500"


@pytest.mark.asyncio
async def test_cancel__executed_batch_routes_to_runner(cancel_harness, executed_runner):
    runner, factory = executed_runner
    caller = UserAPIKeyAuth(api_key="sk-test", user_id="user-cancel-2")
    resp = await call_cancel(cancel_harness, EXECUTED_BATCH_B64, user=caller)

    runner.cancel.assert_awaited_once_with(EXECUTED_BATCH_B64, caller)
    factory.assert_called_once_with(cancel_harness.router, cancel_harness.logging)
    cancel_harness.router_acancel.assert_not_called()
    cancel_harness.litellm_acancel.assert_not_called()
    cancel_harness.creds_resolver.assert_not_called()
    assert resp is runner.cancel.return_value
    assert cancel_harness.update_batch_in_db.call_args.kwargs["operation"] == "cancel"


@pytest.mark.asyncio
async def test_cancel__executed_batch_no_router_500(cancel_harness, executed_runner):
    runner, factory = executed_runner
    with patch.object(  # test-quality-ok: proxy_server module global is the endpoint's only injection point
        proxy_server, "llm_router", None
    ):
        with pytest.raises(ProxyException) as exc:
            await call_cancel(cancel_harness, EXECUTED_BATCH_B64)

    assert exc.value.code == "500"
    factory.assert_not_called()
    runner.cancel.assert_not_called()


# --------------------------------------------------------------------------- #
# SCENARIO 3 - fallback to custom_llm_provider. Rebuilds a CancelBatchRequest
# and forwards only {custom_llm_provider, batch_id}.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_cancel__fallback_default_openai(cancel_harness, openai_env_creds):
    await call_cancel(cancel_harness, "batch-raw-xyz")

    assert cancel_harness.litellm_acancel.call_count == 1
    cancel_harness.router_acancel.assert_not_called()
    cancel_harness.creds_resolver.assert_not_called()  # inverse-bug guard
    # current behavior: enrichment keys dropped; only these two forwarded.
    assert cancel_harness.acancel_kwargs() == {
        "custom_llm_provider": "openai",
        "batch_id": "batch-raw-xyz",
    }
    assert cancel_harness.update_batch_in_db.call_count == 1


@pytest.mark.asyncio
async def test_cancel__fallback_no_creds_404(cancel_harness, no_openai_creds):
    with pytest.raises(ProxyException) as exc:
        await call_cancel(cancel_harness, "batch-raw-xyz")

    assert exc.value.code == "404"
    assert exc.value.type == "invalid_request_error"
    assert exc.value.param is None
    assert exc.value.message == "No batch found with id 'batch-raw-xyz'."
    cancel_harness.litellm_acancel.assert_not_called()
    cancel_harness.router_acancel.assert_not_called()


@pytest.mark.asyncio
async def test_cancel__fallback_explicit_provider_bypasses_not_found_gate(cancel_harness, no_openai_creds):
    await call_cancel(cancel_harness, "batch-raw-xyz", provider="anthropic")

    assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "anthropic"


@pytest.mark.asyncio
async def test_cancel__fallback_env_key_alone_forwards(cancel_harness, monkeypatch):
    monkeypatch.setattr(litellm, "api_key", None)
    monkeypatch.setattr(litellm, "openai_key", None)
    monkeypatch.setenv("OPENAI_API_KEY", "sk-env-openai")

    await call_cancel(cancel_harness, "batch-raw-xyz")

    assert cancel_harness.acancel_kwargs() == {
        "custom_llm_provider": "openai",
        "batch_id": "batch-raw-xyz",
    }


@pytest.mark.asyncio
async def test_cancel__fallback_provider_path_param(cancel_harness):
    await call_cancel(cancel_harness, "batch-raw-xyz", provider="anthropic")

    cancel_harness.creds_resolver.assert_not_called()
    assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "anthropic"


@pytest.mark.asyncio
async def test_cancel__fallback_provider_from_data_body(cancel_harness):
    await call_cancel(cancel_harness, "batch-raw-xyz", data_extra={"custom_llm_provider": "bedrock"})

    assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "bedrock"


@pytest.mark.asyncio
async def test_cancel__fallback_provider_from_header(cancel_harness):
    cancel_harness.provider_from_headers.return_value = "vertex_ai"

    await call_cancel(cancel_harness, "batch-raw-xyz")

    assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "vertex_ai"


@pytest.mark.asyncio
async def test_cancel__fallback_provider_from_query(cancel_harness):
    cancel_harness.provider_from_query.return_value = "azure"

    await call_cancel(cancel_harness, "batch-raw-xyz")

    assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "azure"


@pytest.mark.asyncio
async def test_cancel__fallback_provider_precedence_path_over_body(cancel_harness):
    """Intended contract: provider path param beats a body custom_llm_provider.
    Regression guard: the body value is popped from data before the fallback
    chain, so it never collides with the explicit kwarg."""
    await call_cancel(
        cancel_harness,
        "batch-raw-xyz",
        provider="anthropic",
        data_extra={"custom_llm_provider": "bedrock"},
    )

    assert cancel_harness.acancel_kwargs()["custom_llm_provider"] == "anthropic"


# --------------------------------------------------------------------------- #
# Cross-cutting: enrichment route_type and failure-hook on provider error.
# --------------------------------------------------------------------------- #


@pytest.mark.asyncio
async def test_cancel__uses_acancel_batch_route_type(cancel_harness, openai_env_creds):
    await call_cancel(cancel_harness, "batch-raw-xyz")

    assert cancel_harness.pre_call.call_args.kwargs["route_type"] == "acancel_batch"


@pytest.mark.asyncio
async def test_cancel__exception_calls_failure_hook(cancel_harness, openai_env_creds):
    cancel_harness.litellm_acancel.side_effect = ValueError("provider boom")

    with pytest.raises(ProxyException):
        await call_cancel(cancel_harness, "batch-raw-xyz")

    cancel_harness.logging.post_call_failure_hook.assert_called_once()
    assert (
        cancel_harness.logging.post_call_failure_hook.call_args.kwargs["original_exception"].args[0] == "provider boom"
    )


# =========================================================================== #
# Router-required 500 guards (one per endpoint branch that calls the router).
# These pin the defensive checks that fire when llm_router is unset.
# =========================================================================== #


@pytest.mark.asyncio
async def test_create__loadbalancing_no_router_500(harness):
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
            "model": "lb-model",
        },
    )
    harness.is_known_model.return_value = True
    with (
        patch.object(litellm, "enable_loadbalancing_on_batch_endpoints", True),
        patch.object(proxy_server, "llm_router", None),
    ):
        with pytest.raises(ProxyException) as exc:
            await call_create(harness)

    assert exc.value.code == "500"
    harness.router_acreate.assert_not_called()
    harness.litellm_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_create__unified_no_router_500(harness):
    set_body(
        harness,
        {
            "input_file_id": "litellm_proxy_unified_id",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    with (
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value="unified-xyz"),
        patch.object(endpoints, "get_models_from_unified_file_id", return_value=["gpt-4o-mini"]),
        patch.object(proxy_server, "llm_router", None),
    ):
        with pytest.raises(ProxyException) as exc:
            await call_create(harness)

    assert exc.value.code == "500"


@pytest.mark.asyncio
async def test_retrieve__unified_no_router_500(retrieve_harness):
    with (
        patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID),
        patch.object(proxy_server, "llm_router", None),
    ):
        with pytest.raises(ProxyException) as exc:
            await call_retrieve(retrieve_harness, "batch-unified-blob")

    assert exc.value.code == "500"
    retrieve_harness.router_aretrieve.assert_not_called()
    retrieve_harness.litellm_aretrieve.assert_not_called()


# =========================================================================== #
# SCENARIO 3 + configured deployments: a provider-only call (custom-llm-provider
# header, no model anywhere) must resolve the gateway/team deployment's named
# credential for that provider and attach it to the provider call kwargs,
# instead of silently falling through to the host environment's default
# credentials (regression: vertex batch jobs landing in the hosting env's GCP
# project because litellm_credential_name never reached the call).
# =========================================================================== #

VERTEX_NAMED_CREDENTIAL = CredentialItem(
    credential_name="vertex-named-cred",
    credential_info={},
    credential_values={
        "vertex_project": "customer-project",
        "vertex_location": "us-central1",
        "vertex_credentials": "/creds/customer-sa.json",
    },
)


def vertex_named_credential_router() -> Router:
    return Router(
        model_list=[
            {
                "model_name": "gemini-2.5-pro",
                "litellm_params": {
                    "model": "vertex_ai/gemini-2.5-pro",
                    "litellm_credential_name": "vertex-named-cred",
                },
            }
        ]
    )


@pytest.mark.asyncio
async def test_create__provider_only_resolves_named_vertex_credentials(harness):
    """Provider-only create must attach the configured named credential, and must
    NOT turn the call into a model-routed one (no model kwarg injected)."""
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    harness.provider_from_headers.return_value = "vertex_ai"

    with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
        with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
            await call_create(harness)

    assert harness.acreate_kwargs() == {
        "custom_llm_provider": "vertex_ai",
        "input_file_id": "file-plain",
        "endpoint": "/v1/chat/completions",
        "completion_window": "24h",
        "metadata": None,
        "vertex_project": "customer-project",
        "vertex_location": "us-central1",
        "vertex_credentials": "/creds/customer-sa.json",
    }


@pytest.mark.asyncio
async def test_create__provider_only_ignores_other_provider_deployments(harness):
    """A provider-only vertex call must not pick up credentials from deployments
    of a different provider; with no vertex deployment the payload is exactly the
    pre-fix env-var fallback."""
    set_body(
        harness,
        {
            "input_file_id": "file-plain",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    harness.provider_from_headers.return_value = "vertex_ai"
    openai_only_router = Router(
        model_list=[
            {
                "model_name": "gpt-4o",
                "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-global-openai"},
            }
        ]
    )

    with patch.object(proxy_server, "llm_router", openai_only_router):
        await call_create(harness)

    assert harness.acreate_kwargs() == {
        "custom_llm_provider": "vertex_ai",
        "input_file_id": "file-plain",
        "endpoint": "/v1/chat/completions",
        "completion_window": "24h",
        "metadata": None,
    }


@pytest.mark.asyncio
async def test_retrieve__provider_only_resolves_named_vertex_credentials(retrieve_harness):
    retrieve_harness.provider_from_headers.return_value = "vertex_ai"

    with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
        with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
            await call_retrieve(retrieve_harness, "batch-raw-xyz")

    assert retrieve_harness.aretrieve_kwargs() == {
        "custom_llm_provider": "vertex_ai",
        "batch_id": "batch-raw-xyz",
        "vertex_project": "customer-project",
        "vertex_location": "us-central1",
        "vertex_credentials": "/creds/customer-sa.json",
    }


@pytest.mark.asyncio
async def test_list__provider_only_resolves_named_vertex_credentials(list_harness):
    list_harness.provider_from_headers.return_value = "vertex_ai"

    with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
        with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
            await call_list(list_harness)

    assert list_harness.alist_kwargs() == {
        "custom_llm_provider": "vertex_ai",
        "after": None,
        "limit": None,
        "vertex_project": "customer-project",
        "vertex_location": "us-central1",
        "vertex_credentials": "/creds/customer-sa.json",
    }


@pytest.mark.asyncio
async def test_cancel__provider_only_resolves_named_vertex_credentials(cancel_harness):
    cancel_harness.provider_from_headers.return_value = "vertex_ai"

    with patch.object(litellm, "credential_list", [VERTEX_NAMED_CREDENTIAL]):
        with patch.object(proxy_server, "llm_router", vertex_named_credential_router()):
            await call_cancel(cancel_harness, "batch-raw-xyz")

    assert cancel_harness.acancel_kwargs() == {
        "custom_llm_provider": "vertex_ai",
        "batch_id": "batch-raw-xyz",
        "vertex_project": "customer-project",
        "vertex_location": "us-central1",
        "vertex_credentials": "/creds/customer-sa.json",
    }


# =========================================================================== #
# require_managed_files - raw provider ids must not reach the provider.        #
#                                                                              #
# Ownership rows only exist for LiteLLM managed ids. A raw provider id sent to #
# these routes is forwarded under the shared provider credentials with no      #
# tenant check, so any caller who learns another tenant's id can read its      #
# batch, reuse its file as batch input, or cancel its job. These lock the      #
# guard on every batches route that accepts a caller-supplied id.              #
# =========================================================================== #


def _unified_batch_id(model_id: str = "azure/gpt-4o", batch_id: str = "batch-provider-id") -> str:
    import base64

    from litellm.types.utils import SpecialEnums

    unified = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format(model_id, batch_id)
    return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")


def _unified_file_id() -> str:
    import base64

    from litellm.types.utils import SpecialEnums

    unified = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format(
        "application/json", "managed-id", "gpt-4o-mini", "file-provider-id", "gpt-4o-mini-id"
    )
    return base64.urlsafe_b64encode(unified.encode()).decode().rstrip("=")


@dataclass(frozen=True)
class ManagedResourceAccessCheckerStub:
    file_access: bool = True
    object_access: bool = True

    async def can_user_call_unified_file_id(
        self,
        unified_file_id: str,
        user_api_key_dict: UserAPIKeyAuth,
    ) -> bool:
        return self.file_access

    async def can_user_call_unified_object_id(
        self,
        unified_object_id: str,
        user_api_key_dict: UserAPIKeyAuth,
    ) -> bool:
        return self.object_access


@pytest.mark.asyncio
async def test_create__raw_input_file_id_rejected_when_managed_files_required(harness):
    set_body(
        harness,
        {
            "input_file_id": "file-victim-abc123",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    with patch.object(litellm, "require_managed_files", True):
        with pytest.raises(ProxyException) as exc:
            await call_create(harness)

    assert exc.value.code == "400"
    harness.litellm_acreate.assert_not_called()
    harness.router_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_create__model_encoded_input_file_id_rejected_when_managed_files_required(harness):
    """A model-encoded id is client-forgeable and has no ownership row, so it is
    not a managed file id and must be rejected like any other raw id."""
    set_body(
        harness,
        {
            "input_file_id": AZURE_FILE_ID,
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    with patch.object(litellm, "require_managed_files", True):
        with pytest.raises(ProxyException) as exc:
            await call_create(harness)

    assert exc.value.code == "400"
    harness.litellm_acreate.assert_not_called()
    harness.router_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_create__raw_input_file_id_allowed_when_managed_files_not_required(harness, openai_env_creds):
    set_body(
        harness,
        {
            "input_file_id": "file-victim-abc123",
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    with patch.object(litellm, "require_managed_files", False):
        await call_create(harness)

    assert harness.acreate_kwargs()["input_file_id"] == "file-victim-abc123"


@pytest.mark.asyncio
async def test_create__other_teams_unified_input_file_id_rejected(harness):
    set_body(
        harness,
        {
            "input_file_id": _unified_file_id(),
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )
    harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub(file_access=False)

    with patch.object(litellm, "require_managed_files", True):
        with pytest.raises(ProxyException) as exc:
            await call_create(harness)

    assert exc.value.code == "403"
    harness.litellm_acreate.assert_not_called()
    harness.router_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_retrieve__raw_batch_id_rejected_when_managed_files_required(retrieve_harness):
    with patch.object(litellm, "require_managed_files", True):
        with pytest.raises(ProxyException) as exc:
            await call_retrieve(retrieve_harness, "batch-victim-abc123")

    assert exc.value.code == "400"
    retrieve_harness.litellm_aretrieve.assert_not_called()
    retrieve_harness.router_aretrieve.assert_not_called()


@pytest.mark.asyncio
async def test_retrieve__unified_batch_id_allowed_when_managed_files_required(retrieve_harness):
    retrieve_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub()

    with patch.object(litellm, "require_managed_files", True):
        await call_retrieve(retrieve_harness, _unified_batch_id())

    assert retrieve_harness.router_aretrieve.call_count == 1


@pytest.mark.asyncio
async def test_cancel__raw_batch_id_rejected_when_managed_files_required(cancel_harness):
    with patch.object(litellm, "require_managed_files", True):
        with pytest.raises(ProxyException) as exc:
            await call_cancel(cancel_harness, "batch-victim-abc123")

    assert exc.value.code == "400"
    cancel_harness.litellm_acancel.assert_not_called()
    cancel_harness.router_acancel.assert_not_called()


@pytest.mark.asyncio
async def test_cancel__unified_batch_id_allowed_when_managed_files_required(cancel_harness):
    cancel_harness.logging.get_proxy_hook.return_value = ManagedResourceAccessCheckerStub()

    with patch.object(litellm, "require_managed_files", True):
        await call_cancel(cancel_harness, _unified_batch_id())

    assert cancel_harness.router_acancel.call_count == 1


@pytest.mark.asyncio
async def test_retrieve__managed_batch_defers_cost_to_the_poller_when_it_is_running(retrieve_harness):
    with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)):
        await call_retrieve(retrieve_harness, _unified_batch_id())

    assert retrieve_harness.router.aretrieve_batch.await_count == 1
    metadata = retrieve_harness.router.aretrieve_batch.await_args.kwargs.get("litellm_metadata") or {}
    assert metadata.get("batch_ignore_default_logging") is True


@pytest.mark.asyncio
async def test_retrieve__managed_batch_still_accounts_inline_without_a_poller(retrieve_harness):
    with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=False)):
        await call_retrieve(retrieve_harness, _unified_batch_id())

    assert retrieve_harness.router.aretrieve_batch.await_count == 1
    metadata = retrieve_harness.router.aretrieve_batch.await_args.kwargs.get("litellm_metadata") or {}
    assert metadata.get("batch_ignore_default_logging") is None


@pytest.mark.asyncio
async def test_retrieve__raw_batch_id_is_untouched_by_the_poller_handoff(retrieve_harness, openai_env_creds):
    with patch.object(endpoints, "batch_cost_poller_is_active", MagicMock(return_value=True)):
        await call_retrieve(retrieve_harness, "batch-raw-xyz")

    metadata = retrieve_harness.litellm_aretrieve.await_args.kwargs.get("litellm_metadata") or {}
    assert metadata.get("batch_ignore_default_logging") is None


def _key_restricted_to(*models: str) -> UserAPIKeyAuth:
    return UserAPIKeyAuth(api_key="sk-restricted", team_id="team-a", team_models=list(models), models=list(models))


@pytest.mark.asyncio
async def test_create__header_model_rejects_key_without_model_grant(harness):
    """A key not granted the model named in x-litellm-model must not receive that deployment's credentials."""
    set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"})

    with pytest.raises(ProxyException) as exc_info:
        await call_create(harness, user=_key_restricted_to("azure/gpt-4o"), headers={"x-litellm-model": "vertex-model"})

    assert exc_info.value.code == "403"
    harness.creds_resolver.assert_not_called()
    harness.litellm_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_create__header_model_allows_key_with_model_grant(harness):
    set_body(harness, {"input_file_id": "file-plain", "endpoint": "/v1/chat/completions", "completion_window": "24h"})

    await call_create(harness, user=_key_restricted_to("vertex-model"), headers={"x-litellm-model": "vertex-model"})

    harness.creds_resolver.assert_called_once_with(model_id="vertex-model")
    assert harness.acreate_kwargs()["custom_llm_provider"] == "vertex_ai"


@pytest.mark.asyncio
async def test_retrieve__model_encoded_id_rejects_key_without_model_grant(retrieve_harness):
    """The model embedded in a batch id is caller-controlled, so it is checked against the key's grants too."""
    with pytest.raises(ProxyException) as exc_info:
        await call_retrieve(retrieve_harness, AZURE_BATCH_ID, user=_key_restricted_to("vertex-model"))

    assert exc_info.value.code == "403"
    retrieve_harness.creds_resolver.assert_not_called()
    retrieve_harness.litellm_aretrieve.assert_not_called()


@pytest.mark.asyncio
async def test_cancel__model_encoded_id_rejects_key_without_model_grant(cancel_harness):
    with pytest.raises(ProxyException) as exc_info:
        await call_cancel(cancel_harness, AZURE_BATCH_ID, user=_key_restricted_to("vertex-model"))

    assert exc_info.value.code == "403"
    cancel_harness.creds_resolver.assert_not_called()
    cancel_harness.litellm_acancel.assert_not_called()


def _b64_unified_id(decoded: str) -> str:
    return base64.urlsafe_b64encode(decoded.encode()).decode().rstrip("=")


UNIFIED_FILE_ID_FOR_GPT4O_MINI = _b64_unified_id(
    "litellm_proxy:application/octet-stream;unified_id,c4843482-b176-4901-8292-7523fd0f2c6e;"
    "target_model_names,gpt-4o-mini;llm_output_file_id,file-provider;llm_output_file_model_id,dep-1"
)
UNIFIED_BATCH_ID_FOR_GPT4O_MINI = _b64_unified_id(UNIFIED_BATCH_ID)


@pytest.mark.asyncio
async def test_create__unified_file_id_rejects_key_without_model_grant(harness):
    """The model carried inside a unified file id is caller-controlled too, so it is checked against the key's grants."""
    set_body(
        harness,
        {
            "input_file_id": UNIFIED_FILE_ID_FOR_GPT4O_MINI,
            "endpoint": "/v1/chat/completions",
            "completion_window": "24h",
        },
    )

    with pytest.raises(ProxyException) as exc_info:
        await call_create(harness, user=_key_restricted_to("vertex-model"))

    assert exc_info.value.code == "403"
    harness.router_acreate.assert_not_called()
    harness.litellm_acreate.assert_not_called()


@pytest.mark.asyncio
async def test_retrieve__unified_batch_id_rejects_key_without_model_grant(retrieve_harness):
    with pytest.raises(ProxyException) as exc_info:
        await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to("vertex-model"))

    assert exc_info.value.code == "403"
    retrieve_harness.router_aretrieve.assert_not_called()
    retrieve_harness.creds_resolver.assert_not_called()


@pytest.mark.asyncio
async def test_retrieve__unified_batch_id_rejects_key_without_model_grant_before_db_terminal_shortcut(
    retrieve_harness,
):
    retrieve_harness.get_batch_from_db.return_value = (MagicMock(), make_batch(id="batch-from-db", status="completed"))

    with pytest.raises(ProxyException) as exc_info:
        await call_retrieve(retrieve_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to("vertex-model"))

    assert exc_info.value.code == "403"
    retrieve_harness.logging.post_call_success_hook.assert_not_called()
    retrieve_harness.ensure_managed_files.assert_not_called()
    retrieve_harness.router_aretrieve.assert_not_called()


@pytest.mark.asyncio
async def test_cancel__unified_batch_id_rejects_key_without_model_grant(cancel_harness):
    with pytest.raises(ProxyException) as exc_info:
        await call_cancel(cancel_harness, UNIFIED_BATCH_ID_FOR_GPT4O_MINI, user=_key_restricted_to("vertex-model"))

    assert exc_info.value.code == "403"
    cancel_harness.router_acancel.assert_not_called()


@pytest.mark.asyncio
async def test_cancel__executed_batch_rejects_key_without_model_grant(cancel_harness, executed_runner):
    runner, factory = executed_runner

    with pytest.raises(ProxyException) as exc_info:
        await call_cancel(cancel_harness, EXECUTED_BATCH_B64, user=_key_restricted_to("vertex-model"))

    assert exc_info.value.code == "403"
    factory.assert_not_called()
    runner.cancel.assert_not_called()
    cancel_harness.router_acancel.assert_not_called()