litellm/tests/integration/security/test_api_key_secret_reference.py

from __future__ import annotations

import base64
import json
import threading
import time
import uuid
from collections.abc import Callable, Iterator, Mapping, Sequence
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from types import MappingProxyType
from typing import Final

import httpx
import pytest
import yaml
from integration._support.client import (
    JSON_OBJECT,
    Gateway,
    Scenario,
    eventually,
    gateway_from_environment,
    list_value,
    object_value,
    string_value,
)
from integration._support.openai_wire import chat_reply, openai_error, responses_reply
from integration._support.process import (
    OwnedProxy,
    WorkerHealthcheck,
    graceful_stop_seconds,
    owned_proxy_process,
)
from integration._support.wire import Reply, Request, wire_server
from integration.security._sinks import Recorder
from pydantic import JsonValue

pytestmark: Final = pytest.mark.timeout(int(2 * graceful_stop_seconds() + 120))

_CLOSE: Final = MappingProxyType({"Connection": "close"})
_FRESH: Final = httpx.Limits(max_keepalive_connections=0)
_CHAT_MODEL: Final = "openai/gpt-4o-mini"
_EMBEDDING_MODEL: Final = "openai/text-embedding-3-small"
_REALTIME_MODEL: Final = "openai/gpt-4o-realtime-preview"
_LISTED_MODEL: Final = "wif-listed-model"
_VECTOR_STORE: Final = "vs_wif_reference"
_SEARCH_PATH: Final = f"/v1/vector_stores/{_VECTOR_STORE}/search"
_BATCH_DEPLOYMENT: Final = "wif-vllm-batch"
_BEDROCK_MODEL: Final = "bedrock/converse/us.anthropic.claude-sonnet-5-5"
_PROVIDER_INPUT_TOKENS: Final = 7919
_UNSUPPORTED_REFERENCE: Final = "oidc/github/https://example.invalid"
_GOOGLE_MODEL: Final = "hosted_vllm/wif-google"
_GOOGLE_AUDIENCES: Final = MappingProxyType(
    {
        "config_inline": "https://vllm-inline.example.test",
        "config_credential": "https://vllm-credential.example.test",
        "db_credential": "https://vllm-db-credential.example.test",
    }
)
_UNAVAILABLE_AUDIENCE: Final = "https://vllm-unavailable.example.test"
_HUNG_AUDIENCE: Final = "https://vllm-hung.example.test"
_IDENTITY_PATH: Final = "/computeMetadata/v1/instance/service-accounts/default/identity"
_RUN_SECRETS_TOKEN: Final = "/run/secrets/wif-absent-token"
_SOURCES: Final = ("config_credential", "config_inline", "db_credential", "db_inline")
_UNREADABLE: Final = MappingProxyType(
    {"missing_file": "wif-missing", "outside_allowlist": "wif-outside", "unsupported_kind": "wif-unsupported"}
)
_ATTACKS: Final = (
    "client_api_base_inline",
    "client_api_base_credential",
    "body_reference",
    "body_credential_name",
    "user_config",
)
_CREDENTIAL_ROW: Final = (
    'SELECT credential_values::text AS stored, updated_at::text AS updated_at FROM "LiteLLM_CredentialsTable" '
    "WHERE credential_name = %s"
)
_MODEL_ROW: Final = (
    'SELECT litellm_params::text AS stored, updated_at::text AS updated_at FROM "LiteLLM_ProxyModelTable" '
    "WHERE model_id = %s"
)
_ENDPOINTS: Final = MappingProxyType(
    {"embeddings": ("/v1/embeddings", "wif-embedding"), "responses": ("/v1/responses", "wif-config-inline")}
)
_TEAM_WRITES: Final = MappingProxyType(
    {
        ("team_admin", "plain_key"): 200,
        ("team_admin", "names_credential"): 403,
        ("team_admin", "inline_reference"): 400,
        ("proxy_admin", "inline_reference"): 400,
    }
)
_MODEL_ROWS: Final = (
    "SELECT model_id FROM \"LiteLLM_ProxyModelTable\" WHERE model_info::jsonb ->> 'team_public_model_name' = %s"
)


@dataclass(frozen=True, slots=True)
class _Secrets:
    allowed: Path
    token: Path
    missing: Path
    outside: Path

    def current(self) -> str:
        return self.token.read_text()

    def rotate(self) -> str:
        value: Final = f"wif-token-{uuid.uuid4().hex}"
        self.token.write_text(value)
        return value

    def overrides(self) -> Mapping[str, str]:
        return MappingProxyType(
            {"LITELLM_OIDC_ALLOWED_CREDENTIAL_DIRS": str(self.allowed), "PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": "2"}
        )


def _reference(path: Path) -> str:
    return f"oidc/file/{path}"


def _google_reference(audience: str) -> str:
    return f"oidc/google/{audience}"


def _jwt_segment(part: Mapping[str, JsonValue]) -> str:
    return base64.urlsafe_b64encode(json.dumps(part).encode()).decode().rstrip("=")


