From d3e1efcbba74348353afbae10be4483f670d22d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 12 Oct 2026 21:31:36 -0700 Subject: [PATCH] fix(rust): derive the dangerous-tool-use beta only for hosts with a beta map column DeepSeek and the other hosts without a column in anthropic_beta_headers_config.json keep forwarding safeguards and client betas as sent. Adds the wire-level audit tests. --- .../tests/messages/request.rs | 42 +- .../src/anthropic/messages/transformation.rs | 55 ++- .../src/deepseek/messages/transformation.rs | 23 + .../_anthropic_safeguards_beta_support.py | 284 ++++++++++++ .../test_anthropic_safeguards_beta_chaos.py | 158 +++++++ .../test_anthropic_safeguards_beta_wire.py | 411 ++++++++++++++++++ 6 files changed, 962 insertions(+), 11 deletions(-) create mode 100644 tests/integration/providers/_anthropic_safeguards_beta_support.py create mode 100644 tests/integration/providers/test_anthropic_safeguards_beta_chaos.py create mode 100644 tests/integration/providers/test_anthropic_safeguards_beta_wire.py diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index 543cdf1d8a1..0d0df2d131c 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -301,10 +301,6 @@ fn sent_betas(request: &wiremock::Request) -> BetaSet { json!({"tools": [{"type": "advisor_20260301", "name": "advisor", "model": MODEL}]}), &[AnthropicBeta::AdvisorTool20260301] )] -#[case::safeguards( - json!({"safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]}), - &[AnthropicBeta::DangerousToolUse20260903] -)] #[case::several_features_at_once( json!({"speed": "fast", "output_format": {"type": "json_schema"}}), &[AnthropicBeta::StructuredOutputs20251113, AnthropicBeta::FastMode20260201] @@ -347,6 +343,44 @@ async fn feature_betas_join_the_callers_betas_in_one_sorted_header( assert_eq!(sent, expected); } +#[rstest] +#[case::anthropic("anthropic", &[AnthropicBeta::DangerousToolUse20260903])] +#[case::deepseek("deepseek", &[])] +#[case::minimax("minimax", &[])] +#[tokio::test] +async fn safeguards_add_the_dangerous_tool_use_beta_only_where_the_beta_map_has_a_column( + call: MessagesCall, + #[case] provider: &str, + #[case] derived: &[AnthropicBeta], +) { + let upstream = upstream([message_response()]).await; + let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]); + + run_message(with_fields( + MessagesCall { + custom_llm_provider: Some(provider.into()), + api_key: Some("sk".into()), + api_base: Some(upstream.uri()), + extra_headers: headers([("Anthropic-Beta", "caller-beta-2025-01-01")]), + ..call + }, + json!({"safeguards": safeguards}), + )) + .await; + + let request = only_request(&upstream).await; + let expected: BetaSet = derived + .iter() + .cloned() + .chain([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())]) + .collect(); + assert_eq!(sent_betas(&request), expected); + assert_eq!( + request.body_json::().unwrap()["safeguards"], + safeguards + ); +} + #[rstest] #[tokio::test] async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: MessagesCall) { diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 72df83e3bf4..8b16ba10d77 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -109,7 +109,10 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { } fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers { - update_headers_with_anthropic_beta(headers, request) + merge_beta_headers( + headers, + provider_feature_betas(request, BetaProvider::Anthropic), + ) } } @@ -149,15 +152,24 @@ pub(crate) fn update_headers_with_anthropic_beta( merge_beta_headers(headers, feature_betas(request)) } -/// The betas a request needs on a host other than the first-party API: the Anthropic set -/// filtered and renamed through the host's `BetaProvider` policy. +/// The betas a request needs on a host with a column in the beta map, filtered and renamed +/// through that column. A host without one gets `feature_betas` alone, so `safeguards` reaches +/// it as sent. pub(crate) fn provider_feature_betas(request: &MessagesRequest, provider: BetaProvider) -> BetaSet { feature_betas(request) + .union(mapped_host_betas(&request.params)) .iter() .filter_map(|beta| beta.on(provider)) .collect() } +fn mapped_host_betas(params: &MessagesOptionalParams) -> BetaSet { + carries_safeguards(params) + .then_some(AnthropicBeta::DangerousToolUse20260903) + .into_iter() + .collect() +} + fn feature_betas(request: &MessagesRequest) -> BetaSet { let params = &request.params; let tools = params.tools.as_deref(); @@ -171,7 +183,6 @@ fn feature_betas(request: &MessagesRequest) -> BetaSet { .then_some(AnthropicBeta::PerTurnControl20260701), has_advisor_tool(tools).then_some(AnthropicBeta::AdvisorTool20260301), is_tool_search_used(tools).then_some(AnthropicBeta::AdvancedToolUse20251120), - carries_safeguards(params).then_some(AnthropicBeta::DangerousToolUse20260903), ] .into_iter() .flatten() @@ -1035,13 +1046,43 @@ mod tests { )] #[case::safeguards( json!({"safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}]}), - &["dangerous-tool-use-2026-09-03"] + &[] )] - #[case::null_safeguards(json!({"safeguards": null}), &[])] fn feature_betas_follow_the_request(#[case] fields: Value, #[case] expected: &[&str]) { assert_eq!(feature_betas(&request(fields)), betas(expected)); } + #[rstest] + #[case::anthropic(BetaProvider::Anthropic, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])] + #[case::vertex_ai(BetaProvider::VertexAi, json!([{"type": "dangerous_tool_use"}]), &["dangerous-tool-use-2026-09-03"])] + #[case::empty_list(BetaProvider::Anthropic, json!([]), &["dangerous-tool-use-2026-09-03"])] + #[case::not_a_list(BetaProvider::Anthropic, json!(1), &["dangerous-tool-use-2026-09-03"])] + #[case::null(BetaProvider::Anthropic, Value::Null, &[])] + #[case::bedrock_converse_has_a_null_cell(BetaProvider::BedrockConverse, json!([{"type": "dangerous_tool_use"}]), &[])] + #[case::databricks_has_a_null_cell(BetaProvider::Databricks, json!([{"type": "dangerous_tool_use"}]), &[])] + fn safeguards_add_the_dangerous_tool_use_beta_where_the_beta_map_has_a_cell( + #[case] provider: BetaProvider, + #[case] safeguards: Value, + #[case] expected: &[&str], + ) { + assert_eq!( + provider_feature_betas(&request(json!({"safeguards": safeguards})), provider), + betas(expected) + ); + } + + #[test] + fn hosts_without_a_beta_map_column_get_safeguards_without_a_derived_beta() { + let client = headers(&[("anthropic-beta", "claude-code-20250219")]); + assert_eq!( + update_headers_with_anthropic_beta( + client.clone(), + &request(json!({"safeguards": [{"type": "dangerous_tool_use"}]})) + ), + client + ); + } + #[rstest] #[case::no_betas(&[("x-api-key", "k"), ("anthropic-version", "2023-06-01")], json!({}), &[("x-api-key", "k"), ("anthropic-version", "2023-06-01")])] #[case::blank_beta_header(&[("Anthropic-Beta", " , "), ("x-api-key", "k")], json!({}), &[("Anthropic-Beta", " , "), ("x-api-key", "k")])] @@ -1080,7 +1121,7 @@ mod tests { json!({"messages": [{"role": "user", "content": "hi", "output_config": {"effort": "low"}}]}), &[("anthropic-beta", "claude-code-20250219,context-management-2025-06-27,effort-2025-11-24,interleaved-thinking-2025-05-14,per-turn-control-2026-07-01")], )] - #[case::safeguards_join_the_client_betas_without_the_dangerous_tool_use_beta( + #[case::safeguards_add_the_dangerous_tool_use_beta_to_the_client_betas( &[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")], json!({"safeguards": [{"type": "dangerous_tool_use"}]}), &[("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03,interleaved-thinking-2025-05-14")], diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs index 9359184f40e..fefedbbbf5e 100644 --- a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -451,4 +451,27 @@ mod tests { json!({"type": "enabled", "budget_tokens": 1024}) ); } + + #[rstest] + #[case::without_client_betas(&[])] + #[case::with_client_betas(&[("anthropic-beta", "claude-code-20250219,interleaved-thinking-2025-05-14")])] + fn safeguards_and_client_betas_reach_deepseek_as_sent(#[case] client: &[(&str, &str)]) { + let safeguards = json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]); + let body = json!({ + "model": "deepseek-v4-pro", + "max_tokens": 100, + "messages": [{"role": "user", "content": "hi"}], + "safeguards": safeguards + }); + let headers: Headers = client + .iter() + .map(|(name, value)| (name.to_string(), value.to_string())) + .collect(); + assert_eq!( + DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG + .request_headers(headers.clone(), &request_from(body.clone())), + headers + ); + assert_eq!(transformed(body)["safeguards"], safeguards); + } } diff --git a/tests/integration/providers/_anthropic_safeguards_beta_support.py b/tests/integration/providers/_anthropic_safeguards_beta_support.py new file mode 100644 index 00000000000..90ec06bf385 --- /dev/null +++ b/tests/integration/providers/_anthropic_safeguards_beta_support.py @@ -0,0 +1,284 @@ +from __future__ import annotations + +import json +import re +import threading +import uuid +from collections.abc import Callable, Mapping +from dataclasses import dataclass +from pathlib import Path +from queue import SimpleQueue +from types import MappingProxyType +from typing import Final + +import yaml +from integration._support.anthropic_sse import ( + CONTENT_BLOCK_START, + error_body, + message_start, + sse, + text_delta, +) +from integration._support.client import Gateway, object_value +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, Wire +from pydantic import JsonValue, TypeAdapter + +DANGEROUS_TOOL_USE: Final = "dangerous-tool-use-2026-09-03" +WEB_SEARCH: Final = "web-search-2025-03-05" +CLIENT_BETAS: Final = "interleaved-thinking-2025-05-14,claude-code-20250219,interleaved-thinking-2025-05-14" +MERGED_BETAS: Final = f"claude-code-20250219,{DANGEROUS_TOOL_USE},interleaved-thinking-2025-05-14" +BACKEND: Final = "claude-sonnet-4-6" +ANTHROPIC_DEPLOYMENT: Final = "safeguarded-claude" +VERTEX_DEPLOYMENT: Final = "safeguarded-vertex-claude" +DEEPSEEK_DEPLOYMENT: Final = "safeguarded-deepseek" +ANTHROPIC_KEY: Final = "synthetic-anthropic-key" +DEEPSEEK_KEY: Final = "synthetic-deepseek-key" +PROJECT: Final = "scripted-project" +LOCATION: Final = "us-east5" +_VERTEX_MODEL_PATH: Final = f"/v1/projects/{PROJECT}/locations/{LOCATION}/publishers/anthropic/models/{BACKEND}" +VERTEX_TARGET: Final = f"{_VERTEX_MODEL_PATH}:rawPredict" +VERTEX_STREAM_TARGET: Final = f"{_VERTEX_MODEL_PATH}:streamRawPredict?alt=sse" +ANTHROPIC_TARGET: Final = "/v1/messages" +DEEPSEEK_TARGET: Final = "/anthropic/v1/messages" +SCRIPTED_TARGETS: Final = frozenset({ANTHROPIC_TARGET, DEEPSEEK_TARGET, VERTEX_TARGET, VERTEX_STREAM_TARGET}) +REPLY_TEXT: Final = "safeguards beta control" +EXTRA_INPUTS_MESSAGE: Final = "safeguards: Extra inputs are not permitted" +NOT_A_LIST_MESSAGE: Final = "safeguards: Input should be a valid list" +OVERLOADED_MESSAGE: Final = "Overloaded" +SAFEGUARDS: Final[tuple[JsonValue, ...]] = ( + { + "type": "dangerous_tool_use", + "classifier_context": {"v": 1, "permission_mode": "auto"}, + }, +) +SAFEGUARD_RESULTS: Final[tuple[JsonValue, ...]] = ( + {"type": "dangerous_tool_use", "status": {"type": "available", "tool_uses": {}}}, +) +BASH_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "name": "Bash", + "description": "Run a shell command", + "input_schema": { + "type": "object", + "properties": {"command": {"type": "string"}}, + "required": ["command"], + }, + } +) +WEB_SEARCH_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"type": "web_search_20250305", "name": "web_search", "max_uses": 1} +) +RUST_ROUTE_HEADER: Final = "x-litellm-rust" +_MARKER: Final = re.compile(rb"marker-[0-9a-f]{32}") +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_STARTED_WORKER: Final = "Started server process [" + + +@dataclass(frozen=True, slots=True) +class Delivery: + target: str + beta: str | None + safeguards: JsonValue + marker: str + authorization: str | None + api_key: str | None + + +def new_marker() -> str: + return f"marker-{uuid.uuid4().hex}" + + +def marker_of(request: Request) -> str: + found: Final = tuple(match.group(0) for match in _MARKER.finditer(request.body)) + return found[-1].decode() if found else "unmarked" + + +def body_of(request: Request) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_json(request.body) + + +def json_object(text: str) -> dict[str, JsonValue]: + return _JSON_OBJECT.validate_json(text) + + +def betas_of(request: Request) -> frozenset[str]: + return frozenset(part.strip() for part in request.headers.get("anthropic-beta", "").split(",") if part.strip()) + + +def started_workers(log: Path) -> tuple[int, ...]: + return tuple( + int(line.rsplit("[", 1)[1].strip(" ]\r")) for line in log.read_text().splitlines() if _STARTED_WORKER in line + ) + + +def _streams(request: Request) -> bool: + return request.target == VERTEX_STREAM_TARGET or body_of(request).get("stream") is True + + +def _results(safeguards: JsonValue) -> list[JsonValue] | None: + if safeguards is None: + return None + return list(SAFEGUARD_RESULTS) if safeguards else [] + + +def _message(identity: str, results: list[JsonValue] | None) -> dict[str, JsonValue]: + verdict: Final[dict[str, JsonValue]] = {} if results is None else {"safeguard_results": results} + return { + "id": identity, + "type": "message", + "role": "assistant", + "model": BACKEND, + "content": [{"type": "text", "text": REPLY_TEXT}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 3}, + **verdict, + } + + +def _stream(identity: str, results: list[JsonValue] | None) -> tuple[bytes, ...]: + verdict: Final[dict[str, JsonValue]] = {} if results is None else {"safeguard_results": results} + return ( + message_start(identity, BACKEND), + CONTENT_BLOCK_START, + text_delta(REPLY_TEXT), + sse("content_block_stop", {"type": "content_block_stop", "index": 0}), + sse( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None, **verdict}, + "usage": {"output_tokens": 3}, + }, + ), + sse("message_stop", {"type": "message_stop"}), + ) + + +def anthropic_peer(request: Request) -> Reply: + if request.target not in SCRIPTED_TARGETS: + return Reply( + status=404, + body=json.dumps({"error": f"unscripted target {request.target}"}).encode(), + ) + if request.target == DEEPSEEK_TARGET: + return Reply(body=json.dumps(_message(f"msg_{marker_of(request)}", None)).encode()) + safeguards: Final = body_of(request).get("safeguards") + if safeguards is not None and DANGEROUS_TOOL_USE not in betas_of(request): + return Reply(status=400, body=error_body(400, EXTRA_INPUTS_MESSAGE)) + if safeguards is not None and not isinstance(safeguards, list): + return Reply(status=400, body=error_body(400, NOT_A_LIST_MESSAGE)) + identity: Final = f"msg_{marker_of(request)}" + results: Final = _results(safeguards) + if _streams(request): + return Reply(content_type="text/event-stream", chunks=_stream(identity, results)) + return Reply(body=json.dumps(_message(identity, results)).encode()) + + +def outage_peer(outage: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if outage.is_set(): + return Reply(status=529, body=error_body(529, OVERLOADED_MESSAGE)) + return anthropic_peer(request) + + return respond + + +def holding_peer(held: SimpleQueue[str], release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + held.put(marker_of(request)) + assert release.wait(timeout=60), "Held request was never released" + return anthropic_peer(request) + + return respond + + +def deliveries(wire: Wire) -> tuple[Delivery, ...]: + return tuple( + Delivery( + request.target, + request.headers.get("anthropic-beta"), + body_of(request).get("safeguards"), + marker_of(request), + request.headers.get("authorization"), + request.headers.get("x-api-key"), + ) + for request in wire.drain() + ) + + +def one_delivery(wire: Wire, marker: str) -> Delivery: + (sent,) = deliveries(wire) + assert sent.marker == marker, sent + return sent + + +def rust_routed(headers: Mapping[str, str]) -> bool: + return headers.get(RUST_ROUTE_HEADER) == "true" + + +def messages_body(model: str, marker: str, *, stream: bool = False) -> dict[str, JsonValue]: + return { + "model": model, + "max_tokens": 64, + "stream": stream, + "messages": [{"role": "user", "content": marker}], + "tools": [dict(BASH_TOOL)], + } + + +def safeguarded(body: Mapping[str, JsonValue], safeguards: JsonValue) -> dict[str, JsonValue]: + return {**body, "safeguards": safeguards} + + +def chat_body(model: str, marker: str) -> dict[str, JsonValue]: + return {"model": model, "messages": [{"role": "user", "content": marker}]} + + +def responses_body(model: str, marker: str) -> dict[str, JsonValue]: + return {"model": model, "input": marker} + + +def owned_config(directory: Path, gateway: Gateway, api_base: str) -> Path: + config: Final = _JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + merged: Final = { + **config, + "model_list": [ + { + "model_name": ANTHROPIC_DEPLOYMENT, + "litellm_params": { + "model": f"anthropic/{BACKEND}", + "api_base": api_base, + "api_key": ANTHROPIC_KEY, + }, + }, + { + "model_name": VERTEX_DEPLOYMENT, + "litellm_params": { + "model": f"vertex_ai/{BACKEND}", + "api_base": api_base, + "vertex_project": PROJECT, + "vertex_location": LOCATION, + "vertex_credentials": service_account_json(PROJECT, gateway.upstream_url.rstrip("/")), + }, + }, + { + "model_name": DEEPSEEK_DEPLOYMENT, + "litellm_params": { + "model": "deepseek/deepseek-chat", + "api_base": api_base, + "api_key": DEEPSEEK_KEY, + }, + }, + ], + "router_settings": { + **object_value(config["router_settings"]), + "num_retries": 0, + }, + } + path: Final = directory / f"safeguards-beta-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(merged)) + return path diff --git a/tests/integration/providers/test_anthropic_safeguards_beta_chaos.py b/tests/integration/providers/test_anthropic_safeguards_beta_chaos.py new file mode 100644 index 00000000000..950c9f8e03b --- /dev/null +++ b/tests/integration/providers/test_anthropic_safeguards_beta_chaos.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +import threading +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack +from pathlib import Path +from queue import SimpleQueue +from typing import Final + +import httpcore +import httpx +import psutil +import pytest +from integration._support.client import Gateway, eventually +from integration._support.process import graceful_stop_seconds, owned_proxy_process +from integration._support.wire import wire_server +from integration.providers._anthropic_safeguards_beta_support import ( + ANTHROPIC_DEPLOYMENT, + DANGEROUS_TOOL_USE, + SAFEGUARDS, + deliveries, + holding_peer, + messages_body, + new_marker, + owned_config, + rust_routed, + safeguarded, + started_workers, +) +from pydantic import TypeAdapter + +_CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +_RESTART_TIMEOUT: Final = int(4 * graceful_stop_seconds() + 240) +_BURST: Final = 12 +_SECOND_WAVE: Final = 6 +_CLIENT_ADDRESS: Final = TypeAdapter(tuple[str, int]) + + +def _clients(stack: ExitStack, base_url: str, count: int) -> tuple[httpx.Client, ...]: + return tuple( + stack.enter_context(httpx.Client(base_url=base_url, timeout=120, trust_env=False)) for _ in range(count) + ) + + +def _post(client: httpx.Client, key: str, marker: str) -> int: + body: Final = safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker), list(SAFEGUARDS)) + response: Final = client.post("/v1/messages", json=body, headers={"Authorization": f"Bearer {key}"}) + assert response.status_code != 200 or rust_routed(response.headers), dict(response.headers) + return response.status_code + + +def _post_or_dropped(client: httpx.Client, key: str, marker: str) -> int | None: + try: + return _post(client, key, marker) + except httpx.TransportError: + return None + + +def _local_port(client: httpx.Client) -> int: + with client.stream("GET", "/health/liveliness") as response: + stream: Final[object] = response.extensions["network_stream"] # pyright: ignore[reportAny] # httpx types extensions as Any + assert isinstance(stream, httpcore.NetworkStream), stream + port: Final = _CLIENT_ADDRESS.validate_python(stream.get_extra_info("client_addr"))[1] + response.read() + assert response.status_code == 200, response.text + return port + + +def _accepted_client_ports(pid: int, proxy_port: int) -> frozenset[int]: + return frozenset( + connection.raddr.port + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.raddr and connection.laddr.port == proxy_port + ) + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_worker_sigkill_mid_safeguards_burst_leaves_the_replacement_adding_the_beta( + gateway: Gateway, tmp_path: Path +) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(holding_peer(held, release))) + stack.callback(release.set) + config: Final = owned_config(tmp_path, gateway, wire.url) + owned: Final = stack.enter_context( + owned_proxy_process(gateway, tmp_path, {"LITELLM_RUST": "true"}, config=config, workers=2) + ) + proxy_url: Final = owned.gateway.client.base_url + workers: Final = eventually(lambda: started_workers(owned.log), lambda pids: len(pids) == 2, seconds=30) + clients: Final = _clients(stack, str(proxy_url), _BURST) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(clients))) + ports: Final = tuple(_local_port(client) for client in clients) + markers: Final = tuple(new_marker() for _ in clients) + futures: Final = tuple( + pool.submit(_post_or_dropped, client, gateway.key, marker) + for client, marker in zip(clients, markers, strict=True) + ) + eventually(held.qsize, lambda size: size == len(clients), seconds=30) + shares: Final = {pid: _accepted_client_ports(pid, proxy_url.port or 0) & frozenset(ports) for pid in workers} + assert sum(map(len, shares.values())) == len(clients), shares + victim: Final = min((pid for pid in workers if shares[pid]), key=lambda pid: len(shares[pid])) + with owned.replacing_worker(victim) as kill: + kill() + release.set() + results: Final = tuple(future.result(timeout=60) for future in futures) + for port, result in zip(ports, results, strict=True): + assert result == (None if port in shares[victim] else 200), ( + port, + result, + shares, + ) + second_wave: Final = tuple(new_marker() for _ in range(_SECOND_WAVE)) + for marker in second_wave: + assert _post(owned.gateway.client, gateway.key, marker) == 200 + delivered: Final = deliveries(wire) + assert tuple(sent.beta for sent in delivered) == (DANGEROUS_TOOL_USE,) * (len(clients) + len(second_wave)), ( + delivered + ) + seen: Final = tuple(sum(sent.marker == marker for sent in delivered) for marker in markers + second_wave) + assert seen == (1,) * len(seen), seen + assert owned.process.poll() is None + + +@pytest.mark.timeout(_RESTART_TIMEOUT) +def test_proxy_shutdown_mid_safeguards_burst_then_a_fresh_boot_keeps_adding_the_beta( + gateway: Gateway, tmp_path: Path +) -> None: + held: Final[SimpleQueue[str]] = SimpleQueue() + release: Final = threading.Event() + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(holding_peer(held, release))) + stack.callback(release.set) + config: Final = owned_config(tmp_path, gateway, wire.url) + markers: Final = tuple(new_marker() for _ in range(_BURST)) + pool: Final = stack.enter_context(ThreadPoolExecutor(max_workers=len(markers))) + with owned_proxy_process(gateway, tmp_path, {"LITELLM_RUST": "true"}, config=config, workers=2) as first: + clients: Final = _clients(stack, str(first.gateway.client.base_url), len(markers)) + futures: Final = tuple( + pool.submit(_post_or_dropped, client, gateway.key, marker) + for client, marker in zip(clients, markers, strict=True) + ) + eventually(held.qsize, lambda size: size == len(markers), seconds=30) + first.begin_shutdown() + release.set() + results: Final = tuple(future.result(timeout=120) for future in futures) + during: Final = deliveries(wire) + assert all(result in (None, 200) for result in results), results + assert tuple(sent.beta for sent in during) == (DANGEROUS_TOOL_USE,) * len(markers), during + assert sorted(sent.marker for sent in during) == sorted(markers), during + with owned_proxy_process(gateway, tmp_path, {"LITELLM_RUST": "true"}, config=config, workers=2) as second: + after: Final = tuple(new_marker() for _ in range(_SECOND_WAVE)) + statuses: Final = tuple(_post(second.gateway.client, gateway.key, marker) for marker in after) + assert statuses == (200,) * len(after), statuses + delivered: Final = deliveries(wire) + assert tuple(sent.beta for sent in delivered) == (DANGEROUS_TOOL_USE,) * len(after), delivered + assert sorted(sent.marker for sent in delivered) == sorted(after), delivered diff --git a/tests/integration/providers/test_anthropic_safeguards_beta_wire.py b/tests/integration/providers/test_anthropic_safeguards_beta_wire.py new file mode 100644 index 00000000000..d99d4b72919 --- /dev/null +++ b/tests/integration/providers/test_anthropic_safeguards_beta_wire.py @@ -0,0 +1,411 @@ +from __future__ import annotations + +import json +import threading +from collections.abc import Iterator, Mapping +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from typing import Final + +import anthropic +import httpx +import pytest +from integration._support.client import ( + Gateway, + Scenario, + eventually, + gateway_from_environment, +) +from integration._support.database import read_rows +from integration._support.process import ( + OwnedProxy, + graceful_stop_seconds, + owned_proxy_process, +) +from integration._support.wire import Wire, wire_server +from integration.providers._anthropic_safeguards_beta_support import ( + ANTHROPIC_DEPLOYMENT, + ANTHROPIC_KEY, + ANTHROPIC_TARGET, + BACKEND, + BASH_TOOL, + CLIENT_BETAS, + DANGEROUS_TOOL_USE, + DEEPSEEK_DEPLOYMENT, + DEEPSEEK_KEY, + DEEPSEEK_TARGET, + EXTRA_INPUTS_MESSAGE, + MERGED_BETAS, + NOT_A_LIST_MESSAGE, + OVERLOADED_MESSAGE, + REPLY_TEXT, + SAFEGUARD_RESULTS, + SAFEGUARDS, + VERTEX_DEPLOYMENT, + VERTEX_STREAM_TARGET, + VERTEX_TARGET, + WEB_SEARCH, + WEB_SEARCH_TOOL, + anthropic_peer, + chat_body, + deliveries, + json_object, + messages_body, + new_marker, + one_delivery, + outage_peer, + owned_config, + responses_body, + rust_routed, + safeguarded, + started_workers, +) +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.xdist_group("anthropic_safeguards_beta") + +_CELL_TIMEOUT: Final = int(2 * graceful_stop_seconds() + 120) +_BURST: Final = 24 +_RECOVERY: Final = 6 +_LARGE_ENTRY: Final[JsonValue] = [ + { + "type": "dangerous_tool_use", + "classifier_context": {"v": 1, "permission_mode": "auto", "note": "x" * 5120}, + } +] + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + wire: Wire + outage: threading.Event + + @property + def gateway(self) -> Gateway: + return self.owned.gateway + + +@dataclass(frozen=True, slots=True) +class _Answer: + path: str + marker: str + status: int + text: str + rust: bool + + +@pytest.fixture(scope="module") +def shared_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("anthropic-safeguards-beta") + outage: Final = threading.Event() + with ( + gateway_from_environment() as environment, + wire_server(outage_peer(outage)) as wire, + ): + config: Final = owned_config(directory, environment, wire.url) + with owned_proxy_process(environment, directory, {"LITELLM_RUST": "true"}, config=config, workers=2) as owned: + eventually( + lambda: started_workers(owned.log), + lambda pids: len(pids) == 2, + seconds=60, + ) + wire.drain() + yield _Rig(owned, wire, outage) + + +@pytest.fixture +def rig(shared_rig: _Rig) -> _Rig: + shared_rig.outage.clear() + shared_rig.wire.drain() + return shared_rig + + +def _spend_rows(call_id: str) -> list[dict[str, JsonValue]]: + return read_rows( + 'SELECT request_id FROM "LiteLLM_SpendLogs" WHERE litellm_call_id=%s', + (call_id,), + ) + + +def _post(gateway: Gateway, body: Mapping[str, JsonValue], *, beta: str | None = None) -> httpx.Response: + return gateway.request( + "POST", + "/v1/messages", + body, + headers={} if beta is None else {"anthropic-beta": beta}, + ) + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_sdk_messages_with_safeguards_go_upstream_with_the_dangerous_tool_use_beta( + rig: _Rig, +) -> None: + marker: Final = new_marker() + with httpx.Client(trust_env=False, timeout=60) as http: + client: Final = anthropic.Anthropic( + base_url=str(rig.gateway.client.base_url), + api_key=rig.gateway.key, + max_retries=0, + http_client=http, + ) + raw: Final = client.messages.with_raw_response.create( + model=ANTHROPIC_DEPLOYMENT, + max_tokens=64, + messages=[{"role": "user", "content": marker}], + extra_body={"safeguards": list(SAFEGUARDS), "tools": [dict(BASH_TOOL)]}, + ) + message: Final = raw.parse() + assert rust_routed(raw.headers), dict(raw.headers) + received: Final = json_object(message.model_dump_json()) + assert message.id == f"msg_{marker}", received + assert received.get("safeguard_results") == list(SAFEGUARD_RESULTS), received + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.api_key, sent.beta) == ( + ANTHROPIC_TARGET, + ANTHROPIC_KEY, + DANGEROUS_TOOL_USE, + ), sent + assert sent.safeguards == list(SAFEGUARDS), sent + assert "x-litellm-call-id" in raw.headers, dict(raw.headers) + call_id: Final = raw.headers["x-litellm-call-id"] + rows: Final = eventually(lambda: _spend_rows(call_id), lambda found: len(found) == 1, seconds=70) + assert len(rows) == 1, rows + + +@pytest.mark.timeout(_CELL_TIMEOUT) +async def test_sdk_streamed_messages_with_safeguards_go_upstream_with_the_beta( + rig: _Rig, +) -> None: + marker: Final = new_marker() + async with httpx.AsyncClient(trust_env=False, timeout=60) as http: + client: Final = anthropic.AsyncAnthropic( + base_url=str(rig.gateway.client.base_url), + api_key=rig.gateway.key, + max_retries=0, + http_client=http, + ) + async with client.messages.stream( + model=ANTHROPIC_DEPLOYMENT, + max_tokens=64, + messages=[{"role": "user", "content": marker}], + extra_body={"safeguards": list(SAFEGUARDS), "tools": [dict(BASH_TOOL)]}, + ) as stream: + text: Final = "".join([piece async for piece in stream.text_stream]) + final: Final = await stream.get_final_message() + assert rust_routed(stream.response.headers), dict(stream.response.headers) + assert (final.id, text) == (f"msg_{marker}", REPLY_TEXT), (final, text) + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.beta, sent.safeguards) == ( + ANTHROPIC_TARGET, + DANGEROUS_TOOL_USE, + list(SAFEGUARDS), + ), sent + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_raw_messages_merge_the_safeguards_beta_into_the_client_betas(rig: _Rig, stream: bool) -> None: + marker: Final = new_marker() + body: Final = safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker, stream=stream), list(SAFEGUARDS)) + response: Final = _post(rig.gateway, body, beta=CLIENT_BETAS) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text and "safeguard_results" in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.beta) == (ANTHROPIC_TARGET, MERGED_BETAS), sent + + +@pytest.mark.parametrize("extra", [{"safeguards": None}, {}], ids=["null", "absent"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_null_or_absent_safeguards_add_no_beta(rig: _Rig, extra: Mapping[str, JsonValue]) -> None: + marker: Final = new_marker() + response: Final = _post(rig.gateway, {**messages_body(ANTHROPIC_DEPLOYMENT, marker), **extra}) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text and "safeguard_results" not in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.beta, sent.safeguards) == (None, None), sent + + +@pytest.mark.parametrize( + ("safeguards", "status", "fragment"), + [ + (7, 400, NOT_A_LIST_MESSAGE), + ("dangerous_tool_use", 400, NOT_A_LIST_MESSAGE), + ({}, 400, NOT_A_LIST_MESSAGE), + ([], 200, "safeguard_results"), + (_LARGE_ENTRY, 200, "safeguard_results"), + ], + ids=["int", "string", "object", "empty_list", "5kb_entry"], +) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_malformed_safeguards_still_add_the_beta_and_relay_the_vendor_verdict( + rig: _Rig, safeguards: JsonValue, status: int, fragment: str +) -> None: + marker: Final = new_marker() + response: Final = _post( + rig.gateway, + safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker), safeguards), + ) + assert (response.status_code, fragment in response.text) == (status, True), response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.beta, sent.safeguards) == (DANGEROUS_TOOL_USE, safeguards), sent + control: Final = new_marker() + after: Final = _post(rig.gateway, messages_body(ANTHROPIC_DEPLOYMENT, control)) + assert after.status_code == 200, after.text + assert rust_routed(after.headers), dict(after.headers) + assert one_delivery(rig.wire, control).beta is None + assert rig.gateway.request("GET", "/health/liveliness").status_code == 200 + + +@pytest.mark.parametrize("last", ["null", "list"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_duplicated_safeguards_keys_resolve_to_the_last_value(rig: _Rig, last: str) -> None: + marker: Final = new_marker() + opening: Final = json.dumps(messages_body(ANTHROPIC_DEPLOYMENT, marker))[:-1] + null_field: Final = '"safeguards": null' + list_field: Final = f'"safeguards": {json.dumps(list(SAFEGUARDS))}' + ordered: Final = (null_field, list_field) if last == "list" else (list_field, null_field) + raw: Final = f"{opening}, {ordered[0]}, {ordered[1]}}}".encode() + response: Final = rig.gateway.client.post( + "/v1/messages", + content=raw, + headers={ + "Authorization": f"Bearer {rig.gateway.key}", + "content-type": "application/json", + }, + ) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert ("safeguard_results" in response.text) == (last == "list"), response.text + sent: Final = one_delivery(rig.wire, marker) + assert sent.beta == (DANGEROUS_TOOL_USE if last == "list" else None), sent + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_safeguards_request_with_a_bad_key_never_reaches_the_peer(rig: _Rig) -> None: + body: Final = safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, new_marker()), list(SAFEGUARDS)) + response: Final = rig.gateway.request("POST", "/v1/messages", body, key="sk-not-a-key") + assert response.status_code == 401, response.text + assert deliveries(rig.wire) == () + + +@pytest.mark.parametrize("web_search", [False, True], ids=["bash_only", "with_web_search"]) +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_vertex_messages_with_safeguards_keep_the_beta_from_the_shared_helper( + rig: _Rig, stream: bool, web_search: bool +) -> None: + marker: Final = new_marker() + tools: Final[list[JsonValue]] = [ + dict(BASH_TOOL), + *((dict(WEB_SEARCH_TOOL),) if web_search else ()), + ] + body: Final = safeguarded(messages_body(VERTEX_DEPLOYMENT, marker, stream=stream), list(SAFEGUARDS)) + response: Final = _post(rig.gateway, {**body, "tools": tools}) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text and "safeguard_results" in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert sent.target == (VERTEX_STREAM_TARGET if stream else VERTEX_TARGET), sent + assert sent.authorization == "Bearer scripted-token", sent + assert sent.beta == (f"{DANGEROUS_TOOL_USE},{WEB_SEARCH}" if web_search else DANGEROUS_TOOL_USE), sent + + +def _beta_set(header: str | None) -> frozenset[str]: + return frozenset(part.strip() for part in (header or "").split(",") if part.strip()) + + +@pytest.mark.parametrize("beta", [None, CLIENT_BETAS], ids=["no-client-betas", "client-betas"]) +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_deepseek_messages_forward_safeguards_and_client_betas_as_sent(rig: _Rig, beta: str | None) -> None: + marker: Final = new_marker() + response: Final = _post( + rig.gateway, + safeguarded(messages_body(DEEPSEEK_DEPLOYMENT, marker), list(SAFEGUARDS)), + beta=beta, + ) + assert response.status_code == 200, response.text + assert rust_routed(response.headers), dict(response.headers) + assert f"msg_{marker}" in response.text, response.text + sent: Final = one_delivery(rig.wire, marker) + assert (sent.target, sent.api_key, sent.safeguards) == (DEEPSEEK_TARGET, DEEPSEEK_KEY, list(SAFEGUARDS)), sent + assert _beta_set(sent.beta) == _beta_set(beta), sent + + +def test_python_route_leaves_safeguards_without_the_beta(gateway: Gateway, scenario: Scenario) -> None: + with wire_server(anthropic_peer) as wire: + deployment: Final = scenario.model(model=f"anthropic/{BACKEND}", api_base=wire.url, api_key=ANTHROPIC_KEY) + marker: Final = new_marker() + response: Final = _post(gateway, safeguarded(messages_body(deployment, marker), list(SAFEGUARDS))) + assert response.status_code == 400, response.text + assert EXTRA_INPUTS_MESSAGE in response.text, response.text + (sent,) = deliveries(wire) + assert (sent.marker, sent.beta, sent.safeguards) == ( + marker, + None, + list(SAFEGUARDS), + ), sent + assert not rust_routed(response.headers), dict(response.headers) + + +def _burst_request(index: int, marker: str) -> tuple[str, dict[str, JsonValue]]: + if index % 4 == 0: + return "/v1/messages", safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker), list(SAFEGUARDS)) + if index % 4 == 1: + return "/v1/messages", safeguarded(messages_body(ANTHROPIC_DEPLOYMENT, marker, stream=True), list(SAFEGUARDS)) + if index % 4 == 2: + return "/v1/chat/completions", chat_body(ANTHROPIC_DEPLOYMENT, marker) + return "/v1/responses", responses_body(ANTHROPIC_DEPLOYMENT, marker) + + +def _answer(gateway: Gateway, path: str, marker: str, body: Mapping[str, JsonValue]) -> _Answer: + response: Final = gateway.request("POST", path, body) + return _Answer(path, marker, response.status_code, response.text, rust_routed(response.headers)) + + +def _wave(pool: ThreadPoolExecutor, gateway: Gateway, markers: tuple[str, ...], offset: int) -> tuple[_Answer, ...]: + planned: Final = tuple((marker, *_burst_request(offset + index, marker)) for index, marker in enumerate(markers)) + submitted: Final = tuple(pool.submit(_answer, gateway, path, marker, body) for marker, path, body in planned) + return tuple(future.result(timeout=120) for future in submitted) + + +def _assert_served(answer: _Answer) -> None: + assert answer.status == 200, answer + assert REPLY_TEXT in answer.text, answer + assert answer.rust == (answer.path == "/v1/messages"), answer + if answer.path == "/v1/messages": + assert f"msg_{answer.marker}" in answer.text and "safeguard_results" in answer.text, answer + + +@pytest.mark.timeout(_CELL_TIMEOUT) +def test_peer_outage_mid_safeguards_burst_relays_each_verdict_once_and_recovers( + rig: _Rig, +) -> None: + healthy: Final = tuple(new_marker() for _ in range(_BURST // 2)) + outage: Final = tuple(new_marker() for _ in range(_BURST // 2)) + recovery: Final = tuple(new_marker() for _ in range(_RECOVERY)) + with ThreadPoolExecutor(max_workers=8) as pool: + before: Final = _wave(pool, rig.gateway, healthy, 0) + rig.outage.set() + assert rig.gateway.request("GET", "/health/liveliness").status_code == 200 + during: Final = _wave(pool, rig.gateway, outage, len(healthy)) + rig.outage.clear() + after: Final = _wave(pool, rig.gateway, recovery, len(healthy) + len(outage)) + for answer in (*before, *after): + _assert_served(answer) + for answer in during: + assert answer.status >= 500 and OVERLOADED_MESSAGE in answer.text, answer + markers: Final = (*healthy, *outage, *recovery) + safeguarded_markers: Final = frozenset( + marker for index, marker in enumerate(markers) if _burst_request(index, marker)[0] == "/v1/messages" + ) + delivered: Final = deliveries(rig.wire) + seen: Final = tuple(sum(sent.marker == marker for sent in delivered) for marker in markers) + assert seen == (1,) * len(seen), seen + expected: Final = { + marker: (True, DANGEROUS_TOOL_USE) if marker in safeguarded_markers else (False, None) for marker in markers + } + for sent in delivered: + assert (sent.safeguards is not None, sent.beta) == expected[sent.marker], sent