mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-17 04:34:25 +00:00
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.
This commit is contained in:
mateo-berri
parent
1e539895a2
commit
d3e1efcbba
6 changed files with 962 additions and 11 deletions
|
|
@ -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::<Value>().unwrap()["safeguards"],
|
||||
safeguards
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[tokio::test]
|
||||
async fn an_oauth_key_sends_the_browser_access_header_and_the_oauth_beta(call: MessagesCall) {
|
||||
|
|
|
|||
|
|
@ -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")],
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue