litellm/tests/integration/routing/test_ptu_shares.py

import asyncio
import json
import os
import re
import threading
import uuid
from collections.abc import Callable, Iterator, Mapping
from concurrent.futures import Future, ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal

import httpx
import pytest
import yaml
from anthropic import Anthropic
from integration._support.anthropic_sse import (
    LIFECYCLE,
    Attempts,
    delta_text,
    dropping_reply,
    event_types,
    message_id,
    message_json,
    message_stream,
    parse_sse,
    stream_reply,
)
from integration._support.client import (
    JSON_OBJECT,
    Gateway,
    eventually,
    gateway_from_environment,
    object_value,
    string_value,
)
from integration._support.openai_wire import chat_reply, openai_error, responses_reply
from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process
from integration._support.wire import Reply, Request, wire_server
from openai import AsyncOpenAI, OpenAI
from pydantic import JsonValue, TypeAdapter
from redis import Redis

if TYPE_CHECKING:
    from litellm.llms.azure.ptu_capacity import PTUCapacity

pytestmark: Final = [pytest.mark.timeout(2 * graceful_stop_seconds() + 120), pytest.mark.shares_worker_state]

_SHARED: Final = "ptu-shared"
_MIXED: Final = "ptu-mixed"
_PAIR: Final = "ptu-pair"
_CLAUDE: Final = "ptu-claude"
_SHARED_DEPLOYMENT: Final = "shared-gpt41"
_MIXED_SHARED_DEPLOYMENT: Final = "mixed-shared-gpt41"
_MIXED_OPEN_DEPLOYMENT: Final = "mixed-open-gpt41"
_PAIR_DEPLOYMENTS: Final = ("pair-a-gpt41", "pair-b-gpt41")
_CLAUDE_BACKEND: Final = "claude-ptu-backend"
_AZURE_KEY: Final = "synthetic-azure-key"
_ANTHROPIC_KEY: Final = "synthetic-anthropic-key"
_UPSTREAM_URL_PLACEHOLDER: Final = "upstream-url-placeholder"
_ATTRIBUTION_ON: Final = MappingProxyType({"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "True"})
_ATTRIBUTION_OFF: Final = MappingProxyType({"LITELLM_ENABLE_PTU_COST_ATTRIBUTION": "False"})
_RESERVED: Final = "is reserved for the teams holding a PTU share of it"
_TEXT: Final = "served"
_SMALL_MAX_TOKENS: Final = 16
_SETTLE_USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType(
    {"prompt_tokens": 5, "completion_tokens": 700, "total_tokens": 705}
)
_CACHED_USAGE: Final[Mapping[str, JsonValue]] = MappingProxyType(
    {
        "prompt_tokens": 1000,
        "completion_tokens": 10,
        "total_tokens": 1010,
        "prompt_tokens_details": {"cached_tokens": 800},
    }
)
_SHARE_ROLES: Final = (
    "holder",
    "settle",
    "cached",
    "refund",
    "twin",
    "oversize",
    "retry",
    "queue",
    "burst",
    "worker",
)
_ROLES: Final = (*_SHARE_ROLES, "outsider", "mixed", "pair", "messages")
_COUNTER: Final = TypeAdapter(int)
_CONFIG: Final = TypeAdapter(dict[str, JsonValue])

Surface = Literal["chat", "responses", "messages"]
Behavior = Literal["ok", "settle", "cached", "fail", "fail-once", "hold", "drop-once"]


def _gpt41() -> "PTUCapacity":
    from litellm.llms.azure.ptu_capacity import AZURE_PTU_CAPACITY

    return AZURE_PTU_CAPACITY["gpt-4.1"]


def _text_of(content: JsonValue) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(_text_of(part) for part in content)
    if isinstance(content, dict):
        return _text_of(content.get("text", content.get("content", "")))
    return ""


def _prompt(body: Mapping[str, JsonValue]) -> str:
    messages: Final = body.get("messages")
    if isinstance(messages, list) and messages:
        return _text_of(messages[-1])
    return _text_of(body.get("input", ""))


def _deployment_of(request: Request, body: Mapping[str, JsonValue]) -> str:
    _, found, rest = request.target.partition("/openai/deployments/")
    return rest.split("/", 1)[0] if found else str(body.get("model"))


def _usage_reply(identity: str, model: str, usage: Mapping[str, JsonValue]) -> Reply:
    return Reply(
        body=json.dumps(
            {
                "id": identity,
                "object": "chat.completion",
                "created": 1,
                "model": model,
                "choices": [{"index": 0, "message": {"role": "assistant", "content": _TEXT}, "finish_reason": "stop"}],
                "usage": dict(usage),
            }
        ).encode()
    )


def _anthropic_reply(identity: str, behavior: str, attempt: int, stream: bool) -> Reply:
    full: Final = message_stream(identity, _CLAUDE_BACKEND, _TEXT)
    if behavior == "drop-once" and attempt == 1:
        return dropping_reply(full, abort_after=1)
    if stream:
        return stream_reply(full)
    return Reply(body=message_json(identity, _CLAUDE_BACKEND, _TEXT))


def _served_id(deployment: str, marker: str, attempt: int) -> str:
    return f"chatcmpl-{deployment}-{marker}-a{attempt}"


def _openai_reply(request: Request, behavior: str, identity: str, deployment: str, attempt: int, stream: bool) -> Reply:
    match behavior:
        case "settle":
            return _usage_reply(identity, deployment, _SETTLE_USAGE)
        case "cached":
            return _usage_reply(identity, deployment, _CACHED_USAGE)
        case "fail":
            return openai_error(500)
        case "fail-once" if attempt == 1:
            return openai_error(500)
        case _:
            pass
    if "/responses" in request.target:
        return responses_reply(identity, deployment, _TEXT, stream=stream)
    return chat_reply(identity, deployment, _TEXT, stream=stream)


def _respond(attempts: Attempts, release: threading.Event) -> Callable[[Request], Reply]:
    def respond(request: Request) -> Reply:
        assert request.method == "POST", request
        body: Final = JSON_OBJECT.validate_json(request.body)
        behavior, _, marker = _prompt(body).partition(":")
        deployment: Final = _deployment_of(request, body)
        attempt: Final = attempts.record(f"{deployment}:{marker}")
        identity: Final = _served_id(deployment, marker, attempt)
        stream: Final = body.get("stream") is True
        if behavior == "hold":
            release.wait(timeout=30)
        if request.target.startswith("/v1/messages"):
            return _anthropic_reply(identity, behavior, attempt, stream)
        return _openai_reply(request, behavior, identity, deployment, attempt, stream)

    return respond


def _ptu_terms(shares: Mapping[str, int]) -> dict[str, JsonValue]:
    return {
        "base_model": "azure/gpt-4.1",
        "ptu_count": sum(shares.values()),
        "cost_per_ptu_per_hour": 2.0,
        "ptu_effective_from": "2026-01-01T00:00:00Z",
        "ptu_shares": dict(shares),
    }


def _azure(model_name: str, deployment: str, model_info: Mapping[str, JsonValue]) -> dict[str, JsonValue]:
    return {
        "model_name": model_name,
        "litellm_params": {
            "model": f"azure/{deployment}",
            "api_base": _UPSTREAM_URL_PLACEHOLDER,
            "api_key": _AZURE_KEY,
            "api_version": "2025-03-01-preview",
        },
        "model_info": {"id": f"{deployment}-{uuid.uuid4().hex}", **model_info},
    }


def _model_list(teams: Mapping[str, str]) -> list[JsonValue]:
    one_each: Final = {teams[role]: 1 for role in _SHARE_ROLES}
    return [
        _azure(_SHARED, _SHARED_DEPLOYMENT, _ptu_terms(one_each)),
        _azure(_MIXED, _MIXED_SHARED_DEPLOYMENT, _ptu_terms({teams["mixed"]: 1})),
        _azure(_MIXED, _MIXED_OPEN_DEPLOYMENT, {"base_model": "azure/gpt-4.1"}),
        *(_azure(_PAIR, deployment, _ptu_terms({teams["pair"]: 1})) for deployment in _PAIR_DEPLOYMENTS),
        {
            "model_name": _CLAUDE,
            "litellm_params": {
                "model": f"anthropic/{_CLAUDE_BACKEND}",
                "api_base": _UPSTREAM_URL_PLACEHOLDER,
                "api_key": _ANTHROPIC_KEY,
            },
            "model_info": {"id": f"claude-{uuid.uuid4().hex}", **_ptu_terms({teams["messages"]: 1})},
        },
    ]


def _config(directory: Path, wire_url: str, model_list: list[JsonValue], name: str) -> Path:
    config: Final = {
        **_CONFIG.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())),
        "model_list": model_list,
        "router_settings": {"disable_cooldowns": True, "num_retries": 0, "ignore_invalid_deployments": True},
    }
    path: Final = directory / f"{name}.yaml"
    path.write_text(yaml.safe_dump(config).replace(_UPSTREAM_URL_PLACEHOLDER, wire_url))
    return path


def _team_ids() -> Mapping[str, str]:
    run: Final = uuid.uuid4().hex[:12]
    return MappingProxyType({role: f"ptu-{role}-{run}" for role in _ROLES})


def _create_team(proxy: Gateway, team_id: str) -> str:
    proxy.post("/team/new", {"team_id": team_id, "team_alias": team_id})
    return string_value(proxy.post("/key/generate", {"team_id": team_id})["key"])


def _delete_teams(gateway: Gateway, team_ids: tuple[str, ...]) -> None:
    gateway.post("/team/delete", {"team_ids": list(team_ids)})


@dataclass(frozen=True, slots=True)
class _Rig:
    owned: OwnedProxy
    attempts: Attempts
    release: threading.Event
    teams: Mapping[str, str]
    keys: Mapping[str, str]

    @property
    def proxy(self) -> Gateway:
        return self.owned.gateway

    def chat(self, model: str, role: str | None, behavior: Behavior, marker: str, **extra: JsonValue) -> httpx.Response:
        return self.proxy.request(
            "POST",
            "/v1/chat/completions",
            {
                "model": model,
                "max_tokens": _SMALL_MAX_TOKENS,
                "messages": [{"role": "user", "content": f"{behavior}:{marker}"}],
                **extra,
            },
            key=self.keys[role] if role is not None else None,
        )

    def send(self, surface: Surface, model: str, role: str | None, marker: str) -> httpx.Response:
        key: Final = self.keys[role] if role is not None else None
        prompt: Final = f"ok:{marker}"
        match surface:
            case "chat":
                return self.chat(model, role, "ok", marker)
            case "responses":
                return self.proxy.request(
                    "POST",
                    "/v1/responses",
                    {"model": model, "input": prompt, "max_output_tokens": _SMALL_MAX_TOKENS},
                    key=key,
                )
            case "messages":
                return self.proxy.request(
                    "POST",
                    "/v1/messages",
                    {
                        "model": model,
                        "max_tokens": _SMALL_MAX_TOKENS,
                        "messages": [{"role": "user", "content": prompt}],
                    },
                    key=key,
                )

    def hits(self, deployment: str, marker: str) -> int:
        return self.attempts.count(f"{deployment}:{marker}")


@pytest.fixture(scope="module")
def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]:
    attempts: Final = Attempts()
    release: Final = threading.Event()
    teams: Final = _team_ids()
    directory: Final = tmp_path_factory.mktemp("ptu-shares")
    with gateway_from_environment() as gateway, wire_server(_respond(attempts, release)) as wire:
        config: Final = _config(directory, wire.url, _model_list(teams), "ptu-shares")
        try:
            with owned_proxy_process(gateway, directory, _ATTRIBUTION_ON, config=config, workers=2) as owned:
                keys: Final = MappingProxyType(
                    {role: _create_team(owned.gateway, team) for role, team in teams.items()}
                )
                yield _Rig(owned, attempts, release, teams, keys)
        finally:
            release.set()
            _delete_teams(gateway, tuple(teams.values()))


def _ptu_counter(team_id: str, model_group: str) -> int | None:
    with Redis(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]), decode_responses=True) as client:
        keys: Final = tuple(client.scan_iter(match=f"*model_per_team_ptu:{team_id}:{model_group}}}:tokens"))
        assert len(keys) <= 1, keys
        raw: Final[object] = client.get(keys[0]) if keys else None
    return None if raw is None else _COUNTER.validate_python(raw)


def _settled_counter(team_id: str, model_group: str, expected: int) -> int | None:
    return eventually(
        lambda: _ptu_counter(team_id, model_group),
        lambda observed: observed == expected,
        seconds=15,
        return_last_on_timeout=True,
    )


def _error_message(response: httpx.Response) -> str:
    return string_value(object_value(JSON_OBJECT.validate_json(response.content)["error"])["message"])


@pytest.mark.parametrize("client", ("sync-sdk", "async-sdk-stream"))
def test_a_share_holder_is_served_chat_by_the_shared_deployment(rig: _Rig, client: str) -> None:
    marker: Final = uuid.uuid4().hex
    base_url: Final = f"{rig.proxy.client.base_url}/v1"
    if client == "sync-sdk":
        with OpenAI(api_key=rig.keys["holder"], base_url=base_url, max_retries=0, timeout=30) as sdk:
            completion: Final = sdk.chat.completions.create(
                model=_SHARED,
                max_tokens=_SMALL_MAX_TOKENS,
                messages=[{"role": "user", "content": f"ok:{marker}"}],
            )
        assert (completion.id, completion.choices[0].message.content) == (
            _served_id(_SHARED_DEPLOYMENT, marker, 1),
            _TEXT,
        ), completion

    async def stream() -> tuple[str, str]:
        async with AsyncOpenAI(api_key=rig.keys["holder"], base_url=base_url, max_retries=0, timeout=30) as sdk:
            chunks: Final = await sdk.chat.completions.create(
                model=_SHARED,
                max_tokens=_SMALL_MAX_TOKENS,
                messages=[{"role": "user", "content": f"ok:{marker}"}],
                stream=True,
            )
            received: Final = tuple([chunk async for chunk in chunks])
        return (
            received[0].id,
            "".join(choice.delta.content or "" for chunk in received for choice in chunk.choices),
        )

    if client == "async-sdk-stream":
        assert asyncio.run(stream()) == (_served_id(_SHARED_DEPLOYMENT, marker, 1), _TEXT)
    assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1


@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
def test_a_share_holder_is_served_responses_by_the_shared_deployment(rig: _Rig, stream: bool) -> None:
    marker: Final = uuid.uuid4().hex
    response: Final = rig.proxy.request(
        "POST",
        "/v1/responses",
        {"model": _SHARED, "input": f"ok:{marker}", "max_output_tokens": _SMALL_MAX_TOKENS, "stream": stream},
        key=rig.keys["holder"],
    )
    assert response.status_code == 200, response.text
    assert _served_id(_SHARED_DEPLOYMENT, marker, 1) in response.text, response.text
    assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1


def _anthropic_text(sdk: Anthropic, marker: str, *, stream: bool) -> str:
    if stream:
        with sdk.messages.stream(
            model=_SHARED,
            max_tokens=_SMALL_MAX_TOKENS,
            messages=[{"role": "user", "content": f"ok:{marker}"}],
        ) as events:
            return "".join(events.text_stream)
    message: Final = sdk.messages.create(
        model=_SHARED,
        max_tokens=_SMALL_MAX_TOKENS,
        messages=[{"role": "user", "content": f"ok:{marker}"}],
    )
    return "".join(block.text for block in message.content if block.type == "text")