@dataclass(frozen=True, slots=True)
class _Google:
    expires: int

    def token(self, audience: str) -> str:
        claims: Final = _jwt_segment({"iss": "https://accounts.google.com", "aud": audience, "exp": self.expires})
        return f"{_jwt_segment({'alg': 'RS256', 'typ': 'JWT'})}.{claims}.signature"

    def bearers(self) -> frozenset[str]:
        return frozenset(f"Bearer {self.token(audience)}" for audience in _GOOGLE_AUDIENCES.values())


def _metadata_server(google: _Google) -> Callable[[Request], Reply]:
    def respond(request: Request) -> Reply:
        target: Final = httpx.URL(request.target)
        audience: Final = target.params.get("audience")
        if audience is None or (request.method, target.path, request.headers.get("metadata-flavor")) != (
            "GET",
            _IDENTITY_PATH,
            "Google",
        ):
            return Reply(status=404, body=b"not found", content_type="text/plain")
        if audience == _UNAVAILABLE_AUDIENCE:
            return Reply(status=503, body=b"metadata server unavailable", content_type="text/plain")
        if audience == _HUNG_AUDIENCE:
            return Reply(chunks=(b"never sent",), gate_before_first=threading.Event(), content_type="text/plain")
        return Reply(body=google.token(audience).encode(), content_type="text/plain")

    return respond


def _secrets(directory: Path) -> _Secrets:
    allowed: Final = directory / "secrets"
    outside_dir: Final = directory / "outside"
    allowed.mkdir()
    outside_dir.mkdir()
    secrets: Final = _Secrets(
        allowed=allowed, token=allowed / "token", missing=allowed / "absent-token", outside=outside_dir / "token"
    )
    secrets.rotate()
    secrets.outside.write_text(f"outside-secret-{uuid.uuid4().hex}")
    return secrets


def _json_reply(body: Mapping[str, JsonValue]) -> Reply:
    return Reply(body=json.dumps(body).encode())


def _served(request: Request) -> Reply:
    path: Final = request.target.split("?", 1)[0]
    match request.method, path:
        case "POST", "/v1/chat/completions":
            return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", "gpt-4o-mini", "served", stream=False)
        case "POST", "/v1/responses":
            return responses_reply(f"resp_{uuid.uuid4().hex}", "gpt-4o-mini", "served", stream=False)
        case "POST", "/v1/responses/input_tokens":
            return _json_reply({"object": "response.input_tokens", "input_tokens": _PROVIDER_INPUT_TOKENS})
        case "POST", converse if converse.startswith("/model/") and converse.endswith("/converse"):
            return _json_reply(
                {
                    "output": {"message": {"role": "assistant", "content": [{"text": "served"}]}},
                    "stopReason": "end_turn",
                    "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 2},
                }
            )
        case "POST", "/v1/embeddings":
            return _json_reply(
                {
                    "object": "list",
                    "data": [{"object": "embedding", "index": 0, "embedding": [0.25, 0.5]}],
                    "model": "text-embedding-3-small",
                    "usage": {"prompt_tokens": 2, "total_tokens": 2},
                }
            )
        case "GET", "/v1/models":
            return _json_reply(
                {"object": "list", "data": [{"id": _LISTED_MODEL, "object": "model", "created": 0, "owned_by": "it"}]}
            )
        case "POST", searched if searched == _SEARCH_PATH:
            return _json_reply(
                {
                    "object": "vector_store.search_results.page",
                    "search_query": "served",
                    "data": [],
                    "has_more": False,
                    "next_page": None,
                }
            )
        case _:
            return openai_error(404)


def _provider(secrets: _Secrets, google: _Google) -> Callable[[Request], Reply]:
    def respond(request: Request) -> Reply:
        if request.headers.get("authorization") not in {f"Bearer {secrets.current()}", *google.bearers()}:
            return openai_error(401)
        return _served(request)

    return respond


def _intruder(request: Request) -> Reply:
    return chat_reply(f"chatcmpl-{uuid.uuid4().hex}", "gpt-4o-mini", "intercepted", stream=False)


def _deployment(
    name: str, params: Mapping[str, JsonValue], info: Mapping[str, JsonValue] = MappingProxyType({})
) -> dict[str, JsonValue]:
    return {"model_name": name, "litellm_params": dict(params), "model_info": dict(info)}


def _credential(name: str, api_key: str, api_base: str) -> dict[str, JsonValue]:
    return {
        "credential_name": name,
        "credential_values": {"api_key": api_key, "api_base": api_base},
        "credential_info": {"custom_llm_provider": "openai"},
    }


def _named(model: str, credential: str) -> dict[str, JsonValue]:
    return {"model": model, "litellm_credential_name": credential}


def _config(directory: Path, secrets: _Secrets, provider: str) -> Path:
    inline: Final = {"api_key": _reference(secrets.token), "api_base": provider}
    configuration: Final[dict[str, JsonValue]] = {
        "model_list": [
            _deployment("wif-config-credential", _named(_CHAT_MODEL, "wif-config")),
            _deployment("wif-config-inline", {"model": _CHAT_MODEL, **inline}),
            _deployment("wif-embedding", _named(_EMBEDDING_MODEL, "wif-config"), {"mode": "embedding"}),
            _deployment("wif-missing", _named(_CHAT_MODEL, "wif-missing")),
            _deployment("wif-outside", _named(_CHAT_MODEL, "wif-outside")),
            _deployment("wif-unsupported", _named(_CHAT_MODEL, "wif-unsupported")),
            _deployment("wif-run-secrets", _named(_CHAT_MODEL, "wif-run-secrets")),
            _deployment("wif-missing-fallback", _named(_CHAT_MODEL, "wif-missing")),
            _deployment("wif-realtime", _named(_REALTIME_MODEL, "wif-config"), {"mode": "realtime"}),
            _deployment("wif-vllm", {"model": "hosted_vllm/wif-local", **inline}, {"id": _BATCH_DEPLOYMENT}),
            _deployment("openai/*", _named("openai/*", "wif-config")),
            _deployment(
                "wif-bedrock-endpoint",
                {
                    "model": _BEDROCK_MODEL,
                    "aws_region_name": "us-east-1",
                    "api_key": _reference(secrets.token),
                    "aws_bedrock_runtime_endpoint": provider.removesuffix("/v1"),
                },
            ),
            _deployment(
                "wif-google-inline",
                {
                    "model": _GOOGLE_MODEL,
                    "api_key": _google_reference(_GOOGLE_AUDIENCES["config_inline"]),
                    "api_base": provider,
                },
            ),
            _deployment("wif-google-credential", _named(_GOOGLE_MODEL, "wif-google")),
            _deployment(
                "wif-google-unavailable",
                {"model": _GOOGLE_MODEL, "api_key": _google_reference(_UNAVAILABLE_AUDIENCE), "api_base": provider},
            ),
            _deployment(
                "wif-google-unavailable-count",
                {"model": _CHAT_MODEL, "api_key": _google_reference(_UNAVAILABLE_AUDIENCE), "api_base": provider},
            ),
            _deployment(
                "wif-google-hung",
                {"model": _GOOGLE_MODEL, "api_key": _google_reference(_HUNG_AUDIENCE), "api_base": provider},
            ),
            _deployment(
                "wif-mistral-passthrough",
                {
                    "model": "mistral/mistral-small-latest",
                    "api_key": _reference(secrets.missing),
                    "use_in_pass_through": True,
                },
            ),
        ],
        "credential_list": [
            _credential("wif-config", _reference(secrets.token), provider),
            _credential("wif-missing", _reference(secrets.missing), provider),
            _credential("wif-outside", _reference(secrets.outside), provider),
            _credential("wif-unsupported", _UNSUPPORTED_REFERENCE, provider),
            _credential("wif-run-secrets", f"oidc/file/{_RUN_SECRETS_TOKEN}", provider),
            _credential("wif-google", _google_reference(_GOOGLE_AUDIENCES["config_credential"]), provider),
        ],
        "vector_store_registry": [
            {
                "vector_store_name": "wif-vector-store",
                "litellm_params": {
                    "vector_store_id": _VECTOR_STORE,
                    "custom_llm_provider": "openai",
                    "litellm_credential_name": "wif-config",
                },
            }
        ],
        "router_settings": {
            "fallbacks": [{"wif-missing-fallback": ["wif-config-credential"]}],
            "disable_cooldowns": True,
        },
        "litellm_settings": {"check_provider_endpoint": True},
        "general_settings": {"master_key": "os.environ/LITELLM_MASTER_KEY", "allow_client_side_credentials": True},
    }
    path: Final = directory / "api_key_secret_reference_config.yaml"
    path.write_text(yaml.safe_dump(configuration, sort_keys=False))
    return path


@dataclass(frozen=True, slots=True)
class _Stored:
    credential: str
    models: Mapping[str, str]
    model_ids: Mapping[str, str]


def _delete_credential(gateway: Gateway, name: str) -> None:
    response: Final = gateway.request("DELETE", f"/credentials/{name}")
    assert response.status_code in (200, 404), response.text


def _created_model(gateway: Gateway, scenario: Scenario, body: Mapping[str, JsonValue]) -> tuple[str, str]:
    created: Final = gateway.post("/model/new", body)
    identity: Final = string_value(object_value(created["model_info"])["id"])
    scenario.cleanups.callback(scenario.delete_model, identity)
    return string_value(created["model_name"]), identity


def _stored(owned: OwnedProxy, scenario: Scenario, secrets: _Secrets, provider: str) -> _Stored:
    credential: Final = f"db-oidc-{uuid.uuid4().hex}"
    owned.gateway.post("/credentials", _credential(credential, _reference(secrets.token), provider))
    scenario.cleanups.callback(_delete_credential, owned.gateway, credential)
    by_credential: Final = _created_model(
        owned.gateway,
        scenario,
        _deployment(f"db-credential-{uuid.uuid4().hex}", _named(_CHAT_MODEL, credential)),
    )
    inline: Final = _created_model(
        owned.gateway,
        scenario,
        _deployment(
            f"db-inline-{uuid.uuid4().hex}",
            {"model": _CHAT_MODEL, "api_key": _reference(secrets.token), "api_base": provider},
        ),
    )
    for name, identity in (by_credential, inline):
        eventually(partial(owned.serves_only_on_every_worker, name, identity), bool, seconds=120)
    return _Stored(
        credential=credential,
        models=MappingProxyType({"db_credential": by_credential[0], "db_inline": inline[0]}),
        model_ids=MappingProxyType({"db_credential": by_credential[1], "db_inline": inline[1]}),
    )


@dataclass(frozen=True, slots=True)
class _Rig:
    owned: OwnedProxy
    provider: Recorder
    intruder: Recorder
    metadata: Recorder
    google: _Google
    provider_base: str
    intruder_base: str
    secrets: _Secrets
    stored: _Stored
    boot: tuple[Request, ...]
    boot_token: str

    def model(self, source: str) -> str:
        match source:
            case "config_credential":
                return "wif-config-credential"
            case "config_inline":
                return "wif-config-inline"
            case "db_credential" | "db_inline":
                return self.stored.models[source]
            case _:
                pytest.fail(f"unknown source {source}")


def _is_discovery(request: Request) -> bool:
    return (request.method, request.target) == ("GET", "/v1/models")


@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
    directory: Final = tmp_path_factory.mktemp("api_key_secret_reference").resolve()
    secrets: Final = _secrets(directory)
    google: Final = _Google(expires=int(time.time()) + 3600)
    with (
        gateway_from_environment() as gateway,
        wire_server(_provider(secrets, google)) as provider_wire,
        wire_server(_intruder) as intruder_wire,
        wire_server(_metadata_server(google)) as metadata_wire,
    ):
        provider: Final = Recorder(provider_wire)
        provider_base: Final = f"{provider_wire.url}/v1"
        with owned_proxy_process(
            gateway,
            directory,
            {
                **secrets.overrides(),
                "MISTRAL_API_BASE": provider_wire.url,
                "GCE_METADATA_HOST": httpx.URL(metadata_wire.url).netloc.decode(),
                "GCE_METADATA_TIMEOUT": "1",
            },
            config=_config(directory, secrets, provider_base),
            workers=2,
            healthcheck=WorkerHealthcheck.LENIENT,
        ) as owned:
            with owned.gateway.scenario() as scenario:
                stored: Final = _stored(owned, scenario, secrets, provider_base)
                boot: Final = eventually(provider.requests, lambda seen: any(_is_discovery(r) for r in seen))
                yield _Rig(
                    owned=owned,
                    provider=provider,
                    intruder=Recorder(intruder_wire),
                    metadata=Recorder(metadata_wire),
                    google=google,
                    provider_base=provider_base,
                    intruder_base=f"{intruder_wire.url}/v1",
                    secrets=secrets,
                    stored=stored,
                    boot=boot,
                    boot_token=secrets.current(),
                )


def _marker() -> str:
    return f"wif-marker-{uuid.uuid4().hex}"


def _chat(
    gateway: Gateway,
    model: str,
    marker: str,
    extra: Mapping[str, JsonValue] = MappingProxyType({}),
    *,
    key: str | None = None,
) -> httpx.Response:
    return gateway.request(
        "POST",
        "/v1/chat/completions",
        {"model": model, "messages": [{"role": "user", "content": f"say {marker}"}], **extra},
        key=key,
        headers=_CLOSE,
    )


def _carrying(requests: Sequence[Request], marker: str) -> tuple[Request, ...]:
    return tuple(request for request in requests if marker in request.body.decode(errors="replace"))


def _bearers(requests: Sequence[Request]) -> tuple[str | None, ...]:
    return tuple(request.headers.get("authorization") for request in requests)


def _anywhere(requests: Sequence[Request], value: str) -> bool:
    return any(
        value in json.dumps(dict(request.headers)) + request.body.decode(errors="replace") for request in requests
    )


def _reached_on_worker(rig: _Rig, model: str, idle: int) -> tuple[int, tuple[Request, ...]]:
    marker: Final = _marker()
    gateway: Final = rig.owned.gateway
    with (
        rig.owned.suspended(idle),
        httpx.Client(base_url=gateway.client.base_url, timeout=60, trust_env=False, limits=_FRESH) as client,
    ):
        response: Final = client.post(
            "/v1/chat/completions",
            json={"model": model, "messages": [{"role": "user", "content": f"say {marker}"}]},
            headers={"Authorization": f"Bearer {gateway.key}"},
        )
    return response.status_code, _carrying(rig.provider.requests(), marker)


def _chat_on_worker(rig: _Rig, model: str, idle: int) -> tuple[int, tuple[str | None, ...]]:
    status, reached = _reached_on_worker(rig, model, idle)
    return status, _bearers(reached)


@pytest.mark.parametrize("source", _SOURCES)
def test_admin_reference_sends_the_current_token_from_every_worker(rig: _Rig, source: str) -> None:
    first, second = sorted(eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2))
    model: Final = rig.model(source)
    before: Final = rig.secrets.current()
    sent_before: Final = (_chat_on_worker(rig, model, second), _chat_on_worker(rig, model, first))
    rotated: Final = rig.secrets.rotate()
    sent_after: Final = (_chat_on_worker(rig, model, second), _chat_on_worker(rig, model, first))
    assert sent_before == ((200, (f"Bearer {before}",)),) * 2, sent_before
    assert sent_after == ((200, (f"Bearer {rotated}",)),) * 2, sent_after


@pytest.mark.parametrize("endpoint", tuple(_ENDPOINTS))
def test_other_call_types_send_the_current_token(rig: _Rig, endpoint: str) -> None:
    marker: Final = _marker()
    path, model = _ENDPOINTS[endpoint]
    response: Final = rig.owned.gateway.request(
        "POST", path, {"model": model, "input": f"say {marker}"}, headers=_CLOSE
    )
    sent: Final = _bearers(_carrying(rig.provider.requests(), marker))
    assert (response.status_code, sent) == (200, (f"Bearer {rig.secrets.current()}",)), response.text


@pytest.mark.parametrize("shape", ("credential_name", "inline"))
def test_connection_with_an_admin_reference_reaches_the_provider(rig: _Rig, shape: str) -> None:
    params: Final[dict[str, JsonValue]] = (
        _named(_CHAT_MODEL, "wif-config")
        if shape == "credential_name"
        else {"model": _CHAT_MODEL, "api_key": _reference(rig.secrets.token), "api_base": rig.provider_base}
    )
    seen: Final = len(rig.provider.requests())
    response: Final = rig.owned.gateway.request(
        "POST", "/health/test_connection", {"litellm_params": params, "mode": "chat"}, headers=_CLOSE
    )
    sent: Final = frozenset(_bearers(rig.provider.requests()[seen:]))
    assert response.status_code == 200, response.text
    assert JSON_OBJECT.validate_json(response.content)["status"] == "success", response.text
    assert sent == {f"Bearer {rig.secrets.current()}"}, sent


@pytest.mark.parametrize("case", tuple(_UNREADABLE))
def test_unreadable_reference_fails_closed_before_any_upstream_call(rig: _Rig, case: str) -> None:
    marker: Final = _marker()
    response: Final = _chat(rig.owned.gateway, _UNREADABLE[case], marker)
    outside_secret: Final = rig.secrets.outside.read_text()
    reached: Final = _carrying(rig.provider.requests(), marker)
    assert (response.status_code, reached) == (401, ()), (response.text, _bearers(reached))
    assert outside_secret not in response.text
    assert not _anywhere(rig.provider.requests(), outside_secret)


def test_unreadable_reference_falls_over_to_the_next_deployment(rig: _Rig) -> None:
    marker: Final = _marker()
    response: Final = _chat(rig.owned.gateway, "wif-missing-fallback", marker)
    sent: Final = frozenset(_bearers(_carrying(rig.provider.requests(), marker)))
    assert (response.status_code, sent) == (200, {f"Bearer {rig.secrets.current()}"}), response.text


def test_local_token_count_does_not_need_an_unreadable_reference(rig: _Rig) -> None:
    marker: Final = _marker()
    response: Final = rig.owned.gateway.request(
        "POST",
        "/utils/token_counter",
        {"model": _UNREADABLE["missing_file"], "messages": [{"role": "user", "content": f"say {marker}"}]},
        headers=_CLOSE,
    )
    assert response.status_code == 200, response.text
    counted: Final = JSON_OBJECT.validate_json(response.content)["total_tokens"]
    assert isinstance(counted, int) and counted > 0, response.text
    assert _carrying(rig.provider.requests(), marker) == ()


def test_input_token_count_with_an_unreadable_reference_fails_closed(rig: _Rig) -> None:
    marker: Final = _marker()
    response: Final = rig.owned.gateway.request(
        "POST",
        "/v1/responses/input_tokens",
        {"model": _UNREADABLE["missing_file"], "input": f"say {marker}"},
        headers=_CLOSE,
    )
    reached: Final = _carrying(rig.provider.requests(), marker)
    assert (response.status_code, reached) == (401, ()), (response.text, _bearers(reached))
    assert str(rig.secrets.allowed) not in response.text, response.text


def test_messages_with_an_unreadable_reference_hides_the_token_path(rig: _Rig) -> None:
    marker: Final = _marker()
    response: Final = rig.owned.gateway.request(
        "POST",
        "/v1/messages",
        {
            "model": _UNREADABLE["missing_file"],
            "max_tokens": 16,
            "messages": [{"role": "user", "content": f"say {marker}"}],
        },
        headers=_CLOSE,
    )
    reached: Final = _carrying(rig.provider.requests(), marker)
    assert (response.status_code, reached) == (401, ()), (response.text, _bearers(reached))
    assert str(rig.secrets.allowed) not in response.text, response.text


def test_unreadable_reference_under_run_secrets_hides_the_token_path(rig: _Rig) -> None:
    marker: Final = _marker()
    response: Final = _chat(rig.owned.gateway, "wif-run-secrets", marker)
    reached: Final = _carrying(rig.provider.requests(), marker)
    assert (response.status_code, reached) == (401, ()), (response.text, _bearers(reached))
    assert _RUN_SECRETS_TOKEN not in response.text, response.text


@pytest.mark.parametrize("source", ("config_credential", "config_inline"))
def test_provider_token_count_sends_the_current_token(rig: _Rig, source: str) -> None:
    marker: Final = _marker()
    response: Final = rig.owned.gateway.request(
        "POST", "/v1/responses/input_tokens", {"model": rig.model(source), "input": f"say {marker}"}, headers=_CLOSE
    )
    sent: Final = _bearers(_carrying(rig.provider.requests(), marker))
    assert (response.status_code, sent) == (200, (f"Bearer {rig.secrets.current()}",)), response.text
    assert JSON_OBJECT.validate_json(response.content)["input_tokens"] == _PROVIDER_INPUT_TOKENS, response.text


