diff --git a/litellm/anthropic_beta_headers_manager.py b/litellm/anthropic_beta_headers_manager.py index 8b7fa1c9347..96c38b69104 100644 --- a/litellm/anthropic_beta_headers_manager.py +++ b/litellm/anthropic_beta_headers_manager.py @@ -22,17 +22,28 @@ Environment Variables: - LITELLM_ANTHROPIC_BETA_HEADERS_URL: Custom URL for remote config (optional) """ +import asyncio import json import os +import threading +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from functools import cache from importlib.resources import files -from typing import Final +from typing import Final, TypeAlias import httpx +from pydantic import TypeAdapter from litellm._logging import verbose_logger # Cache for the loaded configuration _BETA_HEADERS_CONFIG: dict | None = None +_BETA_HEADERS_CONFIG_LOCK: Final = threading.Lock() +_BETA_HEADERS_CONFIG_LOADER: Final = ThreadPoolExecutor(max_workers=1, thread_name_prefix="litellm-beta-headers-config") + +BetaHeadersConfig: TypeAlias = Mapping[str, str | Mapping[str, str | None]] +_BETA_HEADERS_CONFIG_ADAPTER: Final = TypeAdapter(BetaHeadersConfig) class GetAnthropicBetaHeadersConfig: @@ -174,13 +185,27 @@ def _load_beta_headers_config() -> dict: if _BETA_HEADERS_CONFIG is not None: return _BETA_HEADERS_CONFIG - # Get the URL from environment or use default - from litellm import anthropic_beta_headers_url + with _BETA_HEADERS_CONFIG_LOCK: + if _BETA_HEADERS_CONFIG is not None: + return _BETA_HEADERS_CONFIG - _BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url) - verbose_logger.debug("Loaded and cached beta headers config") + # Get the URL from environment or use default + from litellm import anthropic_beta_headers_url - return _BETA_HEADERS_CONFIG + _BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url) + verbose_logger.debug("Loaded and cached beta headers config") + + return _BETA_HEADERS_CONFIG + + +async def aload_beta_headers_config() -> None: + if _BETA_HEADERS_CONFIG is None: + await asyncio.get_running_loop().run_in_executor(_BETA_HEADERS_CONFIG_LOADER, _load_beta_headers_config) + + +@cache +def _bundled_beta_headers_config() -> BetaHeadersConfig: + return _BETA_HEADERS_CONFIG_ADAPTER.validate_python(GetAnthropicBetaHeadersConfig.load_local_beta_headers_config()) def reload_beta_headers_config() -> dict: @@ -288,6 +313,14 @@ def is_beta_header_supported( return beta_header in provider_mapping and provider_mapping[beta_header] is not None +def is_beta_header_supported_by_bundled_config(beta_header: str, provider: str) -> bool: + config: Final = _bundled_beta_headers_config() + aliases: Final = config.get("provider_aliases", {}) + canonical_provider: Final = (aliases.get(provider) if isinstance(aliases, Mapping) else None) or provider + provider_mapping: Final = config.get(canonical_provider, {}) + return isinstance(provider_mapping, Mapping) and provider_mapping.get(beta_header) is not None + + def get_provider_beta_header( anthropic_beta_header: str, provider: str, diff --git a/litellm/llms/anthropic/pass_through/messages/transformation.py b/litellm/llms/anthropic/pass_through/messages/transformation.py index 640736d6463..ad45485fa56 100644 --- a/litellm/llms/anthropic/pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/pass_through/messages/transformation.py @@ -4,6 +4,7 @@ from typing import ClassVar, Final import httpx from litellm._logging import verbose_logger +from litellm.anthropic_beta_headers_manager import is_beta_header_supported_by_bundled_config from litellm.exceptions import AuthenticationError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.anthropic_messages.transformation import ( @@ -708,6 +709,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params: dict, custom_llm_provider: str = "anthropic", messages: Sequence[object] = (), + beta_map_provider: str | None = None, ) -> dict: """ Auto-inject anthropic-beta headers based on features used. @@ -724,8 +726,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): optional_params: Optional parameters including tools, context_management, output_format, speed custom_llm_provider: Provider name for looking up correct tool search header messages: Request messages, scanned for per-message output_config + beta_map_provider: Provider whose beta map column gates safeguards, defaults to custom_llm_provider """ - beta_values: Final[set] = set() + beta_values: Final[set[str]] = set() existing_beta: Final = tuple( piece.strip() @@ -775,6 +778,12 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): if optional_params.get("speed") == "fast": beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.FAST_MODE_2026_02_01.value) + if optional_params.get("safeguards") is not None and is_beta_header_supported_by_bundled_config( + ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value, + beta_map_provider or custom_llm_provider, + ): + beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.DANGEROUS_TOOL_USE_2026_09_03.value) + if messages_carry_output_config(messages): beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index fffcdd67150..4ff4c43428f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -4,7 +4,7 @@ import json import ssl from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence from contextlib import asynccontextmanager -from functools import lru_cache, partial +from functools import lru_cache, partial, reduce from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -35,7 +35,7 @@ import litellm.litellm_core_utils import litellm.types import litellm.types.utils from litellm._logging import redact_string, verbose_logger -from litellm.anthropic_beta_headers_manager import update_headers_with_filtered_beta +from litellm.anthropic_beta_headers_manager import aload_beta_headers_config, update_headers_with_filtered_beta from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs @@ -49,6 +49,7 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( synthesize_subtitle_document, ) from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params +from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields from litellm.litellm_core_utils.realtime_errors import ( @@ -2304,6 +2305,16 @@ class BaseLLMHTTPHandler: merged_headers.update(extra_headers_from_kwargs) if provider_specific_headers: merged_headers.update(provider_specific_headers) + request_params: Final = reduce( + delete_nested_value, + litellm_params.get("additional_drop_params") or (), + anthropic_messages_optional_request_params, + ) + if ( + request_params.get("safeguards") is not None + and anthropic_messages_provider_config.should_filter_anthropic_beta_headers() + ): + await aload_beta_headers_config() ( headers, api_base, @@ -2311,7 +2322,7 @@ class BaseLLMHTTPHandler: headers=merged_headers or {}, model=model, messages=messages, - optional_params=anthropic_messages_optional_request_params, + optional_params=request_params, litellm_params=dict(litellm_params), api_key=api_key, api_base=api_base, @@ -2342,20 +2353,11 @@ class BaseLLMHTTPHandler: custom_llm_provider=custom_llm_provider, ) - additional_drop_params: Final[list[str]] = litellm_params.get("additional_drop_params") or [] - if additional_drop_params: - from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value - - for path in additional_drop_params: - anthropic_messages_optional_request_params = delete_nested_value( - anthropic_messages_optional_request_params, path - ) - # Prepare request body request_body: Final = anthropic_messages_provider_config.transform_anthropic_messages_request( model=model, messages=messages, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=request_params, litellm_params=litellm_params, headers=headers, ) @@ -2462,7 +2464,7 @@ class BaseLLMHTTPHandler: held_back_tool_names: Final = self._server_fulfilled_tools_in_request( logging_obj=logging_obj, - tools=anthropic_messages_optional_request_params.get("tools"), + tools=request_params.get("tools"), ) initial_response = AgenticAnthropicStreamingIterator( completion_stream=completion_stream, @@ -2470,7 +2472,7 @@ class BaseLLMHTTPHandler: model=model, messages=messages, anthropic_messages_provider_config=anthropic_messages_provider_config, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=request_params, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, kwargs=kwargs_for_agentic, @@ -2493,7 +2495,7 @@ class BaseLLMHTTPHandler: model=model, messages=messages, anthropic_messages_provider_config=anthropic_messages_provider_config, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=request_params, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, kwargs=kwargs_for_agentic, diff --git a/litellm/llms/openai_like/messages/transformation.py b/litellm/llms/openai_like/messages/transformation.py index 2e9a300e2fd..94f9792da35 100644 --- a/litellm/llms/openai_like/messages/transformation.py +++ b/litellm/llms/openai_like/messages/transformation.py @@ -8,6 +8,7 @@ from litellm.llms.anthropic.pass_through.messages.transformation import ( from litellm.llms.openai_like.json_loader import SimpleProviderConfig from litellm.secret_managers.main import get_secret_str from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import LlmProviders DEFAULT_ANTHROPIC_API_VERSION: Final = "2023-06-01" @@ -57,6 +58,7 @@ class OpenAILikeAnthropicMessagesConfig(AnthropicMessagesConfig): headers=normalized, optional_params=optional_params, messages=messages, + beta_map_provider=LlmProviders.OPENAI_LIKE.value, ) return merged, api_base diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_chaos.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_chaos.py new file mode 100644 index 00000000000..e2f23b7558f --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_chaos.py @@ -0,0 +1,348 @@ +import asyncio +import json +import re +import threading +import uuid +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.messages_endpoint.providers.anthropic.test_anthropic_safeguards_beta_wire import ( + ANTHROPIC_KEY, + ANTHROPIC_MODEL, + BETA, + OMITTED, + anthropic_peer, + marker_of, + messages_body, + stream_frames, +) +from pydantic import JsonValue, TypeAdapter + +_CONFIG_MODEL: Final = "anthropic-safeguards-chaos" +_FOUNDRY_CONFIG_MODEL: Final = "azure-ai-safeguards-chaos" +_YAML_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_STARTED_WORKER: Final = re.compile(r"Started server process \[(\d+)\]") +_MODEL_LIST: Final = json.dumps( + {"object": "list", "data": [{"id": ANTHROPIC_MODEL, "object": "model", "owned_by": "anthropic"}]} +).encode() + + +@dataclass(frozen=True, slots=True) +class _Call: + stream: bool + marker: str + safeguards: bool = True + + +@dataclass(frozen=True, slots=True) +class _Served: + call: _Call + status: int + text: str + + +def _calls(count: int, stream: bool, safeguards: bool = True) -> tuple[_Call, ...]: + return tuple(_Call(stream=stream, marker=uuid.uuid4().hex, safeguards=safeguards) for _ in range(count)) + + +def _streamed_reply(marker: str, abort_after: int | None = None, pause: float = 0) -> Reply: + return Reply( + content_type="text/event-stream", + chunks=stream_frames(marker), + abort_after=abort_after, + pause_between_chunks=pause, + ) + + +async def _send(client: httpx.AsyncClient, key: str, model: str, call: _Call) -> _Served: + async with client.stream( + "POST", + "/v1/messages", + json=( + messages_body(model, call.marker, stream=call.stream) + if call.safeguards + else messages_body(model, call.marker, stream=call.stream, safeguards=OMITTED) + ), + headers={"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"}, + ) as response: + raw: Final = await response.aread() + return _Served(call=call, status=response.status_code, text=raw.decode()) + + +async def _burst( + base_url: str, key: str, model: str, calls: tuple[_Call, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_Served, ...]: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_send(client, key, model, call) for call in calls), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Served)) + + +def _assert_answered_with_its_own_marker(served: _Served) -> None: + assert served.status == 200, served.text + assert set(_MARKER.findall(served.text)) == {served.call.marker}, served.text + assert f"msg_{served.call.marker}" in served.text, served.text + if served.call.stream: + assert "event: message_stop" in served.text, served.text + + +def _assert_every_request_carried_the_beta(received: tuple[Request, ...], markers: frozenset[str]) -> None: + assert sorted(marker_of(request) for request in received) == sorted(markers) + missing: Final = [ + frozenset(request.headers) for request in received if request.headers.get("anthropic-beta") != BETA + ] + assert missing == [], missing + + +def _spend_success_ids(model: str, expected: frozenset[str]) -> None: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE model_group=%s AND status=%s', (model, "success") + ), + lambda found: expected <= {str(row["request_id"]) for row in found}, + seconds=70, + ) + identities: Final = [row["request_id"] for row in rows] + assert len(set(identities)) == len(identities), identities + + +async def test_burst_with_aborted_streams_keeps_the_beta_on_every_forwarded_request(gateway: Gateway) -> None: + calls: Final = tuple(_Call(stream=index % 2 == 0, marker=uuid.uuid4().hex) for index in range(30)) + aborted: Final = frozenset(call.marker for index, call in enumerate(calls) if call.stream and index % 3 == 0) + + def respond(request: Request) -> Reply: + marker: Final = marker_of(request) + if marker in aborted: + return _streamed_reply(marker, abort_after=0) + return anthropic_peer(request) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 30 + for item in served: + if item.call.marker in aborted: + assert item.status == 500, item.text + assert '"type":"api_error"' in item.text and "marker-" not in item.text, item.text + else: + _assert_answered_with_its_own_marker(item) + recovery: Final = _Call(stream=True, marker=uuid.uuid4().hex) + (recovered,) = await _burst(str(gateway.client.base_url), gateway.key, model, (recovery,)) + _assert_answered_with_its_own_marker(recovered) + _assert_every_request_carried_the_beta(wire.drain(), frozenset(call.marker for call in (*calls, recovery))) + answered: Final = frozenset(f"msg_{call.marker}" for call in (*calls, recovery) if call.marker not in aborted) + _spend_success_ids(model, answered) + + +async def test_slow_streams_complete_with_the_beta_on_every_request(gateway: Gateway) -> None: + calls: Final = _calls(10, stream=True) + with ( + wire_server(lambda request: _streamed_reply(marker_of(request), pause=0.3)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + served: Final = await _burst(str(gateway.client.base_url), gateway.key, model, calls) + assert len(served) == 10 + for item in served: + _assert_answered_with_its_own_marker(item) + _assert_every_request_carried_the_beta(wire.drain(), frozenset(call.marker for call in calls)) + _spend_success_ids(model, frozenset(f"msg_{call.marker}" for call in calls)) + + +def _anthropic_deployment(wire: Wire) -> dict[str, JsonValue]: + return { + "model_name": _CONFIG_MODEL, + "litellm_params": {"model": f"anthropic/{ANTHROPIC_MODEL}", "api_base": wire.url, "api_key": ANTHROPIC_KEY}, + } + + +def _foundry_deployment(wire: Wire) -> dict[str, JsonValue]: + return { + "model_name": _FOUNDRY_CONFIG_MODEL, + "litellm_params": { + "model": "azure_ai/claude-sonnet-4-6", + "api_base": wire.url, + "api_key": "synthetic-foundry-key", + }, + } + + +def _chaos_config(tmp_path: Path, *deployments: dict[str, JsonValue]) -> Path: + base: Final = _YAML_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final = {**base, "model_list": list(deployments)} + path: Final = tmp_path / "anthropic-safeguards-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@pytest.mark.timeout(180) +async def test_worker_sigkill_mid_burst_leaves_the_sibling_sending_the_beta(gateway: Gateway, tmp_path: Path) -> None: + calls: Final = _calls(20, stream=False) + release: Final = threading.Event() + held_markers: Final[SimpleQueue[str]] = SimpleQueue() + + def held(request: Request) -> Reply: + if (request.method, request.target) == ("GET", "/v1/models"): + return Reply(body=_MODEL_LIST) + held_markers.put(marker_of(request)) + assert release.wait(timeout=60), "The burst was never released" + return anthropic_peer(request) + + with wire_server(held) as wire: + path: Final = _chaos_config(tmp_path, _anthropic_deployment(wire)) + with owned_proxy_process(gateway, tmp_path, {}, config=path, workers=2) as owned: + candidate: Final = owned.gateway + workers: Final = eventually( + lambda: tuple(int(found.group(1)) for found in _STARTED_WORKER.finditer(owned.log.read_text())), + lambda pids: len(pids) == 2, + seconds=30, + ) + burst: Final = asyncio.create_task( + _burst( + str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, calls, tolerate_transport_errors=True + ) + ) + await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == 20, 60) + held_by: Final = MappingProxyType({pid: _open_upstream_connections(pid, wire.url) for pid in workers}) + assert sum(held_by.values()) == 20, held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert held_by[survivor_pid] >= 10, held_by + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _assert_answered_with_its_own_marker(item) + follow_up: Final = _Call(stream=False, marker=uuid.uuid4().hex) + (answered,) = await _burst(str(candidate.client.base_url), candidate.key, _CONFIG_MODEL, (follow_up,)) + _assert_answered_with_its_own_marker(answered) + received: Final = wire.drain() + posts: Final = tuple(request for request in received if request.method == "POST") + probes: Final[list[tuple[str, str]]] = [ + (request.method, request.target) for request in received if request.method != "POST" + ] + assert set(probes) <= {("GET", "/v1/models")}, probes + _assert_every_request_carried_the_beta(posts, frozenset(call.marker for call in (*calls, follow_up))) + + +def _beta_map_gets(received: tuple[Request, ...]) -> int: + return sum(1 for request in received if (request.method, request.target) == ("GET", "/beta.json")) + + +async def _chat(base_url: str, key: str, marker: str) -> httpx.Response: + async with httpx.AsyncClient(base_url=base_url, timeout=60, trust_env=False) as client: + return await client.post( + "/v1/chat/completions", + json={ + "model": _CONFIG_MODEL, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"Question marker-{marker}"}], + }, + headers={"Authorization": f"Bearer {key}"}, + ) + + +async def _answered_while_the_beta_map_is_held( + base_url: str, key: str, guarded: tuple[_Call, ...], chat_markers: tuple[str, ...] +) -> None: + served: Final = await _burst(base_url, key, _CONFIG_MODEL, guarded) + chats: Final = await asyncio.gather(*(_chat(base_url, key, marker) for marker in chat_markers)) + for item in served: + _assert_answered_with_its_own_marker(item) + for marker, chat in zip(chat_markers, chats, strict=True): + assert chat.status_code == 200, chat.text + assert f"marker-{marker}" in chat.text, chat.text + + +@pytest.mark.timeout(240) +async def test_cold_workers_fetch_the_beta_map_only_for_filtered_safeguards_requests_off_the_shared_threads( + gateway: Gateway, tmp_path: Path +) -> None: + beta_map: Final = Path("litellm/anthropic_beta_headers_config.json").read_bytes() + release: Final = threading.Event() + + def held_beta_map(_: Request) -> Reply: + assert release.wait(timeout=180), "The beta map was never released" + return Reply(body=beta_map) + + def upstream(request: Request) -> Reply: + if (request.method, request.target) == ("GET", "/v1/models"): + return Reply(body=_MODEL_LIST) + return anthropic_peer(request) + + plain: Final = _calls(8, stream=False, safeguards=False) + plain_foundry: Final = _calls(8, stream=False, safeguards=False) + direct: Final = _calls(8, stream=False) + foundry: Final = _calls(80, stream=False) + direct_while_held: Final = _calls(8, stream=False) + chat_markers: Final = tuple(uuid.uuid4().hex for _ in range(4)) + with wire_server(held_beta_map) as beta_map_wire, wire_server(upstream) as wire: + path: Final = _chaos_config(tmp_path, _anthropic_deployment(wire), _foundry_deployment(wire)) + with owned_proxy_process( + gateway, + tmp_path, + {"LITELLM_ANTHROPIC_BETA_HEADERS_URL": f"{beta_map_wire.url}/beta.json"}, + remove_environment=("LITELLM_LOCAL_ANTHROPIC_BETA_HEADERS",), + config=path, + workers=2, + ) as owned: + base_url: Final = str(owned.gateway.client.base_url) + key: Final = owned.gateway.key + cold: Final = ( + *await _burst(base_url, key, _CONFIG_MODEL, plain), + *await _burst(base_url, key, _FOUNDRY_CONFIG_MODEL, plain_foundry), + *await _burst(base_url, key, _CONFIG_MODEL, direct), + ) + for served_cold in cold: + _assert_answered_with_its_own_marker(served_cold) + assert beta_map_wire.received.qsize() == 0, "a request that never filters betas fetched the beta map" + + held: Final = asyncio.create_task(_burst(base_url, key, _FOUNDRY_CONFIG_MODEL, foundry)) + try: + await asyncio.to_thread(eventually, beta_map_wire.received.qsize, lambda size: size >= 1, 60) + await asyncio.wait_for( + _answered_while_the_beta_map_is_held(base_url, key, direct_while_held, chat_markers), timeout=45 + ) + assert not held.done(), "Foundry safeguards requests were answered before the beta map loaded" + finally: + release.set() + for served_foundry in await held: + _assert_answered_with_its_own_marker(served_foundry) + + received: Final = tuple(request for request in wire.drain() if request.method == "POST") + beta_by_marker: Final = MappingProxyType( + {marker_of(request): request.headers.get("anthropic-beta") for request in received} + ) + assert {beta_by_marker[call.marker] for call in (*direct, *direct_while_held, *foundry)} == {BETA}, ( + beta_by_marker + ) + assert {beta_by_marker[call.marker] for call in (*plain, *plain_foundry)} == {None}, beta_by_marker + fetches: Final = _beta_map_gets(beta_map_wire.drain()) + assert 1 <= fetches <= 2, fetches diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py new file mode 100644 index 00000000000..0c41efffe73 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_safeguards_beta_wire.py @@ -0,0 +1,809 @@ +import itertools +import json +import re +import time +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.process import owned_proxy +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue, TypeAdapter + +BETA: Final = "dangerous-tool-use-2026-09-03" +SAFEGUARDS: Final[list[JsonValue]] = [ + {"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}} +] +ANTHROPIC_MODEL: Final = "claude-sonnet-4-5-20250929" +ANTHROPIC_KEY: Final = "synthetic-anthropic-key" +_CLAUDE_CODE_BETAS: Final = ( + "claude-code-20250219", + "interleaved-thinking-2025-05-14", + "fine-grained-tool-streaming-2025-05-14", + "context-management-2025-06-27", +) +_CLIENT_BETAS: Final = ",".join(_CLAUDE_CODE_BETAS) +_CLIENT_BETAS_WITH_DANGEROUS: Final = ",".join((*_CLAUDE_CODE_BETAS, BETA)) +_MERGED_BETAS: Final = ",".join(sorted((*_CLAUDE_CODE_BETAS, BETA))) +_FOUNDRY_MODEL: Final = "azure_ai/claude-sonnet-4-6" +_FOUNDRY_KEY: Final = "synthetic-foundry-key" +_FOUNDRY_TARGET: Final = "/anthropic/v1/messages" +_FOUNDRY_ARRAY_ERROR: Final = "safeguards: Input should be a valid array" +_FOUNDRY_EXTRA_INPUT_ERROR: Final = "safeguards: Extra inputs are not permitted" +_ANTHROPIC_VERSION: Final = MappingProxyType({"anthropic-version": "2023-06-01"}) +_MESSAGES_KEYS: Final = frozenset({"max_tokens", "messages", "model", "stream"}) +_RESPONSES_KEYS: Final = frozenset({"include", "input", "max_output_tokens", "model"}) +_SDK_EXTRA: Final[dict[str, JsonValue]] = {"safeguards": SAFEGUARDS, "cache": {"no-cache": True}, "num_retries": 0} +_MARKER: Final = re.compile(r"marker-([0-9a-f]{32})") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +@dataclass(frozen=True, slots=True) +class _Omitted: + pass + + +OMITTED: Final = _Omitted() + + +@dataclass(frozen=True, slots=True) +class _Observed: + method: str + target: str + marker: str + headers: Mapping[str, str] + body: Mapping[str, JsonValue] + + @property + def beta(self) -> str | None: + return self.headers.get("anthropic-beta") + + @property + def keys(self) -> frozenset[str]: + return frozenset(self.body) + + @property + def header_names(self) -> frozenset[str]: + return frozenset(self.headers) + + +def marker_of(request: Request) -> str: + found: Final = _MARKER.search(request.body.decode()) + assert found is not None, request.body + return found.group(1) + + +def _message(marker: str) -> dict[str, JsonValue]: + return { + "id": f"msg_{marker}", + "type": "message", + "role": "assistant", + "model": ANTHROPIC_MODEL, + "content": [{"type": "text", "text": f"answer marker-{marker}"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 5}, + } + + +def stream_frames(marker: str) -> tuple[bytes, ...]: + opening: Final[dict[str, JsonValue]] = { + **_message(marker), + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 12, "output_tokens": 0}, + } + events: Final[tuple[tuple[str, dict[str, JsonValue]], ...]] = ( + ("message_start", {"message": opening}), + ("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}), + ("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": f"answer marker-{marker}"}}), + ("content_block_stop", {"index": 0}), + ("message_delta", {"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 5}}), + ("message_stop", {}), + ) + return tuple( + f"event: {kind}\ndata: {json.dumps({'type': kind, **payload})}\n\n".encode() for kind, payload in events + ) + + +def anthropic_peer(request: Request) -> Reply: + marker: Final = marker_of(request) + if _JSON_OBJECT.validate_json(request.body).get("stream") is True: + return Reply(content_type="text/event-stream", chunks=stream_frames(marker)) + return Reply(body=json.dumps(_message(marker)).encode()) + + +def _responses_peer(request: Request) -> Reply: + marker: Final = marker_of(request) + return Reply( + body=json.dumps( + { + "id": f"resp_{marker}", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": f"msg_{marker}", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": f"answer marker-{marker}", "annotations": []}], + } + ], + "usage": {"input_tokens": 12, "output_tokens": 5, "total_tokens": 17}, + } + ).encode() + ) + + +def _error_peer(status: int, kind: str, message: str) -> Callable[[Request], Reply]: + body: Final = json.dumps({"type": "error", "error": {"type": kind, "message": message}}).encode() + return lambda request: Reply(status=status, body=body) + + +def messages_body( + model: str, + marker: str, + *, + safeguards: JsonValue | _Omitted = SAFEGUARDS, + stream: bool = False, + no_cache: bool = True, +) -> dict[str, JsonValue]: + base: Final[dict[str, JsonValue]] = { + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": f"Question marker-{marker}"}], + "num_retries": 0, + **({"stream": True} if stream else {}), + **({"cache": {"no-cache": True}} if no_cache else {}), + } + if isinstance(safeguards, _Omitted): + return base + return {**base, "safeguards": safeguards} + + +def _observe(request: Request) -> _Observed: + return _Observed( + request.method, request.target, marker_of(request), request.headers, _JSON_OBJECT.validate_json(request.body) + ) + + +def _only(wire: Wire) -> _Observed: + (request,) = wire.drain() + return _observe(request) + + +def _post(gateway: Gateway, body: Mapping[str, JsonValue], *, beta: str | None = None) -> httpx.Response: + headers: Final = {**_ANTHROPIC_VERSION, **({"anthropic-beta": beta} if beta else {})} + return gateway.request("POST", "/v1/messages", body, headers=headers) + + +def _assert_answered(response: httpx.Response, marker: str) -> None: + assert response.status_code == 200, response.text + assert response.json()["id"] == f"msg_{marker}", response.text + + +def _assert_streamed(response: httpx.Response, marker: str) -> None: + assert response.status_code == 200, response.text + assert f"msg_{marker}" in response.text and f"answer marker-{marker}" in response.text, response.text + assert "event: message_stop" in response.text, response.text + + +def _spend_row(identity: str) -> None: + rows: Final = eventually( + lambda: read_rows('SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE request_id=%s', (identity,)), + lambda found: len(found) == 1, + seconds=70, + ) + assert rows[0]["status"] == "success", rows + + +def _cache_hit_spend_row(identity: str) -> None: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status, cache_hit FROM "LiteLLM_SpendLogs" WHERE request_id LIKE %s', + (f"{identity}\\_cache\\_hit%",), + ), + lambda found: len(found) == 1, + seconds=70, + ) + assert (rows[0]["status"], str(rows[0]["cache_hit"])) == ("success", "True"), rows + + +def _responses_spend_row(model: str) -> None: + rows: Final = eventually( + lambda: read_rows( + 'SELECT request_id, status FROM "LiteLLM_SpendLogs" WHERE model_group = %s AND call_type = %s', + (model, "aresponses"), + ), + lambda found: len(found) == 1, + seconds=70, + ) + assert rows[0]["status"] == "success", rows + + +def _expect_messages( + observed: _Observed, + *, + target: str, + auth: tuple[str, str] | None, + beta: str | None, + safeguards: JsonValue | _Omitted = SAFEGUARDS, + keys: frozenset[str] = _MESSAGES_KEYS, +) -> None: + assert (observed.method, observed.target) == ("POST", target) + if auth is not None: + credential_matches: Final = observed.headers.get(auth[0]) == auth[1] + assert credential_matches, f"{auth[0]} differs from the deployment credential; headers {observed.header_names}" + assert observed.beta == beta, observed.header_names + if isinstance(safeguards, _Omitted): + assert observed.keys == keys + return + assert observed.keys == keys | {"safeguards"} + assert observed.body["safeguards"] == safeguards + + +def _expect_anthropic_version(observed: _Observed) -> None: + assert observed.headers.get("anthropic-version") == "2023-06-01", observed.header_names + + +@dataclass(frozen=True, slots=True) +class _ShapeCase: + sent: JsonValue | _Omitted + forwarded: JsonValue | _Omitted + client_beta: str | None + expected_beta: str | None + stream: bool = False + + +_ANTHROPIC_SHAPES: Final = MappingProxyType( + { + "D1_safeguards_without_client_header": _ShapeCase(SAFEGUARDS, SAFEGUARDS, None, BETA), + "D2_safeguards_with_claude_code_header": _ShapeCase( + SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS_WITH_DANGEROUS, _MERGED_BETAS + ), + "D3_safeguards_with_claude_code_header_lacking_the_beta": _ShapeCase( + SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS, _MERGED_BETAS + ), + "D4_empty_safeguards": _ShapeCase([], [], None, BETA), + "D5_null_safeguards": _ShapeCase(None, OMITTED, None, None), + "D6_missing_safeguards": _ShapeCase(OMITTED, OMITTED, None, None), + } +) +_FOUNDRY_SHAPES: Final = MappingProxyType( + { + "D9_streamed_safeguards_with_claude_code_header": _ShapeCase( + SAFEGUARDS, SAFEGUARDS, _CLIENT_BETAS_WITH_DANGEROUS, BETA, stream=True + ), + "D11_empty_safeguards": _ShapeCase([], [], None, BETA), + "D12_null_safeguards": _ShapeCase(None, OMITTED, None, None), + "D13_missing_safeguards_with_claude_code_header_lacking_the_beta": _ShapeCase( + OMITTED, OMITTED, _CLIENT_BETAS, None + ), + } +) + + +@pytest.mark.parametrize("row", tuple(_ANTHROPIC_SHAPES)) +def test_anthropic_messages_derive_the_dangerous_tool_use_beta_from_safeguards(gateway: Gateway, row: str) -> None: + case: Final = _ANTHROPIC_SHAPES[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + response: Final = _post(gateway, messages_body(model, marker, safeguards=case.sent), beta=case.client_beta) + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages( + observed, + target="/v1/messages", + auth=("x-api-key", ANTHROPIC_KEY), + beta=case.expected_beta, + safeguards=case.forwarded, + ) + _expect_anthropic_version(observed) + _spend_row(f"msg_{marker}") + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _DropCase: + deployment: Callable[[Gateway], Mapping[str, JsonValue]] + target: str + auth: tuple[str, str] | None + keys: frozenset[str] = _MESSAGES_KEYS + + +_VERTEX_PROJECT: Final = "scripted-project" +_VERTEX_LOCATION: Final = "us-east5" +_VERTEX_MODEL: Final = "claude-opus-5-5" + +_DROP_CASES: Final = MappingProxyType( + { + "anthropic": _DropCase( + deployment=lambda _: {"model": f"anthropic/{ANTHROPIC_MODEL}", "api_key": ANTHROPIC_KEY}, + target="/v1/messages", + auth=("x-api-key", ANTHROPIC_KEY), + ), + "azure_ai": _DropCase( + deployment=lambda _: {"model": _FOUNDRY_MODEL, "api_key": _FOUNDRY_KEY}, + target=_FOUNDRY_TARGET, + auth=("x-api-key", _FOUNDRY_KEY), + ), + "bedrock_claude_platform": _DropCase( + deployment=lambda _: { + "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", + "api_key": "synthetic-platform-key", + "workspace_id": "ws-synthetic", + "aws_region_name": "us-east-1", + }, + target="/v1/messages", + auth=("x-api-key", "synthetic-platform-key"), + ), + "vertex_ai": _DropCase( + deployment=lambda gateway: { + "model": f"vertex_ai/{_VERTEX_MODEL}", + "api_key": None, + "vertex_project": _VERTEX_PROJECT, + "vertex_location": _VERTEX_LOCATION, + "vertex_credentials": service_account_json(_VERTEX_PROJECT, gateway.upstream_url.rstrip("/")), + }, + target=( + f"/v1/projects/{_VERTEX_PROJECT}/locations/{_VERTEX_LOCATION}" + f"/publishers/anthropic/models/{_VERTEX_MODEL}:rawPredict" + ), + auth=None, + keys=frozenset({"anthropic_version", "max_tokens", "messages", "stream"}), + ), + } +) + + +@pytest.mark.parametrize("row", tuple(_DROP_CASES)) +def test_a_deployment_dropping_safeguards_sends_neither_the_field_nor_the_beta(gateway: Gateway, row: str) -> None: + case: Final = _DROP_CASES[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + api_base=wire.url, additional_drop_params=["safeguards"], **case.deployment(gateway) + ) + response: Final = _post(gateway, messages_body(model, marker)) + _assert_answered(response, marker) + _expect_messages(_only(wire), target=case.target, auth=case.auth, beta=None, safeguards=OMITTED, keys=case.keys) + _spend_row(f"msg_{marker}") + + +def test_a_deployment_dropping_context_management_sends_neither_the_field_nor_its_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{ANTHROPIC_MODEL}", + api_base=wire.url, + api_key=ANTHROPIC_KEY, + additional_drop_params=["context_management"], + ) + body: Final = { + **messages_body(model, marker, safeguards=OMITTED), + "context_management": {"edits": [{"type": "clear_tool_uses_20250919"}]}, + } + response: Final = _post(gateway, body) + _assert_answered(response, marker) + _expect_messages( + _only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=None, safeguards=OMITTED + ) + _spend_row(f"msg_{marker}") + + +@pytest.mark.parametrize("row", tuple(_FOUNDRY_SHAPES)) +def test_azure_ai_messages_derive_the_beta_instead_of_forwarding_the_client_header(gateway: Gateway, row: str) -> None: + case: Final = _FOUNDRY_SHAPES[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + response: Final = _post( + gateway, messages_body(model, marker, safeguards=case.sent, stream=case.stream), beta=case.client_beta + ) + if case.stream: + _assert_streamed(response, marker) + else: + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages( + observed, + target=_FOUNDRY_TARGET, + auth=("x-api-key", _FOUNDRY_KEY), + beta=case.expected_beta, + safeguards=case.forwarded, + ) + _expect_anthropic_version(observed) + _spend_row(f"msg_{marker}") + + +def test_anthropic_sdk_sync_safeguards_derive_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + reply: Final = client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": f"Question marker-{marker}"}], + extra_body=_SDK_EXTRA, + ) + assert reply.id == f"msg_{marker}", reply + _expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + + +async def test_anthropic_sdk_async_stream_safeguards_derive_the_beta_and_log_spend(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = anthropic.AsyncAnthropic( + base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0 + ) + async with client.messages.with_streaming_response.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": f"Question marker-{marker}"}], + stream=True, + extra_body=_SDK_EXTRA, + ) as streamed: + pieces: Final = [ + event.delta.text + async for event in await streamed.parse() + if event.type == "content_block_delta" and event.delta.type == "text_delta" + ] + assert "".join(pieces) == f"answer marker-{marker}", pieces + _expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + + +def test_azure_ai_sdk_sync_safeguards_derive_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + client: Final = anthropic.Anthropic(base_url=str(gateway.client.base_url), api_key=gateway.key, max_retries=0) + reply: Final = client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": f"Question marker-{marker}"}], + extra_body=_SDK_EXTRA, + ) + assert reply.id == f"msg_{marker}", reply + _expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + + +@dataclass(frozen=True, slots=True, kw_only=True) +class _ProviderCase: + parameters: Mapping[str, JsonValue] + auth: tuple[str, str] | None + expected_beta: str | None + target: str = "/v1/messages" + api_base_suffix: str = "" + model_info: Mapping[str, JsonValue] | None = None + client_beta: str | None = None + keys: frozenset[str] = _MESSAGES_KEYS + extra_headers: Mapping[str, str] = MappingProxyType({}) + body_beta: list[str] | None = None + + +_PROVIDERS: Final = MappingProxyType( + { + "D14_deepseek_with_claude_code_header": _ProviderCase( + parameters={"model": "deepseek/deepseek-chat", "api_key": "synthetic-deepseek-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + target="/anthropic/v1/messages", + auth=("x-api-key", "synthetic-deepseek-key"), + expected_beta=None, + ), + "D15_tencent_with_claude_code_header": _ProviderCase( + parameters={"model": "tencent/deepseek-v4-flash", "api_key": "synthetic-tencent-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("x-api-key", "synthetic-tencent-key"), + expected_beta=None, + ), + "D16_minimax_with_claude_code_header": _ProviderCase( + parameters={"model": "minimax/MiniMax-M2", "api_key": "synthetic-minimax-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("x-api-key", "synthetic-minimax-key"), + expected_beta=None, + ), + "D17_openai_compatible_native_passthrough": _ProviderCase( + parameters={"model": "openai/claude-sonnet-4-5", "api_key": "synthetic-openai-like-key"}, + api_base_suffix="/v1", + model_info={"supported_endpoints": ["/v1/messages"]}, + auth=("authorization", "Bearer synthetic-openai-like-key"), + expected_beta=None, + ), + "D18_sail_with_claude_code_header": _ProviderCase( + parameters={"model": "sail/sail-claude", "api_key": "synthetic-sail-key"}, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("authorization", "Bearer synthetic-sail-key"), + expected_beta=None, + ), + "D19_edenai": _ProviderCase( + parameters={"model": "edenai/claude-sonnet-4-5", "api_key": "synthetic-edenai-key"}, + auth=("authorization", "Bearer synthetic-edenai-key"), + expected_beta=None, + ), + "D21_bedrock_claude_platform": _ProviderCase( + parameters={ + "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", + "api_key": "synthetic-platform-key", + "workspace_id": "ws-synthetic", + "aws_region_name": "us-east-1", + }, + auth=("x-api-key", "synthetic-platform-key"), + expected_beta=BETA, + extra_headers={"anthropic-workspace-id": "ws-synthetic"}, + ), + "D22_bedrock_claude_platform_with_claude_code_header": _ProviderCase( + parameters={ + "model": f"bedrock/claude_platform/{ANTHROPIC_MODEL}", + "api_key": "synthetic-platform-key", + "workspace_id": "ws-synthetic", + "aws_region_name": "us-east-1", + }, + client_beta=_CLIENT_BETAS_WITH_DANGEROUS, + auth=("x-api-key", "synthetic-platform-key"), + expected_beta=_MERGED_BETAS, + extra_headers={"anthropic-workspace-id": "ws-synthetic"}, + ), + "D23_bedrock_invoke": _ProviderCase( + parameters={ + "model": "bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + "api_key": "synthetic-bedrock-bearer", + "aws_region_name": "us-east-1", + }, + target="/model/anthropic.claude-sonnet-4-5-20250929-v1:0/invoke", + auth=("authorization", "Bearer synthetic-bedrock-bearer"), + expected_beta=None, + keys=frozenset({"anthropic_beta", "anthropic_version", "max_tokens", "messages"}), + body_beta=[BETA], + ), + "D24_bedrock_mantle": _ProviderCase( + parameters={ + "model": "bedrock_mantle/anthropic.claude-sonnet-5-v1:0", + "api_key": "synthetic-mantle-bearer", + "aws_region_name": "us-east-1", + }, + target="/anthropic/v1/messages", + auth=("authorization", "Bearer synthetic-mantle-bearer"), + expected_beta=BETA, + keys=frozenset({"max_tokens", "messages", "model"}), + ), + } +) + + +@pytest.mark.parametrize("row", tuple(_PROVIDERS)) +def test_anthropic_compatible_providers_receive_the_beta_only_where_the_proxy_derives_it( + gateway: Gateway, row: str +) -> None: + case: Final = _PROVIDERS[row] + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model_info=case.model_info, api_base=wire.url + case.api_base_suffix, **case.parameters + ) + response: Final = _post(gateway, messages_body(model, marker), beta=case.client_beta) + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages(observed, target=case.target, auth=case.auth, beta=case.expected_beta, keys=case.keys) + for name, value in case.extra_headers.items(): + assert observed.headers.get(name) == value, observed.header_names + assert observed.body.get("anthropic_beta") == case.body_beta, observed.body + _spend_row(f"msg_{marker}") + + +def test_github_copilot_messages_forward_safeguards_without_deriving_the_beta( + gateway: Gateway, tmp_path: Path +) -> None: + marker: Final = uuid.uuid4().hex + token_dir: Final = tmp_path / "copilot" + token_dir.mkdir() + with wire_server(anthropic_peer) as wire: + (token_dir / "api-key.json").write_text( + json.dumps({"token": "synthetic-copilot-token", "expires_at": 4102444800, "endpoints": {"api": wire.url}}) + ) + with ( + owned_proxy(gateway, tmp_path, {"GITHUB_COPILOT_TOKEN_DIR": str(token_dir)}) as owned, + owned.scenario() as scenario, + ): + model: Final = scenario.model(model="github_copilot/claude-sonnet-4-5", api_base=wire.url) + response: Final = _post(owned, messages_body(model, marker), beta=_CLIENT_BETAS_WITH_DANGEROUS) + _assert_answered(response, marker) + observed: Final = _only(wire) + _expect_messages( + observed, target="/v1/messages", auth=("authorization", "Bearer synthetic-copilot-token"), beta=None + ) + assert observed.headers.get("openai-intent") == "messages-proxy", observed.header_names + _spend_row(f"msg_{marker}") + + +def test_non_native_deployments_bridge_to_the_responses_api_without_safeguards_or_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=wire.url + "/v1", api_key="synthetic-openai-key" + ) + response: Final = _post(gateway, messages_body(model, marker), beta=_CLIENT_BETAS_WITH_DANGEROUS) + assert response.status_code == 200 and f"answer marker-{marker}" in response.text, response.text + observed: Final = _only(wire) + assert (observed.method, observed.target) == ("POST", "/v1/responses") + assert observed.headers.get("authorization") == "Bearer synthetic-openai-key", observed.header_names + assert observed.beta is None, observed.header_names + assert observed.keys == _RESPONSES_KEYS, observed.keys + _spend_row(str(_JSON_OBJECT.validate_json(response.content)["id"])) + + +def test_chat_completions_forward_safeguards_verbatim_without_deriving_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = openai.OpenAI( + base_url=str(gateway.client.base_url).rstrip("/") + "/v1", api_key=gateway.key, max_retries=0 + ) + completion: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": f"Question marker-{marker}"}], extra_body=_SDK_EXTRA + ) + assert completion.choices[0].message.content == f"answer marker-{marker}", completion + _expect_messages( + _only(wire), + target="/v1/messages", + auth=("x-api-key", ANTHROPIC_KEY), + beta=None, + keys=frozenset({"max_tokens", "messages", "model"}), + ) + _spend_row(completion.id) + + +async def test_responses_stream_forwards_safeguards_verbatim_without_deriving_the_beta(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{ANTHROPIC_MODEL}", api_base=wire.url, api_key=ANTHROPIC_KEY) + client: Final = openai.AsyncOpenAI( + base_url=str(gateway.client.base_url).rstrip("/") + "/v1", api_key=gateway.key, max_retries=0 + ) + async with client.responses.with_streaming_response.create( + model=model, input=f"Question marker-{marker}", stream=True, extra_body=_SDK_EXTRA + ) as streamed: + events: Final = [event async for event in await streamed.parse()] + deltas: Final = [event.delta for event in events if event.type == "response.output_text.delta"] + assert "".join(deltas) == f"answer marker-{marker}", events + assert any(event.type == "response.completed" for event in events), events + _expect_messages(_only(wire), target="/v1/messages", auth=("x-api-key", ANTHROPIC_KEY), beta=None) + _responses_spend_row(model) + + +_FOUNDRY_REJECTED_SHAPES: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"S1_int": 1, "S2_empty_string": "", "S3_five_kb_string": "x" * 5000} +) + + +@pytest.mark.parametrize("row", tuple(_FOUNDRY_REJECTED_SHAPES)) +def test_foundry_rejecting_a_non_array_safeguards_value_reaches_the_caller_after_one_attempt( + gateway: Gateway, row: str +) -> None: + value: Final = _FOUNDRY_REJECTED_SHAPES[row] + marker: Final = uuid.uuid4().hex + with ( + wire_server(_error_peer(400, "invalid_request_error", _FOUNDRY_ARRAY_ERROR)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + response: Final = _post(gateway, messages_body(model, marker, safeguards=value)) + assert response.status_code == 400 and _FOUNDRY_ARRAY_ERROR in response.text, response.text + _expect_messages( + _only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA, safeguards=value + ) + + +def test_duplicate_safeguards_keys_resolve_to_the_last_value(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + raw: Final = ( + f'{{"model": "{model}", "max_tokens": 16, "num_retries": 0, "cache": {{"no-cache": true}}, ' + f'"messages": [{{"role": "user", "content": "Question marker-{marker}"}}], ' + f'"safeguards": {json.dumps(SAFEGUARDS)}, "safeguards": null}}' + ) + response: Final = gateway.client.post( + "/v1/messages", + content=raw.encode(), + headers={ + "Authorization": f"Bearer {gateway.key}", + "content-type": "application/json", + **_ANTHROPIC_VERSION, + }, + ) + _assert_answered(response, marker) + _expect_messages( + _only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=None, safeguards=OMITTED + ) + _spend_row(f"msg_{marker}") + + +def test_the_same_safeguards_sent_twice_are_two_upstream_calls_with_two_spend_rows(gateway: Gateway) -> None: + markers: Final = (uuid.uuid4().hex, uuid.uuid4().hex) + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + responses: Final = tuple(_post(gateway, messages_body(model, marker)) for marker in markers) + for marker, response in zip(markers, responses, strict=True): + _assert_answered(response, marker) + received: Final = tuple(_observe(request) for request in wire.drain()) + assert len(received) == 2, received + for observed in received: + _expect_messages(observed, target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + assert sorted(item.marker for item in received) == sorted(markers) + for marker in markers: + _spend_row(f"msg_{marker}") + + +_UPSTREAM_ERRORS: Final = MappingProxyType( + { + "S6_unauthorized_deployment": (401, "authentication_error", "invalid x-api-key"), + "S7_foundry_extra_inputs_rejection": (400, "invalid_request_error", _FOUNDRY_EXTRA_INPUT_ERROR), + } +) + + +@pytest.mark.parametrize("row", tuple(_UPSTREAM_ERRORS)) +def test_upstream_errors_reach_the_caller_with_the_beta_sent_once(gateway: Gateway, row: str) -> None: + status, kind, message = _UPSTREAM_ERRORS[row] + marker: Final = uuid.uuid4().hex + with wire_server(_error_peer(status, kind, message)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + response: Final = _post(gateway, messages_body(model, marker)) + assert response.status_code == status and message in response.text, response.text + _expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + + +def _post_noting_whether_upstream_was_called( + gateway: Gateway, body: Mapping[str, JsonValue], wire: Wire +) -> tuple[httpx.Response, bool]: + before: Final = wire.received.qsize() + response: Final = _post(gateway, body) + assert response.status_code == 200, response.text + return response, wire.received.qsize() > before + + +def _until_served_from_cache(gateway: Gateway, body: Mapping[str, JsonValue], wire: Wire) -> tuple[httpx.Response, int]: + deadline: Final = time.monotonic() + 20 + for attempt in itertools.count(1): + response, called_upstream = _post_noting_whether_upstream_was_called(gateway, body, wire) + if not called_upstream: + return response, attempt + assert time.monotonic() < deadline, "the response cache never served the repeated request" + raise AssertionError("unreachable") + + +def test_the_response_cache_serves_the_repeated_request_without_an_upstream_call(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(anthropic_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=_FOUNDRY_MODEL, api_base=wire.url, api_key=_FOUNDRY_KEY) + body: Final = messages_body(model, marker, no_cache=False) + first: Final = _post(gateway, body) + _assert_answered(first, marker) + _expect_messages(_only(wire), target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) + _spend_row(f"msg_{marker}") + hit, attempts = _until_served_from_cache(gateway, body, wire) + assert hit.json()["id"] == first.json()["id"], hit.text + assert hit.json()["content"] == first.json()["content"], hit.text + _cache_hit_spend_row(f"msg_{marker}") + misses: Final = tuple(_observe(request) for request in wire.drain()) + assert len(misses) == attempts - 1, (attempts, misses) + for observed in misses: + _expect_messages(observed, target=_FOUNDRY_TARGET, auth=("x-api-key", _FOUNDRY_KEY), beta=BETA) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py index dcfd1d19ac2..a448a8cd2da 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py @@ -10,6 +10,7 @@ from anthropic.types import Message from integration._support import claude_code as cc from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows +from integration._support.safeguards_rig import DANGEROUS_TOOL_USE_BETA from integration._support.wire import Reply, Request, wire_server from pydantic import JsonValue, TypeAdapter @@ -280,7 +281,9 @@ def test_anthropic_messages_preserves_cli_beta_safeguards_and_result(gateway: Ga def respond(request: Request) -> Reply: assert request.method == "POST" and request.target == "/v1/messages", request.target assert request.headers["x-api-key"] == "synthetic-anthropic-key", request.headers - assert sorted(request.headers["anthropic-beta"].split(",")) == sorted(beta.split(",")), request.headers + assert sorted(request.headers["anthropic-beta"].split(",")) == sorted( + (*beta.split(","), DANGEROUS_TOOL_USE_BETA) + ), request.headers expected_body: Final = { "model": "claude-sonnet-4-6", "max_tokens": 32, diff --git a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py index 866bdf173a9..176dcfc86c9 100644 --- a/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py +++ b/tests/unit/llms/azure_ai/claude/test_azure_anthropic_messages_transformation.py @@ -7,13 +7,17 @@ sys.path.insert( 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) ) +from typing import Final from unittest.mock import patch +import httpx import pytest +import litellm from litellm.llms.azure_ai.anthropic.messages_transformation import ( AzureAnthropicMessagesConfig, ) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.types.router import GenericLiteLLMParams @@ -470,3 +474,80 @@ def test_azure_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_fl and info.get("supports_mid_conversation_system") is not True ] assert missing == [] + + +AUTO_MODE_SAFEGUARDS: Final = ({"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}},) + + +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("client_headers", "optional_params", "expected_beta"), + [ + ({}, {"safeguards": AUTO_MODE_SAFEGUARDS}, "dangerous-tool-use-2026-09-03"), + ( + {"anthropic-beta": "context-1m-2025-08-07"}, + {"safeguards": AUTO_MODE_SAFEGUARDS}, + "context-1m-2025-08-07,dangerous-tool-use-2026-09-03", + ), + ({"anthropic-beta": "context-1m-2025-08-07"}, {}, "context-1m-2025-08-07"), + ({}, {}, None), + ], +) +def test_safeguards_request_carries_the_dangerous_tool_use_beta_without_the_client_header( + client_headers: dict[str, str], optional_params: dict[str, object], expected_beta: str | None +) -> None: + headers, _ = AzureAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers=client_headers, + model="claude-sonnet-4-6", + messages=[{"role": "user", "content": "Run ls -la with the Bash tool"}], + optional_params=optional_params, + litellm_params={"api_key": "test-api-key"}, + ) + assert headers.get("anthropic-beta") == expected_beta, headers + + +def _foundry_answer(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + json={ + "id": "msg_safeguards", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hello from the gateway"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 6}, + }, + request=request, + ) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_outgoing_azure_request_sends_safeguards_with_the_dangerous_tool_use_beta() -> None: + received: Final[list[httpx.Request]] = [] + + def foundry_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return _foundry_answer(request) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(foundry_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=256, + messages=[{"role": "user", "content": "Use the Bash tool to run: echo hello from the gateway"}], + model="azure_ai/claude-sonnet-4-6", + api_key="test-api-key", + api_base="https://test-resource.services.ai.azure.com", + safeguards=AUTO_MODE_SAFEGUARDS, + client=upstream, + ) + finally: + await upstream.close() + + assert len(received) == 1, received + sent: Final = received[0] + assert "dangerous-tool-use-2026-09-03" in sent.headers["anthropic-beta"].split(","), dict(sent.headers) + assert json.loads(sent.content)["safeguards"] == list(AUTO_MODE_SAFEGUARDS) diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 50ed70e09f5..23c142b61ab 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5313,3 +5313,45 @@ async def test_retrieve_file_rejects_unverified_bedrock_range_error( ) finally: sync_client.close() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("local_beta_headers_config") +async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_its_beta() -> None: + received: Final[list[httpx.Request]] = [] + + def anthropic_answers(request: httpx.Request) -> httpx.Response: + received.append(request) + return httpx.Response( + status_code=200, + json={ + "id": "msg_dropped_safeguards", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, + }, + request=request, + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(anthropic_answers)) + try: + await litellm.anthropic.messages.acreate( + max_tokens=16, + messages=[{"role": "user", "content": "hi"}], + model="anthropic/claude-sonnet-4-6", + api_key="sk-ant-test", + safeguards=({"type": "dangerous_tool_use"},), + additional_drop_params=["safeguards"], + client=upstream, + ) + finally: + await upstream.close() + + assert len(received) == 1, received + sent: Final = received[0] + assert "safeguards" not in json.loads(sent.content) + assert "dangerous-tool-use-2026-09-03" not in sent.headers.get("anthropic-beta", ""), dict(sent.headers) diff --git a/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py index 5cbb9a0a7b4..46062f16c9a 100644 --- a/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py +++ b/tests/unit/llms/github_copilot/messages/test_github_copilot_messages_transformation.py @@ -1,3 +1,4 @@ +from typing import Final from unittest.mock import MagicMock import httpx @@ -9,6 +10,10 @@ from litellm.llms.github_copilot.common_utils import GetAPIKeyError from litellm.llms.github_copilot.messages.transformation import ( GithubCopilotAnthropicMessagesConfig, ) +from litellm.llms.github_copilot.per_user_auth import ( + GITHUB_COPILOT_USER_SESSION_KWARG_KEY, + GithubCopilotUserSession, +) def test_github_copilot_anthropic_messages_get_complete_url(): @@ -185,6 +190,26 @@ def test_github_copilot_anthropic_messages_validate_environment_injects_context_ assert "context-management-2025-06-27" in validated_headers["anthropic-beta"] +@pytest.mark.usefixtures("local_beta_headers_config") +def test_github_copilot_safeguards_do_not_derive_the_dangerous_tool_use_beta() -> None: + session: Final = GithubCopilotUserSession(token="gh.test-key", api_base="https://api.githubcopilot.com") + + validated_headers, _ = GithubCopilotAnthropicMessagesConfig().validate_anthropic_messages_environment( + headers={}, + model="github_copilot/claude-haiku-4.5", + messages=[{"role": "user", "content": "Hello"}], + optional_params={ + "safeguards": ({"type": "dangerous_tool_use"},), + "context_management": {"edits": ({"type": "clear_tool_uses_20250919"},)}, + }, + litellm_params={GITHUB_COPILOT_USER_SESSION_KWARG_KEY: session}, + api_key=None, + api_base=None, + ) + + assert validated_headers["anthropic-beta"] == "context-management-2025-06-27" + + def test_github_copilot_anthropic_messages_validate_environment_auth_error(): """Test error handling when authentication fails.""" config = GithubCopilotAnthropicMessagesConfig() diff --git a/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py b/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py index 9167853bf64..1d5ef50ca3a 100644 --- a/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py +++ b/tests/unit/llms/openai_like/messages/test_openai_like_anthropic_messages_transformation.py @@ -516,3 +516,26 @@ def test_request_strips_ttl_only_where_the_messages_api_defines_cache_control(co assert tool_result["content"][0]["cache_control"] == {"type": "ephemeral"} assert payload["messages"][1]["content"][1]["cache_control"] == {"type": "ephemeral"} assert payload["messages"][2] == {"role": "user", "content": "a plain string message"} + + +@pytest.mark.usefixtures("local_beta_headers_config") +@pytest.mark.parametrize( + ("client_headers", "expected_beta"), + [ + ({}, None), + ({"anthropic-beta": "dangerous-tool-use-2026-09-03"}, "dangerous-tool-use-2026-09-03"), + ], +) +def test_safeguards_reach_an_openai_like_host_with_only_the_client_betas( + config: OpenAILikeAnthropicMessagesConfig, client_headers: dict[str, str], expected_beta: str | None +) -> None: + headers, _ = config.validate_anthropic_messages_environment( + headers=client_headers, + model="some-model", + messages=[{"role": "user", "content": "hi"}], + optional_params={"safeguards": ({"type": "dangerous_tool_use"},)}, + litellm_params={}, + api_key="sk-test", + api_base="https://host/v1", + ) + assert headers.get("anthropic-beta") == expected_beta, headers