@pytest.mark.parametrize("stream", (False, True), ids=("json", "stream"))
def test_a_share_holder_is_served_messages_by_the_shared_deployment(rig: _Rig, stream: bool) -> None:
    marker: Final = uuid.uuid4().hex
    with Anthropic(api_key=rig.keys["holder"], base_url=str(rig.proxy.client.base_url), max_retries=0) as sdk:
        received: Final = _anthropic_text(sdk, marker, stream=stream)
    assert received == _TEXT
    assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1


@pytest.mark.parametrize("caller", ("outsider", "master-key"))
@pytest.mark.parametrize("surface", ("chat", "responses", "messages"))
def test_a_caller_without_a_share_is_refused_the_shared_deployment(rig: _Rig, caller: str, surface: Surface) -> None:
    marker: Final = uuid.uuid4().hex
    response: Final = rig.send(surface, _SHARED, "outsider" if caller == "outsider" else None, marker)
    assert response.status_code == 400, response.text
    assert f"Deployment {_SHARED} {_RESERVED}" in response.text, response.text
    assert rig.hits(_SHARED_DEPLOYMENT, marker) == 0


def test_a_mixed_group_serves_each_caller_from_its_own_deployment(rig: _Rig) -> None:
    served: Final = {
        caller: rig.chat(_MIXED, role, "ok", marker)
        for caller, role, marker in (
            ("holder", "mixed", uuid.uuid4().hex),
            ("outsider", "outsider", uuid.uuid4().hex),
            ("master", None, uuid.uuid4().hex),
        )
    }
    assert {caller: response.status_code for caller, response in served.items()} == dict.fromkeys(served, 200), {
        caller: response.text for caller, response in served.items()
    }
    deployments: Final = {
        caller: re.sub(
            r"^chatcmpl-(.+)-[0-9a-f]{32}-a1$", r"\1", string_value(JSON_OBJECT.validate_json(response.content)["id"])
        )
        for caller, response in served.items()
    }
    assert deployments == {
        "holder": _MIXED_SHARED_DEPLOYMENT,
        "outsider": _MIXED_OPEN_DEPLOYMENT,
        "master": _MIXED_OPEN_DEPLOYMENT,
    }


@pytest.mark.parametrize(("model_group", "role", "ptus"), ((_SHARED, "oversize", 1), (_PAIR, "pair", 2)))
def test_a_request_bigger_than_the_whole_share_is_refused_up_front(
    rig: _Rig, model_group: str, role: str, ptus: int
) -> None:
    capacity: Final = _gpt41()
    ceiling: Final = capacity.input_tpm_for(ptus)
    marker: Final = uuid.uuid4().hex
    response: Final = rig.chat(model_group, role, "ok", marker, max_tokens=ceiling)
    assert response.status_code == 400, response.text
    found: Final = re.search(
        rf"This request needs (\d+) normalized tokens, more than your team's PTU share of {model_group} "
        rf"allows in a minute \({ceiling}\), so it can never be admitted\. Lower max_tokens or shorten the prompt\.",
        _error_message(response),
    )
    assert found is not None, response.text
    assert int(found.group(1)) >= capacity.output_to_input_ratio * ceiling, response.text
    assert rig.hits(_SHARED_DEPLOYMENT, marker) + sum(rig.hits(name, marker) for name in _PAIR_DEPLOYMENTS) == 0
    assert _ptu_counter(rig.teams[role], model_group) in (None, 0)


def test_usage_settles_in_normalized_tokens_and_throttles_the_next_request(rig: _Rig) -> None:
    capacity: Final = _gpt41()
    ceiling: Final = capacity.input_tpm_for(1)
    team: Final = rig.teams["settle"]
    first: Final = uuid.uuid4().hex
    settled: Final = rig.chat(_SHARED, "settle", "settle", first)
    assert settled.status_code == 200, settled.text
    assert settled.json()["usage"]["completion_tokens"] == _SETTLE_USAGE["completion_tokens"], settled.text
    normalized: Final = round(
        int(str(_SETTLE_USAGE["prompt_tokens"]))
        + capacity.output_to_input_ratio * int(str(_SETTLE_USAGE["completion_tokens"]))
    )
    assert _settled_counter(team, _SHARED, normalized) == normalized
    max_tokens: Final = int((ceiling - normalized) / capacity.output_to_input_ratio) + 1
    second: Final = uuid.uuid4().hex
    throttled: Final = rig.chat(_SHARED, "settle", "ok", second, max_tokens=max_tokens)
    assert throttled.status_code == 429, throttled.text
    assert (
        f"Rate limit exceeded for model_per_team_ptu: {team}:{_SHARED}. Limit type: tokens. Current limit: {ceiling},"
        in _error_message(throttled)
    ), throttled.text
    assert int(str(_SETTLE_USAGE["total_tokens"])) + max_tokens * capacity.output_to_input_ratio < ceiling
    assert (rig.hits(_SHARED_DEPLOYMENT, first), rig.hits(_SHARED_DEPLOYMENT, second)) == (1, 0)


def test_cached_input_is_weighted_by_the_cached_ratio(rig: _Rig) -> None:
    capacity: Final = _gpt41()
    marker: Final = uuid.uuid4().hex
    response: Final = rig.chat(_SHARED, "cached", "cached", marker)
    assert response.status_code == 200, response.text
    cached: Final = 800
    expected: Final = round(
        (int(str(_CACHED_USAGE["prompt_tokens"])) - cached)
        + capacity.cached_input_ratio * cached
        + capacity.output_to_input_ratio * int(str(_CACHED_USAGE["completion_tokens"]))
    )
    assert _settled_counter(rig.teams["cached"], _SHARED, expected) == expected
    assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1


def test_a_failed_upstream_call_refunds_the_reservation(rig: _Rig) -> None:
    marker: Final = uuid.uuid4().hex
    response: Final = rig.chat(_SHARED, "refund", "fail", marker)
    assert response.status_code == 500, response.text
    assert "scripted 500" in response.text, response.text
    assert _settled_counter(rig.teams["refund"], _SHARED, 0) == 0
    assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1


def test_a_response_cache_hit_settles_against_the_share_like_every_token_limit(rig: _Rig) -> None:
    capacity: Final = _gpt41()
    marker: Final = uuid.uuid4().hex
    first: Final = rig.chat(_SHARED, "twin", "settle", marker)
    normalized: Final = round(
        int(str(_SETTLE_USAGE["prompt_tokens"]))
        + capacity.output_to_input_ratio * int(str(_SETTLE_USAGE["completion_tokens"]))
    )
    assert _settled_counter(rig.teams["twin"], _SHARED, normalized) == normalized
    repeat: Final = rig.chat(_SHARED, "twin", "settle", marker)
    assert (first.status_code, repeat.status_code) == (200, 200), (first.text, repeat.text)
    assert repeat.json()["id"] == first.json()["id"] == _served_id(_SHARED_DEPLOYMENT, marker, 1)
    assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1
    assert _settled_counter(rig.teams["twin"], _SHARED, 2 * normalized) == 2 * normalized


def test_a_cached_response_is_not_served_to_a_caller_without_a_share(rig: _Rig) -> None:
    marker: Final = uuid.uuid4().hex
    cached: Final = rig.chat(_SHARED, "holder", "ok", marker)
    assert cached.status_code == 200, cached.text
    refused: Final = rig.chat(_SHARED, "outsider", "ok", marker)
    assert refused.status_code == 400, refused.text
    assert _RESERVED in refused.text, refused.text
    assert rig.hits(_SHARED_DEPLOYMENT, marker) == 1


