mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-18 04:44:37 +00:00
fix(rust): add the dangerous-tool-use beta when an anthropic /v1/messages request carries safeguards (#46102)
* fix(rust): add the dangerous-tool-use beta when an anthropic /v1/messages request carries safeguards * test(rust): check the dangerous-tool-use beta reaches the wire with safeguards * 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. * fix(rust): add the safeguards beta on Azure AI and Claude Platform hosts Azure AI and Claude Platform Messages transforms now add the dangerous-tool-use beta for a non-null safeguards field, through the Azure AI and Anthropic beta-map columns. The safeguards wire and chaos tests get unique file names so they do not collide with the Python fix's tests, the Python-route test asserts the fixed behavior behind a BUG skip, and the DeepSeek test checks that no beta is added * fix(rust): keep Claude Platform out of the safeguards beta until a route can prove it Claude Platform has no beta-map column and no QA account, so it goes back to sending safeguards with no derived beta, like DeepSeek. Also release held chaos-test requests before the thread pool waits on them, and mark the no-column test as rstest. * fix(rust): derive the safeguards beta for Claude Platform and keep the null-cell test off Databricks Claude Platform on AWS answers 400 to safeguards without dangerous-tool-use-2026-09-03 and 200 with it, so its Messages transform derives the beta through the Anthropic column again. The null-cell case now uses only Bedrock Converse, since Databricks support for the beta is changing in a sibling PR. * fix(rust): drop a null safeguards key instead of sending it without its beta Anthropic rejects any safeguards key, null included, when the dangerous-tool-use beta is missing, so an explicit "safeguards": null on the Rust gateway still answered 400. The shared Messages transform now drops a null safeguards key on every host, as the proxy already drops None params before calling Rust, so it answers 200 with no beta. The scripted peer in the integration wire tests now rejects the key on presence, matching the real API. * test: cover Bedrock Mantle in the null safeguards case and match the vendor's non-list wording The Rust null-safeguards regression test now also runs for bedrock_mantle, whose beta is keyed on the field being present. The scripted Anthropic peer answers a non-list safeguards value with the real API's "valid array" wording. --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
eb5295b20c
commit
5b042c0db6
9 changed files with 1062 additions and 26 deletions
|
|
@ -343,6 +343,77 @@ 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 null_safeguards_are_dropped_instead_of_sent_without_their_beta(
|
||||
call: MessagesCall,
|
||||
#[values("anthropic", "azure_ai", "bedrock_mantle", "deepseek", "minimax")] provider: &str,
|
||||
) {
|
||||
let upstream = upstream([message_response()]).await;
|
||||
|
||||
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": null}),
|
||||
))
|
||||
.await;
|
||||
|
||||
let request = only_request(&upstream).await;
|
||||
assert_eq!(
|
||||
sent_betas(&request),
|
||||
BetaSet::from_iter([AnthropicBeta::Other("caller-beta-2025-01-01".to_string())])
|
||||
);
|
||||
assert!(
|
||||
!request
|
||||
.body_json::<Map<String, Value>>()
|
||||
.unwrap()
|
||||
.contains_key("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),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -136,12 +139,20 @@ pub(crate) fn transform_messages_request(
|
|||
messages: strip_encrypted_reasoning_blocks(messages),
|
||||
params: MessagesOptionalParams {
|
||||
context_management,
|
||||
extra: without_null_safeguards(request.params.extra),
|
||||
..request.params
|
||||
},
|
||||
..request
|
||||
})
|
||||
}
|
||||
|
||||
fn without_null_safeguards(extra: Map<String, Value>) -> Map<String, Value> {
|
||||
extra
|
||||
.into_iter()
|
||||
.filter(|(key, value)| key != "safeguards" || !value.is_null())
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub(crate) fn update_headers_with_anthropic_beta(
|
||||
headers: Headers,
|
||||
request: &MessagesRequest,
|
||||
|
|
@ -149,12 +160,34 @@ 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.
|
||||
pub(crate) fn update_headers_with_safeguards_beta(
|
||||
headers: Headers,
|
||||
request: &MessagesRequest,
|
||||
provider: BetaProvider,
|
||||
) -> Headers {
|
||||
merge_beta_headers(
|
||||
headers,
|
||||
feature_betas(request).union(on_provider(mapped_host_betas(&request.params), provider)),
|
||||
)
|
||||
}
|
||||
|
||||
/// The betas a request needs on a host with a column in the beta map, filtered and renamed
|
||||
/// through that column.
|
||||
pub(crate) fn provider_feature_betas(request: &MessagesRequest, provider: BetaProvider) -> BetaSet {
|
||||
feature_betas(request)
|
||||
.iter()
|
||||
.filter_map(|beta| beta.on(provider))
|
||||
on_provider(
|
||||
feature_betas(request).union(mapped_host_betas(&request.params)),
|
||||
provider,
|
||||
)
|
||||
}
|
||||
|
||||
fn on_provider(betas: BetaSet, provider: BetaProvider) -> BetaSet {
|
||||
betas.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()
|
||||
}
|
||||
|
||||
|
|
@ -206,6 +239,13 @@ fn uses_structured_output(params: &MessagesOptionalParams) -> bool {
|
|||
.is_some_and(|config| config.format.is_some())
|
||||
}
|
||||
|
||||
fn carries_safeguards(params: &MessagesOptionalParams) -> bool {
|
||||
params
|
||||
.extra
|
||||
.get("safeguards")
|
||||
.is_some_and(|value| !value.is_null())
|
||||
}
|
||||
|
||||
fn messages_carry_output_config(messages: &[Message]) -> bool {
|
||||
messages
|
||||
.iter()
|
||||
|
|
@ -1025,10 +1065,44 @@ mod tests {
|
|||
json!({"messages": [{"role": "user", "content": "hi", "output_config": null}]}),
|
||||
&["per-turn-control-2026-07-01"]
|
||||
)]
|
||||
#[case::safeguards(
|
||||
json!({"safeguards": [{"type": "dangerous_tool_use", "classifier_context": {"v": 1, "permission_mode": "auto"}}]}),
|
||||
&[]
|
||||
)]
|
||||
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"}]), &[])]
|
||||
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)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
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")])]
|
||||
|
|
@ -1067,6 +1141,11 @@ 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_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")],
|
||||
)]
|
||||
fn request_headers_merge_the_feature_betas(
|
||||
#[case] input: &[(&str, &str)],
|
||||
#[case] fields: Value,
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ use litellm_auth::SecretValue;
|
|||
use litellm_http::request::{has_bearer_auth, has_header};
|
||||
use litellm_llms_types::{
|
||||
formats::messages::MessagesRequest,
|
||||
providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH},
|
||||
providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH},
|
||||
};
|
||||
use litellm_router_types::LitellmParams;
|
||||
|
||||
|
|
@ -12,7 +12,7 @@ use crate::{
|
|||
common_utils::API_KEY_PLACEMENT,
|
||||
messages::{
|
||||
handler::shape_anthropic_messages_request,
|
||||
transformation::{transform_messages_request, update_headers_with_anthropic_beta},
|
||||
transformation::{transform_messages_request, update_headers_with_safeguards_beta},
|
||||
},
|
||||
},
|
||||
azure_ai::common_utils::{
|
||||
|
|
@ -104,7 +104,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig {
|
|||
}
|
||||
|
||||
fn request_headers(&self, headers: Headers, request: &MessagesRequest) -> Headers {
|
||||
update_headers_with_anthropic_beta(headers, request)
|
||||
update_headers_with_safeguards_beta(headers, request, BetaProvider::AzureAi)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -544,6 +544,12 @@ mod tests {
|
|||
&[("x-api-key", "k"), ("anthropic-beta", "structured-outputs-2025-11-13,web-search-2025-03-05")]
|
||||
)]
|
||||
#[case::no_feature_needs_a_beta(json!({}), &[], &[("x-api-key", "k")])]
|
||||
#[case::safeguards_need_the_dangerous_tool_use_beta(
|
||||
json!({"safeguards": [{"type": "dangerous_tool_use"}]}),
|
||||
&[("anthropic-beta", "claude-code-20250219")],
|
||||
&[("x-api-key", "k"), ("anthropic-beta", "claude-code-20250219,dangerous-tool-use-2026-09-03")]
|
||||
)]
|
||||
#[case::null_safeguards_need_no_beta(json!({"safeguards": null}), &[], &[("x-api-key", "k")])]
|
||||
fn request_headers_carry_the_anthropic_feature_betas(
|
||||
#[case] fields: serde_json::Value,
|
||||
#[case] forwarded: &[(&str, &str)],
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ use litellm_core_utils::settings::resolve_non_empty;
|
|||
use litellm_http::request::{header_value, with_header, without_headers};
|
||||
use litellm_llms_types::{
|
||||
formats::messages::MessagesRequest,
|
||||
providers::anthropic::{DEFAULT_HEADERS, MESSAGES_PATH},
|
||||
providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH},
|
||||
};
|
||||
use litellm_router_types::LitellmParams;
|
||||
use serde_json::Value;
|
||||
|
|
@ -15,7 +15,7 @@ use crate::{
|
|||
Error, ErrorDetail,
|
||||
anthropic::messages::{
|
||||
handler::shape_anthropic_messages_request,
|
||||
transformation::{transform_messages_request, update_headers_with_anthropic_beta},
|
||||
transformation::{transform_messages_request, update_headers_with_safeguards_beta},
|
||||
},
|
||||
base_llm::{
|
||||
auth::{AuthScheme, Headers, ValidatedEnvironment},
|
||||
|
|
@ -174,7 +174,7 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig {
|
|||
Some(workspace) => with_header(headers, WORKSPACE_HEADER, workspace.into()),
|
||||
None => headers,
|
||||
};
|
||||
update_headers_with_anthropic_beta(headers, request)
|
||||
update_headers_with_safeguards_beta(headers, request, BetaProvider::Anthropic)
|
||||
}
|
||||
|
||||
fn wire_body(&self, body: Value) -> Value {
|
||||
|
|
@ -356,6 +356,21 @@ mod tests {
|
|||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::safeguards(json!([{"type": "dangerous_tool_use"}]), Some("dangerous-tool-use-2026-09-03"))]
|
||||
#[case::null_safeguards(Value::Null, None)]
|
||||
fn safeguards_need_the_dangerous_tool_use_beta(
|
||||
#[case] safeguards: Value,
|
||||
#[case] expected: Option<&str>,
|
||||
) {
|
||||
let request = serde_json::from_value(json!({"model": "claude-test", "max_tokens": 16, "messages": [], "safeguards": safeguards})).unwrap();
|
||||
let headers = BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG.request_headers(vec![], &request);
|
||||
assert_eq!(
|
||||
header_value(&headers, "anthropic-beta").map(|value| value.to_string()),
|
||||
expected.map(str::to_string)
|
||||
);
|
||||
}
|
||||
|
||||
#[rstest]
|
||||
#[case::default(json!({}), false)]
|
||||
#[case::override_empty(json!({"claude_platform_unsupported_params": []}), true)]
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -213,23 +213,14 @@ fn has_web_search_tool(tools: Option<&[Recognized<MessagesTool>]>) -> bool {
|
|||
}
|
||||
|
||||
/// Python's Vertex `validate_anthropic_messages_environment` beta set: the shared feature
|
||||
/// betas under Vertex's policy, plus the web search and safeguards betas Vertex needs spelled
|
||||
/// out.
|
||||
/// betas under Vertex's policy, plus the web search beta Vertex needs spelled out.
|
||||
fn vertex_feature_betas(request: &MessagesRequest) -> BetaSet {
|
||||
let params = &request.params;
|
||||
provider_feature_betas(request, BetaProvider::VertexAi).union(
|
||||
[
|
||||
has_web_search_tool(params.tools.as_deref())
|
||||
.then_some(AnthropicBeta::WebSearch20250305),
|
||||
params
|
||||
.extra
|
||||
.get("safeguards")
|
||||
.is_some_and(|value| !value.is_null())
|
||||
.then_some(AnthropicBeta::DangerousToolUse20260903),
|
||||
]
|
||||
.into_iter()
|
||||
.flatten()
|
||||
.collect(),
|
||||
has_web_search_tool(params.tools.as_deref())
|
||||
.then_some(AnthropicBeta::WebSearch20250305)
|
||||
.into_iter()
|
||||
.collect(),
|
||||
)
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,285 @@
|
|||
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 array"
|
||||
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())
|
||||
body: Final = body_of(request)
|
||||
safeguards: Final = body.get("safeguards")
|
||||
if "safeguards" in body 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)))
|
||||
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)))
|
||||
stack.callback(release.set)
|
||||
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)))
|
||||
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)))
|
||||
stack.callback(release.set)
|
||||
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,408 @@
|
|||
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,
|
||||
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_without_adding_a_beta(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) - {DANGEROUS_TOOL_USE}, sent
|
||||
|
||||
|
||||
def test_python_route_adds_the_beta_for_safeguards(gateway: Gateway, scenario: Scenario) -> None:
|
||||
pytest.skip("BUG: Python /v1/messages sends safeguards without the dangerous-tool-use beta, Anthropic answers 400")
|
||||
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 == 200, response.text
|
||||
assert "safeguard_results" in response.text, response.text
|
||||
(sent,) = deliveries(wire)
|
||||
assert (sent.marker, sent.safeguards) == (marker, list(SAFEGUARDS)), sent
|
||||
assert DANGEROUS_TOOL_USE in _beta_set(sent.beta), 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