def test_pass_through_with_an_unreadable_reference_hides_the_token_path(rig: _Rig) -> None:
    marker: Final = _marker()
    response: Final = rig.owned.gateway.request(
        "POST",
        "/mistral/v1/chat/completions",
        {"model": "mistral-small-latest", "messages": [{"role": "user", "content": f"say {marker}"}]},
        headers=_CLOSE,
    )
    assert (response.status_code, _carrying(rig.provider.requests(), marker)) == (401, ()), response.text
    assert str(rig.secrets.allowed) not in response.text, response.text


def test_admin_configured_endpoint_receives_the_current_token(rig: _Rig) -> None:
    marker: Final = _marker()
    response: Final = _chat(rig.owned.gateway, "wif-bedrock-endpoint", marker)
    sent: Final = _bearers(_carrying(rig.provider.requests(), marker))
    assert (response.status_code, sent) == (200, (f"Bearer {rig.secrets.current()}",)), response.text


def _hosts_and_bearers(requests: Sequence[Request]) -> tuple[tuple[str | None, str | None], ...]:
    return tuple((request.headers.get("host"), request.headers.get("authorization")) for request in requests)


def test_updated_api_base_carries_the_token_to_the_new_host(rig: _Rig) -> None:
    first, second = sorted(eventually(rig.owned.worker_pids, lambda pids: len(pids) == 2))
    moved_base: Final = rig.provider_base.replace("127.0.0.1", "localhost")
    with rig.owned.gateway.scenario() as scenario:
        name, identity = _created_model(
            rig.owned.gateway,
            scenario,
            _deployment(
                f"db-moved-{uuid.uuid4().hex}",
                {"model": _CHAT_MODEL, "api_key": _reference(rig.secrets.token), "api_base": rig.provider_base},
            ),
        )
        eventually(partial(rig.owned.serves_only_on_every_worker, name, identity), bool, seconds=120)
        patched: Final = rig.owned.gateway.request(
            "PATCH", f"/model/{identity}/update", {"litellm_params": {"api_base": moved_base}}
        )
        assert patched.status_code == 200, patched.text
        expected: Final = ((200, ((httpx.URL(moved_base).netloc.decode(), f"Bearer {rig.secrets.current()}"),)),) * 2
        reached: Final = eventually(
            lambda: tuple(
                (status, _hosts_and_bearers(requests))
                for status, requests in (_reached_on_worker(rig, name, idle) for idle in (second, first))
            ),
            lambda seen: seen == expected,
            seconds=120,
            return_last_on_timeout=True,
        )
        assert reached == expected, reached


def _attack(rig: _Rig, attack: str) -> tuple[str, dict[str, JsonValue]]:
    redirected: Final[dict[str, JsonValue]] = {"api_base": rig.intruder_base}
    match attack:
        case "client_api_base_inline":
            return "wif-config-inline", redirected
        case "client_api_base_credential":
            return "wif-config-credential", redirected
        case "body_reference":
            return "wif-config-inline", {**redirected, "api_key": _reference(rig.secrets.token)}
        case "body_credential_name":
            return "wif-config-inline", {**redirected, "litellm_credential_name": "wif-config"}
        case "user_config":
            injected: Final = {"model": _CHAT_MODEL, "api_key": _reference(rig.secrets.token), **redirected}
            return "wif-injected", {
                "user_config": {"model_list": [{"model_name": "wif-injected", "litellm_params": injected}]}
            }
        case _:
            pytest.fail(f"unknown attack {attack}")


@pytest.mark.parametrize("attack", _ATTACKS)
def test_caller_cannot_send_the_admin_token_to_its_own_host(rig: _Rig, attack: str) -> None:
    marker: Final = _marker()
    model, extra = _attack(rig, attack)
    response: Final = _chat(rig.owned.gateway, model, marker, extra)
    intruded: Final = rig.intruder.requests()
    assert not _anywhere(intruded, rig.secrets.current())
    assert _carrying(intruded, marker) == (), _bearers(_carrying(intruded, marker))
    assert _carrying(rig.provider.requests(), marker) == ()
    assert response.status_code in (400, 401), response.text


def test_caller_with_its_own_key_still_reaches_its_own_host(rig: _Rig) -> None:
    marker: Final = _marker()
    own_key: Final = f"sk-caller-{uuid.uuid4().hex}"
    response: Final = _chat(
        rig.owned.gateway, "wif-config-inline", marker, {"api_base": rig.intruder_base, "api_key": own_key}
    )
    intruded: Final = rig.intruder.requests()
    assert (response.status_code, _bearers(_carrying(intruded, marker))) == (200, (f"Bearer {own_key}",)), response.text
    assert not _anywhere(intruded, rig.secrets.current())


def _team_key(scenario: Scenario, team_id: str, role: str) -> str:
    return scenario.key(user_id=scenario.member(team_id, role=role), team_id=team_id)