def test_a_retry_after_an_upstream_error_stays_on_the_shared_deployment(rig: _Rig) -> None:
    marker: Final = uuid.uuid4().hex
    response: Final = rig.chat(_SHARED, "retry", "fail-once", marker, num_retries=1)
    assert response.status_code == 200, response.text
    assert response.json()["id"] == _served_id(_SHARED_DEPLOYMENT, marker, 2), response.text
    assert rig.hits(_SHARED_DEPLOYMENT, marker) == 2


def test_the_priority_scheduler_serves_share_holders_and_refuses_everyone_else(rig: _Rig) -> None:
    served_marker: Final = uuid.uuid4().hex
    refused_marker: Final = uuid.uuid4().hex
    served: Final = rig.chat(_SHARED, "queue", "ok", served_marker, priority=1)
    refused: Final = rig.chat(_SHARED, "outsider", "ok", refused_marker, priority=1)
    assert served.status_code == 200, served.text
    assert served.json()["id"] == _served_id(_SHARED_DEPLOYMENT, served_marker, 1), served.text
    assert "x-litellm-request-prioritization-used" in served.headers, served.headers
    assert refused.status_code == 400, refused.text
    assert _RESERVED in refused.text, refused.text
    assert (rig.hits(_SHARED_DEPLOYMENT, served_marker), rig.hits(_SHARED_DEPLOYMENT, refused_marker)) == (1, 0)


def test_a_dropped_messages_stream_retries_on_the_shared_deployment(rig: _Rig) -> None:
    marker: Final = uuid.uuid4().hex
    response: Final = rig.proxy.request(
        "POST",
        "/v1/messages",
        {
            "model": _CLAUDE,
            "max_tokens": _SMALL_MAX_TOKENS,
            "stream": True,
            "num_retries": 1,
            "messages": [{"role": "user", "content": f"drop-once:{marker}"}],
        },
        key=rig.keys["messages"],
    )
    assert response.status_code == 200, response.text
    events: Final = parse_sse(response.text)
    assert event_types(events) == LIFECYCLE, events
    assert message_id(events) == _served_id(_CLAUDE_BACKEND, marker, 2), events
    assert delta_text(events) == _TEXT, events
    assert rig.hits(_CLAUDE_BACKEND, marker) == 2


def test_concurrent_requests_admit_only_what_fits_under_the_share(rig: _Rig) -> None:
    capacity: Final = _gpt41()
    ceiling: Final = capacity.input_tpm_for(1)
    max_tokens: Final = int(ceiling / (2.5 * capacity.output_to_input_ratio))
    markers: Final = tuple(uuid.uuid4().hex for _ in range(5))
    with ThreadPoolExecutor(max_workers=len(markers)) as pool:
        futures: Final[tuple[Future[httpx.Response], ...]] = tuple(
            pool.submit(rig.chat, _SHARED, "burst", "hold", marker, max_tokens=max_tokens) for marker in markers
        )
        eventually(
            lambda: (
                sum(future.done() for future in futures)
                + sum(rig.hits(_SHARED_DEPLOYMENT, marker) for marker in markers)
            ),
            lambda settled: settled >= len(markers),
            seconds=30,
        )
        rig.release.set()
        responses: Final = tuple(future.result(timeout=60) for future in futures)
    assert sorted(response.status_code for response in responses) == [200, 200, 429, 429, 429], tuple(
        response.text for response in responses
    )
    admitted: Final = tuple(
        marker for marker, response in zip(markers, responses, strict=True) if response.status_code == 200
    )
    assert tuple(rig.hits(_SHARED_DEPLOYMENT, marker) for marker in markers) == tuple(
        1 if marker in admitted else 0 for marker in markers
    )


