mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fix(anthropic): add the dangerous-tool-use beta when a /v1/messages request carries safeguards (#43955)
* fix(anthropic): add the dangerous-tool-use beta when a /v1/messages request carries safeguards * test(azure_ai): pin the outgoing Foundry request keeps the dangerous-tool-use beta next to safeguards * test(azure_ai): inject the Foundry client and scope the beta cache reset to the test * test(azure_ai): own and close the mock transport client, drop the docstring entry * test(integration): audit the safeguards beta derivation across anthropic-compatible providers * fix(github_copilot): keep safeguards from deriving the dangerous-tool-use beta * test(integration): kill the chaos victim worker through the harness helper * fix(anthropic): derive the dangerous-tool-use beta only where the beta map supports it and honor additional_drop_params The safeguards beta is now added only for providers whose beta map column accepts it, so github_copilot, Eden AI and the OpenAI-compatible JSON providers keep their base wire. additional_drop_params now runs before the provider headers are built, so a dropped param no longer drives its derived beta. * fix(anthropic): load the beta map off the event loop and pin a dropped context_management The first /v1/messages request that needs the remote beta map now loads it in a worker thread, so a slow fetch no longer stalls every other request on the worker. An integration cell pins that a deployment dropping context_management sends neither the field nor its derived beta. * fix(anthropic): load the beta map only for safeguards requests, once per worker * fix(anthropic): gate safeguards on the bundled beta map and load the remote map on its own thread --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
5dcf2c2062
commit
eb5295b20c
11 changed files with 1401 additions and 24 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue