From d43a248900cc8ae70a3be9cd41782ad4d59bbbb5 Mon Sep 17 00:00:00 2001 From: yucheng Date: Fri, 25 Sep 2026 21:56:12 +0000 Subject: [PATCH] test(integration): audit cells for the Agent 365 fail-open default across entry points, Entra faults, throttling and two workers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp/test_mcp_agent_365_guardrail.py | 367 +++++++++++++++++- 1 file changed, 350 insertions(+), 17 deletions(-) diff --git a/tests/integration/mcp/test_mcp_agent_365_guardrail.py b/tests/integration/mcp/test_mcp_agent_365_guardrail.py index 53b9efb4818..fe94fd81f73 100644 --- a/tests/integration/mcp/test_mcp_agent_365_guardrail.py +++ b/tests/integration/mcp/test_mcp_agent_365_guardrail.py @@ -1,35 +1,78 @@ """Agent 365 guardrail when its own dependencies fail: Entra and Agent 365 are owned local doubles.""" import json +import os +import signal +import time import uuid from collections.abc import Iterator -from contextlib import contextmanager -from dataclasses import dataclass +from concurrent.futures import ThreadPoolExecutor +from contextlib import ExitStack, contextmanager +from dataclasses import dataclass, replace from hashlib import sha256 from pathlib import Path from typing import Final +from urllib.parse import parse_qs +import httpx +import psutil import yaml from integration._support.client import Gateway, eventually from integration._support.database import read_rows -from integration._support.mcp import McpCaller, McpPeer, echo_tool, register_mcp, scripted_peer, tool_calls +from integration._support.mcp import ( + ENTRY_POINTS, + EntryPoint, + McpCaller, + McpPeer, + Outcome, + echo_tool, + register_mcp, + scripted_peer, + tool_calls, +) from integration._support.process import owned_proxy_process -from integration._support.wire import Reply, Request, wire_server +from integration._support.wire import Reply, Request, Wire, wire_server from prometheus_client.parser import text_string_to_metric_families TENANT: Final = "00000000-0000-4000-8000-0000000a3650" EVALUATE_PATH: Final = "/agents/tool-evaluation/evaluate" GUARDRAIL_ERRORS: Final = "litellm_guardrail_errors_total" CALLER_TOKEN: Final = "eyJhbGciOiJub25lIn0.eyJzdWIiOiJpbnRlZ3JhdGlvbiJ9.synthetic-signature" +GUARDRAIL_TIMEOUT_SECONDS: Final = 1.0 +SLOW_REPLY_SECONDS: Final = 3.0 GUARDRAIL_STATUSES: Final = ( "SELECT metadata->'mcp_tool_call_metadata'->>'name' AS tool, gi->>'guardrail_status' AS status " "FROM \"LiteLLM_SpendLogs\", jsonb_array_elements(metadata->'guardrail_information') gi " "WHERE api_key = %s AND gi->>'guardrail_provider' = 'agent_365'" ) +GUARDRAIL_ROWS: Final = ( + "SELECT metadata->'guardrail_information' AS gi FROM \"LiteLLM_SpendLogs\" " + 'WHERE api_key = %s AND call_type = %s ORDER BY "startTime"' +) +TOOLS: Final = ("outage", "skipped", "denied", "add", "throttled", "rejected", "nonjson", "nobool", "slow") + + +def _caller_token(entra_case: str) -> str: + """A compact JWS whose signature segment tells the Entra double how to answer the OBO exchange.""" + return f"eyJhbGciOiJub25lIn0.eyJzdWIiOiJpbnRlZ3JhdGlvbiJ9.{entra_case}" def _entra(request: Request) -> Reply: assert request.target == f"/{TENANT}/oauth2/v2.0/token", request.target + case: Final = parse_qs(request.body.decode())["assertion"][0].rsplit(".", 1)[-1] + match case: + case "entra-outage": + return Reply(status=503, body=json.dumps({"error": "synthetic Entra outage"}).encode()) + case "entra-nonjson": + return Reply(body=b"synthetic gateway timeout", content_type="text/html") + case "entra-slow": + time.sleep(SLOW_REPLY_SECONDS) + case "entra-misconfigured": + misconfigured: Final = {"error": "invalid_client", "error_codes": [7000215]} + return Reply(status=401, body=json.dumps(misconfigured).encode()) + case "entra-rejected": + rejected: Final = {"error": "invalid_grant", "error_codes": [50013]} + return Reply(status=400, body=json.dumps(rejected).encode()) return Reply(body=json.dumps({"access_token": "obo-" + uuid.uuid4().hex, "expires_in": 3599}).encode()) @@ -44,11 +87,32 @@ def _agent_365(request: Request) -> Reply: case "denied": verdict: Final = {"status": "Evaluated", "verdict": "Block", "message": "synthetic block"} return Reply(body=json.dumps({"allowed": False, "defender": verdict, "correlationId": "denied-1"}).encode()) - case _: - return Reply(body=json.dumps({"allowed": True, "defender": {"status": "Evaluated"}}).encode()) + case "throttled": + return Reply(status=429, body=json.dumps({"error": "synthetic throttle"}).encode()) + case "rejected": + return Reply(status=400, body=json.dumps({"error": "synthetic malformed evaluation request"}).encode()) + case "nonjson": + return Reply(body=b"synthetic upstream error page", content_type="text/html") + case "nobool": + return Reply(body=json.dumps({"allowed": "yes", "defender": {"status": "Evaluated"}}).encode()) + case "slow": + time.sleep(SLOW_REPLY_SECONDS) + return Reply(body=json.dumps({"allowed": True, "defender": {"status": "Evaluated"}}).encode()) -def _config(tmp_path: Path, name: str, entra_url: str, agent_365_url: str, fallback: str | None) -> Path: +def _generic_guardrail_outage(request: Request) -> Reply: + assert request.target == "/beta/litellm_basic_guardrail_api", request.target + return Reply(status=503, body=json.dumps({"error": "synthetic sibling guardrail outage"}).encode()) + + +def _config( + tmp_path: Path, + name: str, + entra_url: str | None, + agent_365_url: str, + fallback: str | None, + sibling_url: str | None, +) -> Path: config: dict = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) config["litellm_settings"]["callbacks"] = ["prometheus"] config["guardrails"] = [ @@ -62,10 +126,26 @@ def _config(tmp_path: Path, name: str, entra_url: str, agent_365_url: str, fallb "client_id": "synthetic-client-id", "client_secret": "synthetic-client-secret", "api_base": agent_365_url, - "authority_host": entra_url, + "timeout": GUARDRAIL_TIMEOUT_SECONDS, + **({"authority_host": entra_url} if entra_url else {}), **({"unreachable_fallback": fallback} if fallback else {}), }, - } + }, + *( + [ + { + "guardrail_name": f"{name}-sibling", + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": True, + "api_base": sibling_url, + }, + } + ] + if sibling_url + else [] + ), ] path: Final = tmp_path / "agent_365.yaml" path.write_text(yaml.safe_dump(config)) @@ -80,30 +160,86 @@ class Rig: alias: str guardrail_name: str peer: McpPeer + entra: Wire + agent_365: Wire + process_id: int + server_id: str + + def caller_for(self, entry: EntryPoint, token: str = CALLER_TOKEN) -> McpCaller: + return McpCaller(self.candidate, self.key, entry, self.alias, headers={"Authorization": f"Bearer {token}"}) + + def every_worker_serves_the_catalog(self, samples: int = 8) -> bool: + """Concurrent fresh connections spread across the accepting workers, so ``samples`` of them cover both.""" + headers: Final = {"Authorization": f"Bearer {CALLER_TOKEN}"} + with ExitStack() as connections, ThreadPoolExecutor(max_workers=samples) as pool: + callers: Final = tuple( + McpCaller( + replace(self.candidate, client=connections.enter_context(httpx.Client(base_url=base))), + self.key, + "mcp", + self.alias, + headers=headers, + ) + for base in (self.candidate.client.base_url,) * samples + ) + return all(f"{self.alias}-add" in listed.tools for listed in pool.map(McpCaller.list_tools, callers)) + + def call_without_bearer(self, tool: str) -> Outcome: + return McpCaller(self.candidate, self.key, "mcp", self.alias).call(f"{self.alias}-{tool}", {"a": 0}) + + def upstream_tool_names(self) -> tuple[str, ...]: + return tuple(str(call["body"]["params"]["name"]) for call in tool_calls(self.peer.drain())) @contextmanager -def _rig(gateway: Gateway, tmp_path: Path, fallback: str | None) -> Iterator[Rig]: +def _rig( + gateway: Gateway, + tmp_path: Path, + fallback: str | None, + *, + workers: int = 1, + environment: dict[str, str] | None = None, + authority_env: str | None = None, + sibling: bool = False, +) -> Iterator[Rig]: + """``authority_env`` names the environment variable that carries the Entra double instead of the config.""" alias: Final = "a365" + uuid.uuid4().hex[:8] prom_dir: Final = tmp_path / "prom" - prom_dir.mkdir() + prom_dir.mkdir(parents=True) with ( wire_server(_entra) as entra, wire_server(_agent_365) as agent_365, - scripted_peer(echo_tool("outage"), echo_tool("skipped"), echo_tool("denied"), echo_tool("add")) as peer, + wire_server(_generic_guardrail_outage) as sibling_outage, + scripted_peer(*(echo_tool(tool) for tool in TOOLS)) as peer, owned_proxy_process( gateway, tmp_path, - {"PROMETHEUS_MULTIPROC_DIR": str(prom_dir)}, - config=_config(tmp_path, alias, entra.url, agent_365.url, fallback), + { + "PROMETHEUS_MULTIPROC_DIR": str(prom_dir), + "PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": "2", + **(environment or {}), + **({authority_env: entra.url} if authority_env else {}), + }, + config=_config( + tmp_path, + alias, + None if authority_env else entra.url, + agent_365.url, + fallback, + sibling_outage.url if sibling else None, + ), + workers=workers, ) as owned, owned.gateway.scenario() as scenario, ): identity: Final = register_mcp(scenario, peer, alias) key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) caller: Final = McpCaller(owned.gateway, key, "mcp", alias, headers={"Authorization": f"Bearer {CALLER_TOKEN}"}) + rig: Final = Rig(owned.gateway, caller, key, alias, alias, peer, entra, agent_365, owned.process.pid, identity) + if workers > 1: + eventually(rig.every_worker_serves_the_catalog, lambda served: served, seconds=40) peer.drain() - yield Rig(owned.gateway, caller, key, alias, alias, peer) + yield rig def _guardrail_statuses(key: str) -> dict[str, str]: @@ -124,6 +260,12 @@ def _guardrail_error_counts(candidate: Gateway, guardrail_name: str) -> dict[str } +def _chat(rig: Rig, model: str, marker: str) -> httpx.Response: + return rig.candidate.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": marker}]}, key=rig.key + ) + + def test_default_lets_the_call_through_unscanned_and_counts_it_when_agent_365_cannot_evaluate( gateway: Gateway, tmp_path: Path ) -> None: @@ -134,7 +276,7 @@ def test_default_lets_the_call_through_unscanned_and_counts_it_when_agent_365_ca assert skipped.text == '{"a": 2}', f"Defender skipping the call must fail open by default: {skipped.raw}" denied: Final = rig.caller.call(f"{rig.alias}-denied", {"a": 3}) assert denied.error is not None and "Blocked by Microsoft Defender" in denied.raw, denied.raw - assert tuple(call["body"]["params"]["name"] for call in tool_calls(rig.peer.drain())) == ("outage", "skipped") + assert rig.upstream_tool_names() == ("outage", "skipped") statuses: Final = eventually(lambda: _guardrail_statuses(rig.key), lambda seen: len(seen) >= 3, seconds=70) assert statuses == { "outage": "guardrail_failed_to_respond", @@ -154,7 +296,198 @@ def test_explicit_fail_closed_blocks_with_503_and_never_reaches_upstream_when_ag outage: Final = rig.caller.call(f"{rig.alias}-outage", {"a": 1}) assert outage.error is not None and "could not authorize the tool call" in outage.raw, outage.raw assert rig.caller.call(f"{rig.alias}-add", {"a": 2}).text == '{"a": 2}' - assert tuple(call["body"]["params"]["name"] for call in tool_calls(rig.peer.drain())) == ("add",) + assert rig.upstream_tool_names() == ("add",) statuses: Final = eventually(lambda: _guardrail_statuses(rig.key), lambda seen: len(seen) >= 2, seconds=70) assert statuses == {"outage": "guardrail_failed_to_respond", "add": "success"}, statuses assert _guardrail_error_counts(rig.candidate, rig.guardrail_name) == {"HTTPException": 1.0} + + +def test_explicit_fail_open_matches_the_default_and_still_blocks_policy_denials( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path, fallback="fail_open") as rig: + assert rig.caller.call(f"{rig.alias}-outage", {"a": 1}).text == '{"a": 1}' + denied: Final = rig.caller.call(f"{rig.alias}-denied", {"a": 2}) + assert denied.error is not None and "Blocked by Microsoft Defender" in denied.raw, denied.raw + assert rig.upstream_tool_names() == ("outage",) + statuses: Final = eventually(lambda: _guardrail_statuses(rig.key), lambda seen: len(seen) >= 2, seconds=70) + assert statuses == {"outage": "guardrail_failed_to_respond", "denied": "guardrail_intervened"}, statuses + assert _guardrail_error_counts(rig.candidate, rig.guardrail_name) == {"fail_open": 1.0, "HTTPException": 1.0} + + +def test_default_fails_open_on_malformed_or_stalled_agent_365_replies(gateway: Gateway, tmp_path: Path) -> None: + with _rig(gateway, tmp_path, fallback=None) as rig: + for index, tool in enumerate(("nonjson", "nobool", "slow")): + outcome: Final = rig.caller.call(f"{rig.alias}-{tool}", {"a": index}) + assert outcome.text == json.dumps({"a": index}), f"{tool}: {outcome.raw}" + assert rig.upstream_tool_names() == ("nonjson", "nobool", "slow") + statuses: Final = eventually(lambda: _guardrail_statuses(rig.key), lambda seen: len(seen) >= 3, seconds=70) + assert statuses == dict.fromkeys(("nonjson", "nobool", "slow"), "guardrail_failed_to_respond"), statuses + assert _guardrail_error_counts(rig.candidate, rig.guardrail_name) == {"fail_open": 3.0} + + +def test_default_fails_open_when_entra_is_down_stalled_malformed_or_refuses_the_gateway_credentials( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path, fallback=None) as rig: + cases: Final = ("entra-outage", "entra-slow", "entra-nonjson", "entra-misconfigured") + for index, case in enumerate(cases): + outcome: Final = rig.caller_for("mcp", _caller_token(case)).call(f"{rig.alias}-add", {"a": index}) + assert outcome.text == json.dumps({"a": index}), f"{case}: {outcome.raw}" + assert rig.upstream_tool_names() == ("add",) * len(cases) + assert tuple( + sorted(parse_qs(request.body.decode())["assertion"][0].rsplit(".", 1)[-1] for request in rig.entra.drain()) + ) == tuple(sorted(cases)) + assert rig.agent_365.drain() == (), "no OBO token means no evaluation request" + rows: Final = eventually( + lambda: read_rows(GUARDRAIL_ROWS, (sha256(rig.key.encode()).hexdigest(), "call_mcp_tool")), + lambda seen: len(seen) >= len(cases), + seconds=70, + ) + assert [row["gi"][0]["guardrail_status"] for row in rows] == ["guardrail_failed_to_respond"] * len(cases) + assert _guardrail_error_counts(rig.candidate, rig.guardrail_name) == {"fail_open": float(len(cases))} + + +def test_throttling_and_ordinary_4xx_from_agent_365_keep_blocking_under_the_default( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path, fallback=None) as rig: + throttled: Final = rig.caller.call(f"{rig.alias}-throttled", {"a": 1}) + assert throttled.error == "Error: Agent 365 guardrail could not authorize the tool call", throttled.raw + rejected: Final = rig.caller.call(f"{rig.alias}-rejected", {"a": 2}) + assert rejected.error == "Error: Agent 365 rejected the tool evaluation request", rejected.raw + assert rig.upstream_tool_names() == () + statuses: Final = eventually(lambda: _guardrail_statuses(rig.key), lambda seen: len(seen) >= 2, seconds=70) + assert statuses == {"throttled": "guardrail_failed_to_respond", "rejected": "guardrail_intervened"}, statuses + assert _guardrail_error_counts(rig.candidate, rig.guardrail_name) == {"HTTPException": 2.0} + + +def test_caller_authentication_failures_keep_blocking_under_the_default(gateway: Gateway, tmp_path: Path) -> None: + with _rig(gateway, tmp_path, fallback=None) as rig: + rejected: Final = "Error: Agent 365 guardrail rejected the tool call" + missing: Final = rig.call_without_bearer("add") + assert missing.error == rejected, missing.raw + malformed: Final = rig.caller_for("mcp", "not-a-jws").call(f"{rig.alias}-add", {"a": 1}) + assert malformed.error == rejected, malformed.raw + refused: Final = rig.caller_for("mcp", _caller_token("entra-rejected")).call(f"{rig.alias}-add", {"a": 2}) + assert refused.error == rejected, refused.raw + assert len(rig.entra.drain()) == 1, "only the well-formed bearer reaches the OBO exchange" + assert rig.upstream_tool_names() == () + assert rig.agent_365.drain() == () + rows: Final = eventually( + lambda: read_rows(GUARDRAIL_ROWS, (sha256(rig.key.encode()).hexdigest(), "call_mcp_tool")), + lambda seen: len(seen) >= 3, + seconds=70, + ) + assert [row["gi"][0]["guardrail_status"] for row in rows] == ["guardrail_intervened"] * 3 + assert _guardrail_error_counts(rig.candidate, rig.guardrail_name) == {"HTTPException": 3.0} + + +def test_every_mcp_entry_point_fails_open_on_outage_and_blocks_denials(gateway: Gateway, tmp_path: Path) -> None: + with _rig(gateway, tmp_path, fallback=None) as rig: + for entry in ENTRY_POINTS: + caller: Final = rig.caller_for(entry) + passed: Final = caller.call(f"{rig.alias}-outage", {"entry": entry}, server_id=rig.server_id) + assert passed.text == json.dumps({"entry": entry}), f"{entry}: {passed.raw}" + denied: Final = caller.call(f"{rig.alias}-denied", {"entry": entry}, server_id=rig.server_id) + assert denied.error is not None and "Blocked by Microsoft Defender" in denied.raw, f"{entry}: {denied.raw}" + assert rig.upstream_tool_names() == ("outage",) * len(ENTRY_POINTS) + counts: Final = eventually( + lambda: _guardrail_error_counts(rig.candidate, rig.guardrail_name), + lambda seen: sum(seen.values()) >= 2 * len(ENTRY_POINTS), + ) + assert counts == {"fail_open": float(len(ENTRY_POINTS)), "HTTPException": float(len(ENTRY_POINTS))}, counts + + +def test_chat_completions_never_touch_agent_365_while_a_sibling_guardrail_keeps_its_fail_closed_default( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path, fallback=None, sibling=True) as rig, rig.candidate.scenario() as scenario: + model: Final = scenario.model() + chat: Final = _chat(rig, model, "sibling-" + uuid.uuid4().hex) + assert chat.status_code == 500 and "Generic Guardrail API failed" in chat.text, chat.text + assert rig.entra.drain() == () and rig.agent_365.drain() == () + assert rig.caller.call(f"{rig.alias}-outage", {"a": 1}).text == '{"a": 1}' + assert _guardrail_error_counts(rig.candidate, rig.guardrail_name) == {"fail_open": 1.0} + + +def test_chat_completions_are_unaffected_by_the_mcp_guardrail(gateway: Gateway, tmp_path: Path) -> None: + with _rig(gateway, tmp_path, fallback=None) as rig, rig.candidate.scenario() as scenario: + model: Final = scenario.model() + marker: Final = "unaffected-" + uuid.uuid4().hex + chat: Final = _chat(rig, model, marker) + assert chat.status_code == 200, chat.text + assert rig.entra.drain() == () and rig.agent_365.drain() == () + rows: Final = eventually( + lambda: read_rows(GUARDRAIL_ROWS, (sha256(rig.key.encode()).hexdigest(), "acompletion")), + lambda seen: len(seen) >= 1, + seconds=70, + ) + assert rows[0]["gi"] is None, rows + assert _guardrail_error_counts(rig.candidate, rig.guardrail_name) == {} + + +def test_agent365_authority_host_env_wins_over_azure_authority_host_when_config_has_none( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path / "azure", fallback=None, authority_env="AZURE_AUTHORITY_HOST") as azure_only: + assert azure_only.caller.call(f"{azure_only.alias}-add", {"a": 1}).text == '{"a": 1}' + assert len(azure_only.entra.drain()) == 1 + with ( + wire_server(_entra) as decoy, + _rig( + gateway, + tmp_path / "agent365", + fallback=None, + authority_env="AGENT365_AUTHORITY_HOST", + environment={"AZURE_AUTHORITY_HOST": decoy.url}, + ) as both, + ): + assert both.caller.call(f"{both.alias}-add", {"a": 2}).text == '{"a": 2}' + assert len(both.entra.drain()) == 1 and decoy.drain() == () + + +def test_thirty_call_burst_against_a_flapping_agent_365_reaches_upstream_exactly_once_each_on_two_workers( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path, fallback=None, workers=2) as rig: + markers: Final = tuple(("outage" if index % 2 else "add", uuid.uuid4().hex) for index in range(30)) + with ThreadPoolExecutor(max_workers=10) as pool: + outcomes: Final = tuple( + pool.map(lambda pair: rig.caller.call(f"{rig.alias}-{pair[0]}", {"marker": pair[1]}), markers) + ) + for (tool, marker), outcome in zip(markers, outcomes, strict=True): + assert outcome.text == json.dumps({"marker": marker}), f"{tool}: {outcome.raw}" + seen: Final = sorted( + str(call["body"]["params"]["arguments"]["marker"]) for call in tool_calls(rig.peer.drain()) + ) + assert seen == sorted(marker for _, marker in markers) + assert rig.caller.call(f"{rig.alias}-denied", {"a": 1}).error is not None + counts: Final = eventually( + lambda: _guardrail_error_counts(rig.candidate, rig.guardrail_name), lambda seen: len(seen) >= 2 + ) + assert counts == {"fail_open": 15.0, "HTTPException": 1.0}, counts + + +def test_default_survives_a_worker_kill_and_keeps_blocking_denials_on_two_workers( + gateway: Gateway, tmp_path: Path +) -> None: + with _rig(gateway, tmp_path, fallback=None, workers=2) as rig: + workers: Final = eventually( + lambda: psutil.Process(rig.process_id).children(recursive=True), + lambda children: len(children) >= 2, + seconds=30, + ) + os.kill(workers[0].pid, signal.SIGKILL) + eventually( + lambda: tuple(rig.caller.list_tools().tools for _ in range(4)), + lambda seen: all(f"{rig.alias}-add" in tools for tools in seen), + seconds=40, + ) + rig.peer.drain() + for index in range(10): + passed: Final = rig.caller.call(f"{rig.alias}-outage", {"a": index}) + assert passed.text == json.dumps({"a": index}), passed.raw + denied: Final = rig.caller.call(f"{rig.alias}-denied", {"a": index}) + assert denied.error is not None and "Blocked by Microsoft Defender" in denied.raw, denied.raw + assert rig.upstream_tool_names() == ("outage",) * 10