@pytest.mark.timeout(4 * graceful_stop_seconds() + 240)
def test_a_replacement_worker_rebuilds_the_share_rules(rig: _Rig) -> None:
    victim: Final = min(rig.owned.worker_pids())
    replacement: Final = rig.owned.kill_worker(victim)
    assert replacement != victim
    answers: Final = tuple(
        (rig.chat(_SHARED, "worker", "ok", uuid.uuid4().hex), rig.chat(_SHARED, "outsider", "ok", uuid.uuid4().hex))
        for _ in range(6)
    )
    assert tuple((served.status_code, refused.status_code) for served, refused in answers) == ((200, 400),) * 6, tuple(
        (served.text, refused.text) for served, refused in answers
    )
    assert all(
        string_value(JSON_OBJECT.validate_json(served.content)["id"]).startswith(f"chatcmpl-{_SHARED_DEPLOYMENT}-")
        for served, _ in answers
    )


def test_with_attribution_off_the_access_rule_holds_without_a_ceiling(tmp_path: Path) -> None:
    attempts: Final = Attempts()
    teams: Final = _team_ids()
    capacity: Final = _gpt41()
    with gateway_from_environment() as gateway, wire_server(_respond(attempts, threading.Event())) as wire:
        config: Final = _config(tmp_path, wire.url, _model_list(teams), "ptu-shares-off")
        try:
            with owned_proxy_process(gateway, tmp_path, _ATTRIBUTION_OFF, config=config, workers=2) as owned:
                holder: Final = _create_team(owned.gateway, teams["oversize"])
                outsider: Final = _create_team(owned.gateway, teams["outsider"])
                served_marker: Final = uuid.uuid4().hex
                served: Final = owned.gateway.request(
                    "POST",
                    "/v1/chat/completions",
                    {
                        "model": _SHARED,
                        "max_tokens": capacity.input_tpm_for(1),
                        "messages": [{"role": "user", "content": f"ok:{served_marker}"}],
                    },
                    key=holder,
                )
                refused: Final = owned.gateway.request(
                    "POST",
                    "/v1/chat/completions",
                    {"model": _SHARED, "messages": [{"role": "user", "content": f"ok:{uuid.uuid4().hex}"}]},
                    key=outsider,
                )
        finally:
            _delete_teams(gateway, (teams["oversize"], teams["outsider"]))
    assert served.status_code == 200, served.text
    assert served.json()["id"] == _served_id(_SHARED_DEPLOYMENT, served_marker, 1), served.text
    assert refused.status_code == 400, refused.text
    assert _RESERVED in refused.text, refused.text
    assert _ptu_counter(teams["oversize"], _SHARED) is None


def test_a_split_that_does_not_add_up_is_dropped_at_load_with_attribution_off(tmp_path: Path) -> None:
    attempts: Final = Attempts()
    holder: Final = _team_ids()["oversize"]
    marker: Final = uuid.uuid4().hex
    with gateway_from_environment() as gateway, wire_server(_respond(attempts, threading.Event())) as wire:
        split: Final = _azure(_SHARED, _SHARED_DEPLOYMENT, {**_ptu_terms({holder: 1}), "ptu_count": 2})
        config: Final = _config(tmp_path, wire.url, [split], "ptu-split-mismatch")
        try:
            with owned_proxy_process(gateway, tmp_path, _ATTRIBUTION_OFF, config=config, workers=2) as owned:
                key: Final = _create_team(owned.gateway, holder)
                refused: Final = owned.gateway.request(
                    "POST",
                    "/v1/chat/completions",
                    {"model": _SHARED, "messages": [{"role": "user", "content": f"ok:{marker}"}]},
                    key=key,
                )
                log: Final = owned.log.read_text()
        finally:
            _delete_teams(gateway, (holder,))
    assert refused.status_code == 400, refused.text
    assert _RESERVED not in refused.text, refused.text
    assert attempts.count(f"{_SHARED_DEPLOYMENT}:{marker}") == 0
    assert (
        f"PTU configuration on model '{_SHARED}' is invalid: ptu_shares must add up to ptu_count (1 of 2 allocated)"
        in log
    ), log[-4000:]