import json import uuid from pathlib import Path from typing import Final import pytest import yaml from integration._support.client import Gateway, eventually from integration._support.database import read_rows from integration._support.mcp import mcp_peer, register_mcp, tool_names from integration._support.process import owned_proxy from integration._support.wire import Reply, Request, wire_server @pytest.mark.covers("other.observability.guardrails.rewrite_reaches_correct_anthropic_positions") def test_guardrail_rewrites_system_and_user_in_actual_anthropic_request(gateway: Gateway, tmp_path: Path) -> None: identity: Final = "guardrail" + uuid.uuid4().hex originals: Final = ["synthetic private system", "synthetic private user", "unchanged sibling"] replacements: Final = ["permitted system", "permitted user", "unchanged sibling"] def guardrail(request: Request) -> Reply: assert request.target == "/beta/litellm_basic_guardrail_api" body: Final = json.loads(request.body) assert body["texts"] == originals return Reply(body=json.dumps({"action": "GUARDRAIL_INTERVENED", "texts": replacements}).encode()) def provider(request: Request) -> Reply: assert request.target == "/v1/messages" body: Final = json.loads(request.body) assert body["system"] == [{"type": "text", "text": replacements[0]}] assert body["messages"] == [ { "role": "user", "content": [{"type": "text", "text": replacements[1]}, {"type": "text", "text": replacements[2]}], } ] assert all(text.encode() not in request.body for text in originals[:2]) return Reply( body=json.dumps( { "id": identity, "type": "message", "role": "assistant", "model": "claude-sonnet-4-5-20250929", "content": [{"type": "text", "text": "permitted response"}], "stop_reason": "end_turn", "stop_sequence": None, "usage": {"input_tokens": 11, "output_tokens": 4}, } ).encode() ) with wire_server(guardrail) as policy, wire_server(provider) as upstream: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["guardrails"] = [ { "guardrail_name": identity, "litellm_params": { "guardrail": "generic_guardrail_api", "mode": "pre_call", "default_on": True, "api_base": policy.url, "api_key": "synthetic-guardrail-key", }, } ] path: Final = tmp_path / "rewrite.yaml" path.write_text(yaml.safe_dump(config)) with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario: model: Final = scenario.model( model="anthropic/claude-sonnet-4-5-20250929", api_base=upstream.url, api_key="synthetic-anthropic-key" ) response: Final = candidate.request( "POST", "/v1/chat/completions", { "model": model, "max_tokens": 16, "messages": [ {"role": "system", "content": originals[0]}, {"role": "user", "content": [{"type": "text", "text": text} for text in originals[1:]]}, ], }, ) assert response.status_code == 200, response.text assert response.json()["choices"][0]["message"]["content"] == "permitted response" assert response.json()["choices"][0]["finish_reason"] == "stop" assert response.json()["usage"]["total_tokens"] == 15 assert len(policy.drain()) == len(upstream.drain()) == 1 @pytest.mark.covers("other.observability.guardrails.denial_prevents_provider_with_allowed_control") def test_guardrail_denial_prevents_provider_and_preserves_allowed_control(gateway: Gateway, tmp_path: Path) -> None: identity: Final = "guardrail" + uuid.uuid4().hex def guardrail(request: Request) -> Reply: assert request.target == "/beta/litellm_basic_guardrail_api" body: Final = json.loads(request.body) assert body["texts"] in (["synthetic denied marker"], ["synthetic allowed marker"]) result: Final = ( {"action": "BLOCKED", "blocked_reason": "synthetic policy denial"} if body["texts"] == ["synthetic denied marker"] else {"action": "NONE"} ) return Reply(body=json.dumps(result).encode()) with wire_server(guardrail) as policy: config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["guardrails"] = [ { "guardrail_name": identity, "litellm_params": { "guardrail": "generic_guardrail_api", "mode": "pre_call", "default_on": True, "api_base": policy.url, "api_key": "synthetic-guardrail-key", }, } ] path: Final = tmp_path / "deny.yaml" path.write_text(yaml.safe_dump(config)) with owned_proxy(gateway, tmp_path, {}, config=path) as candidate, candidate.scenario() as scenario: model: Final = scenario.model() key: Final = scenario.key(models=[model]) import httpx with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as observed: observed.get("/__observations") denied: Final = candidate.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "synthetic denied marker"}]}, key=key, ) assert denied.status_code == 400 and "synthetic policy denial" in denied.text, denied.text assert observed.get("/__observations").json()["requests"] == [] allowed: Final = candidate.chat(model, text="synthetic allowed marker", key=key) assert allowed["usage"]["total_tokens"] == 40 assert ( allowed["choices"][0]["message"]["content"] == "Hello! This is a mock response from the fake OpenAI endpoint." ) assert len(observed.get("/__observations").json()["requests"]) == 1 assert len(policy.drain()) == 2 @pytest.mark.covers("other.mcp.guardrails.request_selection_blocks_resolved_tool_without_execution") def test_request_selected_mcp_guardrail_blocks_direct_and_virtual_calls(gateway: Gateway, tmp_path: Path) -> None: guardrail = "mcp-policy-" + uuid.uuid4().hex config = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["guardrails"] = [ { "guardrail_name": guardrail, "litellm_params": { "guardrail": "custom_code", "mode": "pre_mcp_call", "default_on": False, "custom_code": ( "def apply_guardrail(inputs, request_data, input_type):\n" ' if inputs.get("tools", [{}])[0].get("function", {}).get("name") == "add":\n' ' return block("integration resolved add denied")\n' " return allow()\n" ), }, } ] path = tmp_path / "mcp-guardrail.yaml" path.write_text(yaml.safe_dump(config)) with ( owned_proxy(gateway, tmp_path, {}, config=path) as candidate, mcp_peer() as peer, candidate.scenario() as scenario, ): identity = register_mcp(scenario, peer, "guardrail" + uuid.uuid4().hex) permission = {"mcp_servers": [identity], "mcp_tool_search_enabled": True} key = scenario.key(object_permission=permission) key_selected = scenario.key(object_permission=permission, guardrails=[guardrail]) team = scenario.team(guardrails=[guardrail], object_permission={"mcp_servers": [identity]}) team_selected = scenario.key(team_id=team, object_permission=permission) catalog_key = scenario.key(object_permission={"mcp_servers": [identity]}) names = tool_names(candidate, catalog_key, identity) assert set(names) == {"add", "multiply", "fail"} for virtual in (False, True): for caller, selected, tool, expected in ( (key, [], "add", 8), (key, [guardrail], "add", None), (key_selected, [], "add", None), (team_selected, [], "add", None), (key, [guardrail], "multiply", 15), ): arguments = {"a": 3, "b": 5} peer.drain() response = candidate.client.post( "/mcp-rest/tools/call", headers={"x-litellm-api-key": caller}, json={ "server_id": identity, "name": "mcp_tool_call" if virtual else names[tool], "arguments": {"tool_name": names[tool], "arguments": arguments} if virtual else arguments, "guardrails": selected, }, ) calls = tuple(item for item in peer.drain() if item["body"].get("method") == "tools/call") if expected is None: assert response.status_code == 400, response.text assert "integration resolved add denied" in response.text, response.text assert calls == (), "pre-call denial must prevent upstream execution" else: assert response.status_code == 200, response.text assert response.json()["isError"] is False assert response.json()["content"][0]["text"] == str(expected), response.text assert len(calls) == 1 assert calls[0]["body"]["params"]["name"] == tool assert calls[0]["body"]["params"]["arguments"] == arguments