def _team_body(name: str, team_id: str, params: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
    return _deployment(name, {"model": _CHAT_MODEL, **params}, {"team_id": team_id})


def _team_write_params(rig: _Rig, shape: str) -> dict[str, JsonValue]:
    match shape:
        case "plain_key":
            return {"api_key": f"sk-team-{uuid.uuid4().hex}", "api_base": rig.provider_base}
        case "names_credential":
            return {"litellm_credential_name": "wif-config"}
        case "inline_reference":
            return {"api_key": _reference(rig.secrets.token), "api_base": rig.provider_base}
        case _:
            pytest.fail(f"unknown shape {shape}")


def _delete_written(scenario: Scenario, response: httpx.Response) -> None:
    if response.status_code == 200:
        scenario.delete_model(
            string_value(object_value(JSON_OBJECT.validate_json(response.content)["model_info"])["id"])
        )


@pytest.mark.parametrize(("writer", "shape"), tuple(_TEAM_WRITES))
def test_team_model_write_carrying_a_reference_is_refused(rig: _Rig, writer: str, shape: str) -> None:
    with rig.owned.gateway.scenario() as scenario:
        team_id: Final = scenario.team()
        name: Final = f"wif-team-{uuid.uuid4().hex}"
        key: Final = None if writer == "proxy_admin" else _team_key(scenario, team_id, "admin")
        response: Final = rig.owned.gateway.request(
            "POST", "/model/new", _team_body(name, team_id, _team_write_params(rig, shape)), key=key
        )
        stored: Final = len(rig.owned.gateway.rows(_MODEL_ROWS, (name,)))
        _delete_written(scenario, response)
        expected: Final = _TEAM_WRITES[(writer, shape)]
        assert (response.status_code, stored) == (expected, int(expected == 200)), response.text


@pytest.mark.parametrize("host", ("credential_host", "other_host"))
def test_team_row_naming_the_credential_sends_the_token_only_to_its_host(rig: _Rig, host: str) -> None:
    redirected: Final[dict[str, JsonValue]] = {} if host == "credential_host" else {"api_base": rig.intruder_base}
    with rig.owned.gateway.scenario() as scenario:
        team_id: Final = scenario.team()
        name: Final = f"wif-team-{uuid.uuid4().hex}"
        stored_name, identity = _created_model(
            rig.owned.gateway,
            scenario,
            _team_body(name, team_id, {"litellm_credential_name": "wif-config", **redirected}),
        )
        eventually(lambda: rig.owned.serves_only_on_every_worker(stored_name, identity), bool, seconds=120)
        marker: Final = _marker()
        response: Final = _chat(rig.owned.gateway, name, marker, key=_team_key(scenario, team_id, "user"))
        intruded: Final = rig.intruder.requests()
        sent: Final = _bearers(_carrying(rig.provider.requests(), marker))
        assert not _anywhere(intruded, rig.secrets.current())
        assert _carrying(intruded, marker) == ()
        expected: Final = (200, (f"Bearer {rig.secrets.current()}",)) if host == "credential_host" else (401, ())
        assert (response.status_code, sent) == expected, response.text


def test_startup_discovery_sends_the_boot_token(rig: _Rig) -> None:
    sent: Final = frozenset(_bearers(tuple(request for request in rig.boot if _is_discovery(request))))
    assert sent - rig.google.bearers() == {f"Bearer {rig.boot_token}"}, sent


def _discovery_bearers(rig: _Rig) -> frozenset[str | None]:
    return frozenset(_bearers(tuple(request for request in rig.provider.requests() if _is_discovery(request))))


def test_startup_discovery_sends_the_google_token(rig: _Rig) -> None:
    expected: Final = frozenset(
        f"Bearer {rig.google.token(_GOOGLE_AUDIENCES[source])}" for source in ("config_inline", "config_credential")
    )
    sent: Final = eventually(
        partial(_discovery_bearers, rig), lambda seen: expected <= seen, seconds=30, return_last_on_timeout=True
    )
    assert expected <= sent, sent


def _google_model(rig: _Rig, scenario: Scenario, source: str) -> str:
    match source:
        case "config_inline":
            return "wif-google-inline"
        case "config_credential":
            return "wif-google-credential"
        case "db_credential":
            credential: Final = f"db-google-{uuid.uuid4().hex}"
            rig.owned.gateway.post(
                "/credentials",
                _credential(credential, _google_reference(_GOOGLE_AUDIENCES[source]), rig.provider_base),
            )
            scenario.cleanups.callback(_delete_credential, rig.owned.gateway, credential)
            name, identity = _created_model(
                rig.owned.gateway,
                scenario,
                _deployment(f"db-google-{uuid.uuid4().hex}", _named(_GOOGLE_MODEL, credential)),
            )
            eventually(partial(rig.owned.serves_only_on_every_worker, name, identity), bool, seconds=120)
            return name
        case _:
            pytest.fail(f"unknown source {source}")


def _asked_audiences(rig: _Rig) -> frozenset[str | None]:
    return frozenset(httpx.URL(request.target).params.get("audience") for request in rig.metadata.requests())


@pytest.mark.parametrize("source", tuple(_GOOGLE_AUDIENCES))
def test_google_reference_sends_the_metadata_server_token_for_its_audience(rig: _Rig, source: str) -> None:
    audience: Final = _GOOGLE_AUDIENCES[source]
    with rig.owned.gateway.scenario() as scenario:
        marker: Final = _marker()
        response: Final = _chat(rig.owned.gateway, _google_model(rig, scenario, source), marker)
        sent: Final = _bearers(_carrying(rig.provider.requests(), marker))
    assert (response.status_code, sent) == (200, (f"Bearer {rig.google.token(audience)}",)), response.text
    assert audience in _asked_audiences(rig)


def test_google_reference_fails_closed_when_the_metadata_server_fails(rig: _Rig) -> None:
    marker: Final = _marker()
    response: Final = _chat(rig.owned.gateway, "wif-google-unavailable", marker)
    reached: Final = _carrying(rig.provider.requests(), marker)
    assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached))
    assert "oidc/google secret reference could not be resolved" in response.text, response.text
    assert _UNAVAILABLE_AUDIENCE.removeprefix("https://") not in response.text, response.text
    assert _UNAVAILABLE_AUDIENCE in _asked_audiences(rig)


def test_google_reference_answers_503_when_the_metadata_server_never_finishes_its_reply(rig: _Rig) -> None:
    marker: Final = _marker()
    response: Final = _chat(rig.owned.gateway, "wif-google-hung", marker)
    reached: Final = _carrying(rig.provider.requests(), marker)
    assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached))
    assert "the metadata server did not return an ID token" in response.text, response.text
    assert _HUNG_AUDIENCE in _asked_audiences(rig)


def test_input_token_count_answers_503_when_the_metadata_server_fails(rig: _Rig) -> None:
    marker: Final = _marker()
    response: Final = rig.owned.gateway.request(
        "POST",
        "/v1/responses/input_tokens",
        {"model": "wif-google-unavailable-count", "input": f"say {marker}"},
        headers=_CLOSE,
    )
    reached: Final = _carrying(rig.provider.requests(), marker)
    assert (response.status_code, reached) == (503, ()), (response.text, _bearers(reached))
    assert "oidc/google secret reference could not be resolved" in response.text, response.text


def _side_call(rig: _Rig, seam: str, marker: str) -> tuple[httpx.Response, Callable[[Request], bool]]:
    gateway: Final = rig.owned.gateway
    match seam:
        case "realtime_health":
            return (
                gateway.request("GET", "/health", params={"model": "wif-realtime"}, headers=_CLOSE),
                lambda request: request.target.startswith("/v1/realtime"),
            )
        case "vector_store_search":
            return (
                gateway.request(
                    "POST", f"/v1/vector_stores/{_VECTOR_STORE}/search", {"query": f"say {marker}"}, headers=_CLOSE
                ),
                lambda request: marker in request.body.decode(errors="replace"),
            )
        case "batch_probe":
            return (
                gateway.request(
                    "POST",
                    "/v1/batches",
                    {
                        "input_file_id": f"file-{marker}",
                        "endpoint": "/v1/chat/completions",
                        "completion_window": "24h",
                        "model": _BATCH_DEPLOYMENT,
                    },
                    headers=_CLOSE,
                ),
                lambda request: request.target == "/v1/files",
            )
        case _:
            pytest.fail(f"unknown seam {seam}")


_SIDE_STATUS: Final = MappingProxyType({"realtime_health": 503, "vector_store_search": 200, "batch_probe": 400})


@pytest.mark.parametrize("seam", tuple(_SIDE_STATUS))
def test_calls_outside_the_client_wrapper_send_the_current_token(rig: _Rig, seam: str) -> None:
    seen: Final = len(rig.provider.requests())
    response, matches = _side_call(rig, seam, _marker())
    sent: Final = frozenset(_bearers(tuple(request for request in rig.provider.requests()[seen:] if matches(request))))
    assert (response.status_code, sent) == (_SIDE_STATUS[seam], {f"Bearer {rig.secrets.current()}"}), response.text


def test_wildcard_listing_shows_the_models_the_provider_lists(rig: _Rig) -> None:
    response: Final = rig.owned.gateway.request("GET", "/v1/models", headers=_CLOSE)
    assert response.status_code == 200, response.text
    entries: Final = list_value(JSON_OBJECT.validate_json(response.content)["data"])
    listed: Final = frozenset(string_value(object_value(entry)["id"]) for entry in entries)
    asked: Final = frozenset(_bearers(tuple(request for request in rig.provider.requests() if _is_discovery(request))))
    assert any(_LISTED_MODEL in identity for identity in listed), sorted(listed)
    assert not any(bearer.startswith("Bearer oidc/") for bearer in asked if bearer is not None), asked


def _stored_rows(rig: _Rig) -> tuple[list[dict[str, JsonValue]], ...]:
    return (
        rig.owned.gateway.rows(_CREDENTIAL_ROW, (rig.stored.credential,)),
        *(rig.owned.gateway.rows(_MODEL_ROW, (identity,)) for identity in rig.stored.model_ids.values()),
    )


def test_resolved_token_is_never_written_back(rig: _Rig) -> None:
    before: Final = _stored_rows(rig)
    calls: Final = tuple(
        _chat(rig.owned.gateway, rig.model(source), _marker()).status_code for source in ("db_credential", "db_inline")
    )
    after: Final = _stored_rows(rig)
    assert tuple(len(rows) for rows in before) == (1, 1, 1), before
    assert after == before, (calls, before, after)