From 82ca7c9cf2b1a259e77a0ace77d76e8d71faf2fb Mon Sep 17 00:00:00 2001 From: Imran Siddique Date: Wed, 10 Jun 2026 16:05:03 -0700 Subject: [PATCH 1/2] fix(cli): wire bearer token, audit store, TEE anchor, and attestation timestamps into the live server run_startup() validated the bearer token, opened the SQLite audit store, and produced attestation timestamps, but cli.start() dropped all of them: - MCPServer was built without bearer_token, so every protected endpoint (/mcp, /audit/export, /catalog/exception, /sessions/*) was reachable unauthenticated in production (AUTH-001 dead in the live path). - AuditChain was built directly instead of via SessionManager, so entries never reached the SQLite store (AUDIT-001 inert) and the chain was never TEE-anchored (AUDIT-002). - CMCPProxy never received attestation_generated_at, so the staleness check could not fire (fail-open). Extract the composition into cli.build_server(ctx) and add regression tests that exercise the real entrypoint wiring, including a TestClient 401 check. Also serialise SqliteAuditStore access with a threading.Lock: the single connection is shared across async handlers and worker threads with check_same_thread=False. Co-Authored-By: Claude Fable 5 --- src/cmcp_runtime/audit/store.py | 59 +++++++++-------- src/cmcp_runtime/cli.py | 79 +++++++++++++++-------- tests/unit/test_cli_wiring.py | 109 ++++++++++++++++++++++++++++++++ 3 files changed, 195 insertions(+), 52 deletions(-) create mode 100644 tests/unit/test_cli_wiring.py diff --git a/src/cmcp_runtime/audit/store.py b/src/cmcp_runtime/audit/store.py index 79e744d8..6a3aa6b4 100644 --- a/src/cmcp_runtime/audit/store.py +++ b/src/cmcp_runtime/audit/store.py @@ -5,6 +5,7 @@ import json import logging import sqlite3 +import threading from dataclasses import asdict from pathlib import Path @@ -40,7 +41,10 @@ class SqliteAuditStore: def __init__(self, db_path: Path) -> None: self._db_path = db_path + # check_same_thread=False allows use from async handlers and worker + # threads; all access is serialised through self._lock. self._conn = sqlite3.connect(str(db_path), check_same_thread=False) + self._lock = threading.Lock() self._conn.execute("PRAGMA journal_mode=WAL") self._conn.execute("PRAGMA synchronous=FULL") self._conn.executescript(_CREATE_TABLE) @@ -49,21 +53,22 @@ def __init__(self, db_path: Path) -> None: def append(self, entry: AuditEntry) -> None: payload = json.dumps(asdict(entry), sort_keys=True, separators=(",", ":")) - self._conn.execute( - "INSERT INTO audit_entries " - "(sequence_number, session_id, entry_id, entry_type, entry_hash, prev_entry_hash, payload) " - "VALUES (?, ?, ?, ?, ?, ?, ?)", - ( - entry.sequence_number, - entry.session_id, - entry.entry_id, - entry.entry_type, - entry.entry_hash, - entry.prev_entry_hash, - payload, - ), - ) - self._conn.commit() + with self._lock: + self._conn.execute( + "INSERT INTO audit_entries " + "(sequence_number, session_id, entry_id, entry_type, entry_hash, prev_entry_hash, payload) " + "VALUES (?, ?, ?, ?, ?, ?, ?)", + ( + entry.sequence_number, + entry.session_id, + entry.entry_id, + entry.entry_type, + entry.entry_hash, + entry.prev_entry_hash, + payload, + ), + ) + self._conn.commit() def find_orphaned_sessions(self) -> list[str]: """ @@ -72,16 +77,18 @@ def find_orphaned_sessions(self) -> list[str]: These represent sessions that were open when the gateway last stopped, either due to a crash or an unclean shutdown. """ - cur = self._conn.execute( - """ - SELECT DISTINCT session_id FROM audit_entries - WHERE entry_type = 'session_start' - AND session_id NOT IN ( - SELECT session_id FROM audit_entries WHERE entry_type = 'session_end' - ) - """ - ) - return [row[0] for row in cur.fetchall()] + with self._lock: + cur = self._conn.execute( + """ + SELECT DISTINCT session_id FROM audit_entries + WHERE entry_type = 'session_start' + AND session_id NOT IN ( + SELECT session_id FROM audit_entries WHERE entry_type = 'session_end' + ) + """ + ) + return [row[0] for row in cur.fetchall()] def close(self) -> None: - self._conn.close() + with self._lock: + self._conn.close() diff --git a/src/cmcp_runtime/cli.py b/src/cmcp_runtime/cli.py index e9ac4f52..119ea085 100644 --- a/src/cmcp_runtime/cli.py +++ b/src/cmcp_runtime/cli.py @@ -2,10 +2,62 @@ from __future__ import annotations +from typing import TYPE_CHECKING + import click from cmcp_runtime import __version__ +if TYPE_CHECKING: + from cmcp_runtime.mcp.server import MCPServer + from cmcp_runtime.startup import RuntimeContext + + +def build_server(ctx: RuntimeContext) -> MCPServer: + """ + Compose the running gateway from a validated RuntimeContext. + + All components validated by run_startup() MUST be wired here — a component + that is validated but not passed through is silently inert in production + (the AUTH-001 bearer token and AUDIT-001 store were both lost this way). + """ + from cmcp_runtime.audit.trace_claim import _PROVIDER_MAP + from cmcp_runtime.mcp.proxy import CMCPProxy + from cmcp_runtime.mcp.server import MCPServer + from cmcp_runtime.policy.evaluator import PolicyEvaluator + from cmcp_runtime.session.manager import SessionManager + + # Resolve provider string to canonical platform name for Cedar context. + # Falls back to the raw provider string if not in the map (e.g. future providers). + attestation_platform = _PROVIDER_MAP.get( + ctx.attestation_report.provider, ctx.attestation_report.provider + ) + + # AUDIT-001/AUDIT-002: sessions MUST be created through SessionManager so the + # chain is backed by the durable SQLite store and TEE-anchored at creation. + session_manager = SessionManager(ctx) + session, audit_chain = session_manager.create_session() + policy_evaluator = PolicyEvaluator(bundle=ctx.policy_bundle, config=ctx.config) + proxy = CMCPProxy( + catalog=ctx.catalog, + policy_evaluator=policy_evaluator, + session=session, + audit_chain=audit_chain, + config=ctx.config, + attestation_generated_at=ctx.attestation_report.attestation_generated_at, + attestation_validity_seconds=ctx.attestation_report.attestation_validity_seconds, + attestation_platform=attestation_platform, + ) + # AUTH-001: the token validated in run_startup must reach the server, otherwise + # every protected endpoint is reachable unauthenticated. + return MCPServer( + proxy=proxy, + session_manager=session_manager, + audit_chain=audit_chain, + session=session, + bearer_token=ctx.config.bearer_token, + ) + @click.group() @click.version_option(__version__, prog_name="cmcp") @@ -19,16 +71,8 @@ def main() -> None: help="Override attestation.enforcement_mode from config") def start(config: str, enforcement: str | None) -> None: """Start the cMCP Runtime.""" - from uuid import uuid4 - import uvicorn - from cmcp_runtime.audit.chain import AuditChain - from cmcp_runtime.audit.trace_claim import _PROVIDER_MAP - from cmcp_runtime.mcp.proxy import CMCPProxy - from cmcp_runtime.mcp.server import MCPServer - from cmcp_runtime.policy.evaluator import PolicyEvaluator - from cmcp_runtime.session.state import SessionState from cmcp_runtime.startup import run_startup ctx = run_startup(config) @@ -38,24 +82,7 @@ def start(config: str, enforcement: str | None) -> None: from cmcp_runtime.config import EnforcementMode ctx.config.attestation.enforcement_mode = EnforcementMode(enforcement) - # Resolve provider string to canonical platform name for Cedar context. - # Falls back to the raw provider string if not in the map (e.g. future providers). - attestation_platform = _PROVIDER_MAP.get( - ctx.attestation_report.provider, ctx.attestation_report.provider - ) - - session = SessionState(session_id=str(uuid4())) - audit_chain = AuditChain(session_id=session.session_id) - policy_evaluator = PolicyEvaluator(bundle=ctx.policy_bundle, config=ctx.config) - proxy = CMCPProxy( - catalog=ctx.catalog, - policy_evaluator=policy_evaluator, - session=session, - audit_chain=audit_chain, - config=ctx.config, - attestation_platform=attestation_platform, - ) - server = MCPServer(proxy=proxy) + server = build_server(ctx) host, _, port_str = ctx.config.listen_addr.rpartition(":") port = int(port_str) diff --git a/tests/unit/test_cli_wiring.py b/tests/unit/test_cli_wiring.py new file mode 100644 index 00000000..38064c69 --- /dev/null +++ b/tests/unit/test_cli_wiring.py @@ -0,0 +1,109 @@ +""" +Regression tests for cli.build_server() — the production composition path. + +These exist because the previous cli.start() body constructed MCPServer without +the bearer token (AUTH-001 dead in production), built AuditChain without the +SQLite store and TEE anchor (AUDIT-001/AUDIT-002 inert), and never passed +attestation timestamps to the proxy (staleness check dead). Unit tests that +construct MCPServer directly cannot catch wiring gaps in the entrypoint. +""" + +from __future__ import annotations + +import sqlite3 +from datetime import UTC, datetime +from unittest.mock import MagicMock + +import pytest +from starlette.testclient import TestClient + +from cmcp_runtime.audit.store import SqliteAuditStore +from cmcp_runtime.cli import build_server +from cmcp_runtime.config import AttestationConfig, Config +from cmcp_runtime.policy.bundle import PolicyStore +from cmcp_runtime.startup import RuntimeContext + +BEARER = "test-secret-token" + + +@pytest.fixture +def ctx(tmp_path) -> RuntimeContext: + config = Config( + attestation=AttestationConfig(), + bearer_token=BEARER, + dev_mode=True, + ) + + attestation_report = MagicMock() + attestation_report.provider = "software-only" + attestation_report.attestation_generated_at = datetime.now(UTC) + attestation_report.attestation_validity_seconds = 86400 + + bundle = MagicMock() + bundle.bundle_hash = "sha256:" + "0" * 64 + bundle.policy_files = {"allow.cedar": "permit (principal, action, resource);"} + policy_store = MagicMock(spec=PolicyStore) + policy_store.bundle = bundle + policy_store.reload_if_stale = MagicMock() + + catalog = MagicMock() + catalog.entries = {} + catalog.catalog_hash = "sha256:" + "1" * 64 + catalog.exceptions = [] + + return RuntimeContext( + config=config, + tee_provider=MagicMock(), + attestation_report=attestation_report, + signing_key=MagicMock(), + policy_bundle=policy_store, + catalog=catalog, + audit_store=SqliteAuditStore(tmp_path / "audit.db"), + ) + + +def test_bearer_token_reaches_server(ctx): + """AUTH-001: a request without the token must get 401, with it not-401.""" + server = build_server(ctx) + client = TestClient(server.app) + + unauthenticated = client.get("/tools/list") + assert unauthenticated.status_code == 401 + + authenticated = client.get( + "/tools/list", headers={"Authorization": f"Bearer {BEARER}"} + ) + assert authenticated.status_code != 401 + + +def test_health_exempt_from_auth(ctx): + server = build_server(ctx) + client = TestClient(server.app) + assert client.get("/health").status_code != 401 + + +def test_audit_chain_persists_to_store(ctx, tmp_path): + """AUDIT-001: the session_start entry must land in the SQLite DB.""" + build_server(ctx) + conn = sqlite3.connect(tmp_path / "audit.db") + rows = conn.execute( + "SELECT entry_type FROM audit_entries" + ).fetchall() + conn.close() + assert ("session_start",) in rows + + +def test_audit_chain_is_tee_anchored(ctx): + """AUDIT-002: the chain created by the entrypoint must have its anchor set.""" + server = build_server(ctx) + chain = server._audit_chain + assert chain is not None + assert chain.tee_anchor == chain.chain_root + + +def test_proxy_receives_attestation_timestamps(ctx): + """Staleness enforcement requires attestation_generated_at to be wired.""" + server = build_server(ctx) + proxy = server._proxy + assert proxy._attestation_generated_at is not None + assert proxy._attestation_validity_seconds == 86400 From dba786cd46395299065aa4448ab094094c4847e6 Mon Sep 17 00:00:00 2001 From: Imran Siddique <45405841+imran-siddique@users.noreply.github.com> Date: Wed, 10 Jun 2026 17:43:14 -0700 Subject: [PATCH 2/2] feat(policy): surface Cedar @annotation metadata as structured advice on denies (#279) --- src/cmcp_runtime/benchmarks.py | 21 +- src/cmcp_runtime/errors.py | 22 ++ src/cmcp_runtime/mcp/proxy.py | 298 +++++++++++++++++++++--- src/cmcp_runtime/mcp/server.py | 93 +++++++- src/cmcp_runtime/policy/annotations.py | 82 +++++++ src/cmcp_runtime/policy/evaluator.py | 49 +++- tests/soak/run_soak.py | 15 +- tests/unit/conftest.py | 40 ++++ tests/unit/test_benchmarks.py | 26 ++- tests/unit/test_break_glass.py | 6 +- tests/unit/test_call_log_integration.py | 8 +- tests/unit/test_egress_policy.py | 23 +- tests/unit/test_mcp_proxy.py | 88 ++++--- tests/unit/test_mcp_server_auth.py | 1 + tests/unit/test_mid_session_failure.py | 8 +- tests/unit/test_policy_advice.py | 128 ++++++++++ tests/unit/test_session_close.py | 105 +++++++++ tests/unit/test_session_reset.py | 8 +- tests/unit/test_upstream_forwarding.py | 125 ++++++++++ tests/unit/test_workflow_scope.py | 8 +- 20 files changed, 1027 insertions(+), 127 deletions(-) create mode 100644 src/cmcp_runtime/policy/annotations.py create mode 100644 tests/unit/conftest.py create mode 100644 tests/unit/test_policy_advice.py create mode 100644 tests/unit/test_session_close.py create mode 100644 tests/unit/test_upstream_forwarding.py diff --git a/src/cmcp_runtime/benchmarks.py b/src/cmcp_runtime/benchmarks.py index b8ad413c..983b6cd1 100644 --- a/src/cmcp_runtime/benchmarks.py +++ b/src/cmcp_runtime/benchmarks.py @@ -194,12 +194,6 @@ def _make_proxy(bundle: Any, catalog: Any) -> tuple[Any, Any]: session = SessionState(session_id=str(uuid.uuid4())) chain = AuditChain(session_id=session.session_id) - agt_result = MagicMock( - sensitivity_tags=[], - injection_detected=False, - modified_response=b'{"result": "benchmark-ok"}', - ) - with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): proxy = CMCPProxy( @@ -209,8 +203,21 @@ def _make_proxy(bundle: Any, catalog: Any) -> tuple[Any, Any]: audit_chain=chain, config=config, ) + # Mock the gateway seam (pre-call check, upstream forward, response + # scan) so benchmarks measure cmcp overhead, not network latency. proxy._mcp_gateway = MagicMock() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=agt_result) + proxy._mcp_gateway.intercept_tool_call = MagicMock(return_value=(True, "ok")) + proxy._forward_to_upstream = AsyncMock( + return_value='{"result": "benchmark-ok"}' + ) + proxy._mcp_gateway.intercept_tool_response = MagicMock( + return_value=MagicMock( + allowed=True, + content='{"result": "benchmark-ok"}', + threats=[], + action="allowed", + ) + ) return proxy, evaluator diff --git a/src/cmcp_runtime/errors.py b/src/cmcp_runtime/errors.py index e2bd6518..b1aad80f 100644 --- a/src/cmcp_runtime/errors.py +++ b/src/cmcp_runtime/errors.py @@ -38,6 +38,18 @@ class PolicyDeny(CMCPError): code = "POLICY_DENY" http_status = 403 + def __init__( + self, + message: str, + *, + detail: str | None = None, + advice: dict[str, str] | None = None, + ) -> None: + super().__init__(message, detail=detail) + # Annotations of the forbid policies that caused this deny — sourced + # from the hash-pinned policy bundle, safe to reflect to the caller. + self.advice: dict[str, str] = advice or {} + class CatalogToolNameCollision(CMCPError): code = "CATALOG_TOOL_NAME_COLLISION" @@ -84,6 +96,16 @@ class TeeFault(CMCPError): http_status = 500 +class UpstreamUnavailable(CMCPError): + code = "UPSTREAM_UNAVAILABLE" + http_status = 502 + + +class UpstreamToolError(CMCPError): + code = "UPSTREAM_TOOL_ERROR" + http_status = 502 + + class AttestationStale(CMCPError): code = "ATTESTATION_STALE" http_status = 412 diff --git a/src/cmcp_runtime/mcp/proxy.py b/src/cmcp_runtime/mcp/proxy.py index 5fc59ddb..a1a6fdac 100644 --- a/src/cmcp_runtime/mcp/proxy.py +++ b/src/cmcp_runtime/mcp/proxy.py @@ -20,13 +20,14 @@ from datetime import UTC, datetime from typing import Any +import httpx from agent_os.mcp_gateway import GovernancePolicy, MCPGateway # type: ignore[attr-defined] from agent_os.mcp_response_scanner import MCPResponseScanner from cmcp_runtime.audit.chain import AuditChain -from cmcp_runtime.catalog.loader import ToolCatalog +from cmcp_runtime.catalog.loader import CatalogEntry, ToolCatalog from cmcp_runtime.config import Config -from cmcp_runtime.errors import PolicyDeny +from cmcp_runtime.errors import PolicyDeny, UpstreamToolError, UpstreamUnavailable from cmcp_runtime.policy.evaluator import PolicyEvaluator from cmcp_runtime.session.call_log import CallLog, CallRecord, SessionCallLog from cmcp_runtime.session.state import SessionState @@ -46,6 +47,29 @@ class CallResult: deny_reason: str | None latency_us: int audit_entry_hash: str + # Annotations from the forbid policies that matched (deny or advisory). + # Sourced from the hash-pinned policy bundle, safe to reflect to callers. + advice: dict[str, str] | None = None + + +def _cedar_safe(value: Any) -> Any: + """ + Coerce a JSON value into types Cedar can ingest. + + Cedar has no float or null type: a single float anywhere in the request + context makes cedarpy reject the whole request, which fails closed and + denies the call. Floats are preserved as strings; None values are dropped + (policies use `has` checks, so absence is the correct representation). + """ + if isinstance(value, (bool, int, str)): + return value + if isinstance(value, float): + return str(value) + if isinstance(value, dict): + return {k: _cedar_safe(v) for k, v in value.items() if v is not None} + if isinstance(value, (list, tuple)): + return [_cedar_safe(v) for v in value if v is not None] + return str(value) class CMCPProxy: @@ -102,6 +126,81 @@ def __init__( response_scanner=MCPResponseScanner(), ) + # Shared async HTTP client for upstream forwarding; created lazily so + # proxy construction stays sync and tests need no event loop. + self._http: httpx.AsyncClient | None = None + + def rebind_session(self, session: SessionState, audit_chain: AuditChain) -> None: + """ + Point the proxy at a fresh session after the previous one was closed. + + Call logs are recreated for the new session id; catalog, policy + evaluator, and gateway are unchanged. + """ + self._session = session + self._audit = audit_chain + self._call_log = CallLog(session_id=session.session_id) + self._session_call_log = SessionCallLog(session_id=session.session_id) + + async def _forward_to_upstream( + self, + call_id: str, + entry: CatalogEntry, + tool_name: str, + arguments: dict[str, Any], + ) -> str: + """ + Forward the tool call to the attested upstream MCP server (JSON-RPC 2.0 + tools/call over HTTP POST to the catalog entry's server.url). + + Returns the concatenated text content of the MCP result. + + Raises UpstreamUnavailable on transport errors / non-2xx / non-JSON, + UpstreamToolError when the upstream returns a JSON-RPC error object. + """ + if self._http is None: + self._http = httpx.AsyncClient(timeout=httpx.Timeout(30.0)) + payload = { + "jsonrpc": "2.0", + "id": call_id, + "method": "tools/call", + "params": {"name": tool_name, "arguments": arguments}, + } + try: + resp = await self._http.post(entry.server.url, json=payload) + resp.raise_for_status() + body = resp.json() + except httpx.HTTPError as exc: + raise UpstreamUnavailable( + f"Upstream MCP server unreachable: {entry.server.url}", + detail=str(exc), + ) from exc + except ValueError as exc: + raise UpstreamUnavailable( + f"Upstream returned non-JSON body: {entry.server.url}", + detail=str(exc), + ) from exc + if not isinstance(body, dict): + raise UpstreamUnavailable( + f"Upstream returned non-object JSON-RPC body: {entry.server.url}" + ) + if "error" in body: + error = body["error"] if isinstance(body["error"], dict) else {} + raise UpstreamToolError( + f"Upstream tool error from {tool_name}: " + f"{str(error.get('message', 'unknown'))[:200]}" + ) + result = body.get("result", {}) + content = result.get("content", []) if isinstance(result, dict) else [] + texts = [ + c.get("text", "") + for c in content + if isinstance(c, dict) and c.get("type") == "text" + ] + if texts: + return "\n".join(texts) + return json.dumps(result, default=str) + def _check_health(self) -> str | None: """ Check attestation staleness and catalog drift. @@ -159,7 +258,7 @@ def _build_cedar_context( entry = self._catalog.lookup(tool_name) ctx: dict[str, Any] = { "tool_name": tool_name, - "arguments": arguments, + "arguments": _cedar_safe(arguments), "server_identity": entry.server.url if entry else "", "compliance_domain": entry.compliance_domain if entry else "external", "baa_covered": (not entry.requires_baa) if entry else False, @@ -329,10 +428,12 @@ async def call_tool( # Step 2: Cedar policy evaluation cedar_context = self._build_cedar_context(tool_name, arguments, workflow_id) policy_rule: str | None = None + ingress_advice: dict[str, str] = {} try: decision = self._policy.evaluate(cedar_context) policy_rule = decision.rule_matched would_have_denied = decision.would_have_denied + ingress_advice = decision.advice except PolicyDeny as exc: self._audit.append( "tool_call", @@ -367,6 +468,7 @@ async def call_tool( deny_reason=str(exc), latency_us=int(elapsed_ms * 1000), audit_entry_hash=self._audit.chain_tip, + advice=exc.advice or None, ) except Exception as exc: # POLICY-003: Cedar backend raised an unexpected exception (e.g. malformed @@ -387,26 +489,105 @@ async def call_tool( ) raise - # Step 3: AGT MCPGateway enforcement - # AGT handles per-agent rate limiting, parameter sanitization, and - # response scanning. We pass tool_name as the action and arguments as params. + # Step 3a: AGT MCPGateway pre-call interception — per-agent rate + # limiting, parameter sanitization, allow/deny. Fail-closed inside AGT. + agt_allowed, agt_reason = self._mcp_gateway.intercept_tool_call( + agent_id=self._session.session_id, + tool_name=tool_name, + params=arguments, + ) + if not agt_allowed: + logger.warning( + "AGT MCPGateway rejected call: tool=%s reason=%s", tool_name, agt_reason + ) + self._audit.append( + "tool_call", + call_id=call_id, + tool_name=tool_name, + server_identity=entry.server.url, + policy_decision="deny", + policy_rule_matched=f"agt_gateway:{agt_reason[:200]}", + request_payload_hash=request_payload_hash, + session_sensitivity_before=sensitivity_before, + session_sensitivity_after=self._session.max_sensitivity, + workflow_id=workflow_id, + ) + elapsed_ms = (time.perf_counter() - t0) * 1000 + self._record_call( + tool_name=tool_name, + called_at=called_at, + duration_ms=elapsed_ms, + allowed=False, + sensitivity_before=sensitivity_before, + stage_results={"agt_gateway": "deny"}, + call_id=call_id, + catalog_entry=entry, + policy_decision="deny", + ) + return CallResult( + call_id=call_id, + tool_name=tool_name, + allowed=False, + would_have_denied=would_have_denied, + response=None, + deny_reason=agt_reason, + latency_us=int(elapsed_ms * 1000), + audit_entry_hash=self._audit.chain_tip, + ) + + # Step 3b: forward to the attested upstream MCP server. try: - agt_result = await self._mcp_gateway.call_tool( # type: ignore[attr-defined] + response_text = await self._forward_to_upstream( + call_id, entry, tool_name, arguments + ) + except (UpstreamUnavailable, UpstreamToolError) as exc: + logger.warning("Upstream call failed: tool=%s error=%s", tool_name, exc) + self._audit.append( + "fault", + call_id=call_id, tool_name=tool_name, - arguments=arguments, - agent_id=self._session.session_id, + server_identity=entry.server.url, + policy_decision="fault", + policy_rule_matched=f"upstream:{exc.code}", + request_payload_hash=request_payload_hash, + session_sensitivity_before=sensitivity_before, + session_sensitivity_after=self._session.max_sensitivity, + detail={"error_code": exc.code}, ) - except Exception as exc: - # AGT denied or errored — map to our error types - logger.warning("AGT MCPGateway rejected call: tool=%s error=%s", tool_name, exc) + elapsed_ms = (time.perf_counter() - t0) * 1000 + self._record_call( + tool_name=tool_name, + called_at=called_at, + duration_ms=elapsed_ms, + allowed=False, + sensitivity_before=sensitivity_before, + stage_results={"upstream": "fault"}, + call_id=call_id, + catalog_entry=entry, + policy_decision="fault", + ) + return CallResult( + call_id=call_id, + tool_name=tool_name, + allowed=False, + would_have_denied=would_have_denied, + response=None, + deny_reason=f"upstream_error:{exc.code}", + latency_us=int(elapsed_ms * 1000), + audit_entry_hash=self._audit.chain_tip, + ) + + # Step 3c: response size guard (DOS-002) before scanning. + if len(response_text.encode()) > self._config.max_response_size_bytes: self._audit.append( "tool_call", call_id=call_id, tool_name=tool_name, server_identity=entry.server.url, policy_decision="deny", - policy_rule_matched=f"agt_gateway:{type(exc).__name__}", + policy_rule_matched="response_size_exceeded", request_payload_hash=request_payload_hash, + response_inspection_result="size_exceeded", session_sensitivity_before=sensitivity_before, session_sensitivity_after=self._session.max_sensitivity, workflow_id=workflow_id, @@ -418,7 +599,7 @@ async def call_tool( duration_ms=elapsed_ms, allowed=False, sensitivity_before=sensitivity_before, - stage_results={"agt_gateway": "deny"}, + stage_results={"inspection": "size_exceeded"}, call_id=call_id, catalog_entry=entry, policy_decision="deny", @@ -429,45 +610,95 @@ async def call_tool( allowed=False, would_have_denied=would_have_denied, response=None, - deny_reason=str(exc), + deny_reason="response_size_exceeded", + latency_us=int(elapsed_ms * 1000), + audit_entry_hash=self._audit.chain_tip, + ) + + # Step 3d: AGT response interception — injection / credential / PII scan. + scan = self._mcp_gateway.intercept_tool_response( + agent_id=self._session.session_id, + tool_name=tool_name, + response_content=response_text, + ) + injection_detected = bool(scan.threats) + if not scan.allowed: + async with self._session.mutation_lock: + self._session.update_from_inspection( + call_id=call_id, + sensitivity_tags=[entry.sensitivity_level], + injection_detected=injection_detected, + response_allowed=False, + ) + threat_categories = ",".join( + sorted({str(t.get("category", "unknown")) for t in scan.threats}) + ) + self._audit.append( + "tool_call", + call_id=call_id, + tool_name=tool_name, + server_identity=entry.server.url, + policy_decision="deny", + policy_rule_matched=f"response_scan:{threat_categories[:200]}", + request_payload_hash=request_payload_hash, + response_inspection_result="injection_detected", + session_sensitivity_before=sensitivity_before, + session_sensitivity_after=self._session.max_sensitivity, + workflow_id=workflow_id, + ) + elapsed_ms = (time.perf_counter() - t0) * 1000 + self._record_call( + tool_name=tool_name, + called_at=called_at, + duration_ms=elapsed_ms, + allowed=False, + sensitivity_before=sensitivity_before, + stage_results={"response_scan": "deny"}, + call_id=call_id, + catalog_entry=entry, + policy_decision="deny", + ) + return CallResult( + call_id=call_id, + tool_name=tool_name, + allowed=False, + would_have_denied=would_have_denied, + response=None, + deny_reason="response_blocked_by_scanner", latency_us=int(elapsed_ms * 1000), audit_entry_hash=self._audit.chain_tip, ) + # Scanner may have sanitized the content (ResponsePolicy.SANITIZE). + agt_result: str = scan.content if scan.content is not None else response_text # Step 4: session update from response sensitivity # AUTH-002: lock protects against race with concurrent session reset requests. - response_sensitivity = getattr(agt_result, "sensitivity_tags", []) - injection_detected = getattr(agt_result, "injection_detected", False) - # INJECT-003: capture scanner attribution and matched pattern for audit chain detail - injection_scanner = getattr(agt_result, "injection_scanner", None) - injection_pattern = getattr(agt_result, "matched_pattern", None) or getattr( - agt_result, "injection_pattern", None + # Sensitivity comes from the attested catalog entry's declared level. + response_sensitivity = [entry.sensitivity_level] + injection_scanner = "agt_response_scanner" if injection_detected else None + injection_pattern = ( + ",".join(sorted({str(t.get("category", "unknown")) for t in scan.threats})) + if injection_detected + else None ) - # INJECT-007: capture threshold so audit consumers can replay the decision - injection_threshold = getattr(agt_result, "injection_threshold", None) + injection_threshold = None async with self._session.mutation_lock: self._session.update_from_inspection( call_id=call_id, - sensitivity_tags=response_sensitivity or [entry.sensitivity_level], + sensitivity_tags=response_sensitivity, injection_detected=injection_detected, response_allowed=True, ) # Step 5: egress Cedar policy check - # Derive response bytes for size accounting and egress evaluation. - # Prefer a bytes-typed modified_response (Stage 2 redaction output); - # fall back to str() so we never block on an un-serialisable AGT object. - modified = getattr(agt_result, "modified_response", None) - if isinstance(modified, bytes): - response_bytes: bytes = modified - else: - response_bytes = str(agt_result).encode() + response_bytes: bytes = agt_result.encode() try: egress_decision = self._policy.authorize_egress( tool_name, response_bytes, self._session ) egress_would_deny = egress_decision.would_have_denied + egress_advice = egress_decision.advice except PolicyDeny as exc: egress_deny_reason = str(exc) self._audit.append( @@ -490,10 +721,12 @@ async def call_tool( deny_reason=egress_deny_reason, latency_us=int((time.perf_counter() - t0) * 1_000_000), audit_entry_hash=self._audit.chain_tip, + advice=exc.advice or None, ) # Merge egress advisory flag into the overall would_have_denied would_have_denied = would_have_denied or egress_would_deny + advisory_advice = {**ingress_advice, **egress_advice} # Step 6: audit chain write policy_decision: Any = "advisory_deny" if would_have_denied else "allow" @@ -548,4 +781,5 @@ async def call_tool( deny_reason=None, latency_us=latency_us, audit_entry_hash=self._audit.chain_tip, + advice=advisory_advice or None, ) diff --git a/src/cmcp_runtime/mcp/server.py b/src/cmcp_runtime/mcp/server.py index cba16057..5dae3005 100644 --- a/src/cmcp_runtime/mcp/server.py +++ b/src/cmcp_runtime/mcp/server.py @@ -153,6 +153,9 @@ def __init__( self._session = session self._max_request_bytes = max_request_bytes self._audit = audit_chain + # Chains of closed sessions, kept so /audit/export still serves them + # after the live session rotates. + self._closed_chains: dict[str, AuditChain] = {} self._kernel = StatelessKernel() # NET-002: rate-limit unauthenticated /health before auth middleware runs. # Starlette applies middleware outermost-first (first in list = first to run). @@ -188,6 +191,11 @@ def __init__( self._session_reset, methods=["POST"], ), + Route( + "/sessions/{session_id}/close", + self._session_close, + methods=["POST"], + ), Route("/catalog/exception", self._catalog_exception, methods=["POST"]), ], middleware=middleware, @@ -300,6 +308,23 @@ async def _handle_tool_call(self, rpc_id: Any, params: dict[str, Any]) -> Respon ) if not result.allowed: + # Upstream transport/tool failure is a 502, not a policy deny. + if (result.deny_reason or "").startswith("upstream_error:"): + return JSONResponse( + { + "jsonrpc": "2.0", + "error": { + "code": -32000, + "message": "Upstream MCP server error", + "data": { + "error_code": result.deny_reason.removeprefix("upstream_error:"), + "call_id": call_id, + }, + }, + "id": rpc_id, + }, + status_code=502, + ) _HEALTH_REASONS = {"attestation_stale", "catalog_drift"} if result.deny_reason in _HEALTH_REASONS: return JSONResponse( @@ -327,16 +352,22 @@ async def _handle_tool_call(self, rpc_id: Any, params: dict[str, Any]) -> Respon "POLICY_DENY: call_id=%s error_code=%s reason=%s", call_id, error_code, result.deny_reason, ) + error_data: dict[str, Any] = { + "error_code": error_code, + "call_id": call_id, + } + # Advice annotations come from the hash-pinned policy bundle + # (operator-authored, not caller input), so reflecting them does + # not violate INJECT-003. They carry e.g. HITL escalation payloads. + if result.advice: + error_data["advice"] = result.advice return JSONResponse( { "jsonrpc": "2.0", "error": { "code": -32000, "message": "Request denied by policy", - "data": { - "error_code": error_code, - "call_id": call_id, - }, + "data": error_data, }, "id": rpc_id, }, @@ -349,6 +380,10 @@ async def _handle_tool_call(self, rpc_id: Any, params: dict[str, Any]) -> Respon "would_have_denied": result.would_have_denied, "latency_us": result.latency_us, } + if self._session is not None: + cmcp_meta["session_id"] = self._session.session_id + if result.would_have_denied and result.advice: + cmcp_meta["advice"] = result.advice if workflow_id is not None: cmcp_meta["workflow_id"] = workflow_id return JSONResponse({ @@ -447,10 +482,10 @@ async def _audit_export(self, request: Request) -> Response: {"error": "query parameter 'session_id' is required"}, status_code=400, ) + # Closed sessions keep their chain available for export after rotation. + chain = self._closed_chains.get(session_id, self._audit_chain) try: - bundle = self._session_manager.get_audit_bundle( - session_id, self._audit_chain - ) + bundle = self._session_manager.get_audit_bundle(session_id, chain) except ValueError as exc: logger.error( "Audit chain integrity failure: session_id=%s error=%s", @@ -462,6 +497,50 @@ async def _audit_export(self, request: Request) -> Response: ) return JSONResponse(bundle) + async def _session_close(self, request: Request) -> Response: + """POST /sessions/{session_id}/close — close the session, return its signed TRACE Claim. + + Appends the session_end audit entry, signs the claim, then rotates the + gateway onto a fresh session so subsequent tool calls keep working. + The closed session's claim stays available at + GET /sessions/{session_id}/trace-claim and its audit bundle at + GET /audit/export?session_id=. + """ + if ( + self._session_manager is None + or self._session is None + or self._audit_chain is None + ): + return JSONResponse( + {"error": "session management not available"}, status_code=501 + ) + session_id: str = request.path_params["session_id"] + if session_id != self._session.session_id: + return JSONResponse( + {"error": f"unknown or already closed session_id={session_id}"}, + status_code=404, + ) + + claim = self._session_manager.close_session( + session_id, + self._session, + self._audit_chain, + call_log=getattr(self._proxy, "_call_log", None), + session_call_log=getattr(self._proxy, "_session_call_log", None), + ) + self._closed_chains[session_id] = self._audit_chain + + # Rotate onto a fresh session so the gateway keeps serving. + new_session, new_chain = self._session_manager.create_session() + self._session = new_session + self._audit_chain = new_chain + self._audit = new_chain + self._proxy.rebind_session(new_session, new_chain) + logger.info( + "Session closed via API: closed=%s new=%s", session_id, new_session.session_id + ) + return JSONResponse(claim) + async def _catalog_exception(self, request: Request) -> Response: """POST /catalog/exception — add a break-glass catalog exception at runtime. diff --git a/src/cmcp_runtime/policy/annotations.py b/src/cmcp_runtime/policy/annotations.py new file mode 100644 index 00000000..6f1539ba --- /dev/null +++ b/src/cmcp_runtime/policy/annotations.py @@ -0,0 +1,82 @@ +""" +Cedar policy annotation extraction. + +Cedar has no first-class "advice" construct: annotations (``@key("value")``) +attached to a policy are the supported way to carry structured metadata such +as HITL escalation instructions. cedarpy reports which policies determined a +decision (``diagnostics.reasons`` as implicit ids ``policy0``, ``policy1``, +... in source order) but only surfaces the ``@id`` annotation value, so the +full annotation set must be recovered from the policy source. + +This module parses annotations from a combined Cedar policy string, keyed by +the same implicit ids cedarpy assigns, so a deny decision can be mapped back +to the matched policies' annotations. +""" + +from __future__ import annotations + +import re + +# @key("value") — value may contain escaped quotes/backslashes. +_ANNOTATION_RE = re.compile( + r'@([A-Za-z_][A-Za-z0-9_]*)\s*\(\s*"((?:[^"\\]|\\.)*)"\s*\)' +) + +# A policy statement: zero or more annotations followed by permit/forbid. +# Matched against comment-stripped source, in order, so the Nth match is +# cedarpy's implicit id "policyN". +_POLICY_RE = re.compile( + r'((?:@[A-Za-z_][A-Za-z0-9_]*\s*\(\s*"(?:[^"\\]|\\.)*"\s*\)\s*)*)' + r"\b(permit|forbid)\s*\(" +) + + +def _strip_line_comments(text: str) -> str: + """Remove ``// ...`` comments, ignoring ``//`` inside string literals.""" + out_lines: list[str] = [] + for line in text.splitlines(): + in_string = False + escaped = False + cut = len(line) + for i, ch in enumerate(line): + if escaped: + escaped = False + continue + if ch == "\\" and in_string: + escaped = True + continue + if ch == '"': + in_string = not in_string + continue + if ch == "/" and not in_string and line[i : i + 2] == "//": + cut = i + break + out_lines.append(line[:cut]) + return "\n".join(out_lines) + + +def _unescape(value: str) -> str: + return value.replace('\\"', '"').replace("\\\\", "\\") + + +def parse_policy_annotations(policy_text: str) -> dict[str, dict[str, str]]: + """ + Map implicit Cedar policy ids to their annotations. + + Returns ``{"policy0": {"reason": "...", ...}, "policy2": {...}}``; policies + without annotations are omitted. Ordering follows statement order in the + source, matching the ids cedarpy assigns when parsing the same string. + """ + stripped = _strip_line_comments(policy_text) + annotations: dict[str, dict[str, str]] = {} + for index, match in enumerate(_POLICY_RE.finditer(stripped)): + block = match.group(1) + if not block: + continue + parsed = { + key: _unescape(value) + for key, value in _ANNOTATION_RE.findall(block) + } + if parsed: + annotations[f"policy{index}"] = parsed + return annotations diff --git a/src/cmcp_runtime/policy/evaluator.py b/src/cmcp_runtime/policy/evaluator.py index 2c48c96f..70fff06c 100644 --- a/src/cmcp_runtime/policy/evaluator.py +++ b/src/cmcp_runtime/policy/evaluator.py @@ -16,6 +16,7 @@ from cmcp_runtime.config import Config, EnforcementMode from cmcp_runtime.errors import PolicyDeny +from cmcp_runtime.policy.annotations import parse_policy_annotations from cmcp_runtime.policy.bundle import PolicyBundle, PolicyStore from cmcp_runtime.session.state import SENSITIVITY_ORDER @@ -73,6 +74,8 @@ def __init__(self, bundle: PolicyBundle | PolicyStore, config: Config) -> None: policy_content=combined_policy, mode="auto", # cedarpy > cli > builtin ) + self._combined_policy = combined_policy + self._annotations = parse_policy_annotations(combined_policy) logger.info( "PolicyEvaluator ready: bundle_hash=%s enforcement=%s backend=%s", initial_bundle.bundle_hash, @@ -89,9 +92,45 @@ def _maybe_reload(self) -> None: content for _, content in sorted(bundle.policy_files.items()) ) self._backend = CedarBackend(policy_content=combined_policy, mode="auto") + self._combined_policy = combined_policy + self._annotations = parse_policy_annotations(combined_policy) self._current_hash = bundle.bundle_hash logger.info("PolicyEvaluator backend refreshed: new_hash=%s", self._current_hash) + def _advice_for_deny(self, context: dict[str, Any]) -> dict[str, str]: + """ + Best-effort: recover the annotations of the forbid policies that caused + a deny, to return as structured advice (e.g. HITL escalation payloads). + + AGT's CedarBackend does not expose cedarpy's diagnostics.reasons, so + this re-evaluates the same request directly with cedarpy purely for + diagnostics — the authorization decision itself is never taken from + here. Runs only on the deny path; returns {} on any failure. + """ + if not self._annotations: + return {} + try: + import cedarpy + + request = self._backend._build_cedar_request(context) + response = cedarpy.is_authorized( + request={ + "principal": request["principal"], + "action": request["action"], + "resource": request["resource"], + "context": request.get("context", {}), + }, + policies=self._combined_policy, + entities=getattr(self._backend, "_entities", []), + ) + advice: dict[str, str] = {} + for policy_id in response.diagnostics.reasons: + advice.update(self._annotations.get(policy_id, {})) + return advice + except Exception: + logger.debug("Advice extraction failed", exc_info=True) + return {} + def evaluate(self, context: dict[str, Any]) -> PolicyDecision: """ Evaluate a tool call against the Cedar policy bundle. @@ -121,11 +160,15 @@ def evaluate(self, context: dict[str, Any]) -> PolicyDecision: evaluation_ms=evaluation_ms, ) - # Cedar denied — apply enforcement mode + # Cedar denied — recover advice annotations from the matched policies, + # then apply enforcement mode. + advice = self._advice_for_deny(context) + if self._mode == EnforcementMode.ENFORCING: raise PolicyDeny( f"Policy denied tool call: {context.get('tool_name', '?')}", detail=f"rule={rule} eval_ms={evaluation_ms:.2f}", + advice=advice, ) if self._mode == EnforcementMode.ADVISORY: @@ -137,7 +180,7 @@ def evaluate(self, context: dict[str, Any]) -> PolicyDecision: allowed=True, enforcement_mode=self._mode, rule_matched=rule, - advice={}, + advice=advice, evaluation_ms=evaluation_ms, would_have_denied=True, ) @@ -147,7 +190,7 @@ def evaluate(self, context: dict[str, Any]) -> PolicyDecision: allowed=True, enforcement_mode=self._mode, rule_matched=rule, - advice={}, + advice=advice, evaluation_ms=evaluation_ms, would_have_denied=True, ) diff --git a/tests/soak/run_soak.py b/tests/soak/run_soak.py index e39a2a66..0cc75091 100644 --- a/tests/soak/run_soak.py +++ b/tests/soak/run_soak.py @@ -144,10 +144,7 @@ def _make_soak_gateway( with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(io.StringIO()): evaluator = PolicyEvaluator(bundle=bundle, config=config) - agt_result = MagicMock( - sensitivity_tags=[], injection_detected=False, - modified_response=b'{"result": "soak-ok"}', - ) + _soak_response = '{"result": "soak-ok"}' with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): @@ -160,8 +157,16 @@ def _make_soak_gateway( attestation_generated_at=datetime.now(UTC), attestation_validity_seconds=attestation_validity_seconds, ) + # Mock the AGT gateway seam + upstream forwarding (new proxy step 3 seam) proxy._mcp_gateway = MagicMock() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=agt_result) + proxy._mcp_gateway.intercept_tool_call = MagicMock(return_value=(True, "ok")) + proxy._forward_to_upstream = AsyncMock(return_value=_soak_response) + proxy._mcp_gateway.intercept_tool_response = MagicMock(return_value=MagicMock( + allowed=True, + content=_soak_response, + threats=[], + action="allowed", + )) with patch("cmcp_runtime.mcp.server.StatelessKernel"): server = MCPServer(proxy, session=session, audit_chain=chain, bearer_token=bearer_token) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 00000000..eed03960 --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,40 @@ +"""Shared unit-test helpers for mocking the CMCPProxy gateway seam. + +The proxy's step 3 calls three seams: + 1. _mcp_gateway.intercept_tool_call(agent_id=, tool_name=, params=) -> (bool, str) [sync] + 2. proxy._forward_to_upstream(call_id, entry, tool_name, arguments) -> str [async] + 3. _mcp_gateway.intercept_tool_response(agent_id=, tool_name=, response_content=) + -> object with .allowed (bool), .content (str|None), .threats (list[dict]), + .action (str) [sync] +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + + +def wire_mock_gateway( + proxy, + *, + response_text: str = "tool response", + call_allowed: bool = True, + call_reason: str = "ok", + scan_allowed: bool = True, + threats: list[dict] | None = None, + scan_content: str | None = None, +): + """Replace the proxy's AGT gateway + upstream forwarding with mocks.""" + proxy._mcp_gateway = MagicMock() + proxy._mcp_gateway.intercept_tool_call = MagicMock( + return_value=(call_allowed, call_reason) + ) + proxy._forward_to_upstream = AsyncMock(return_value=response_text) + proxy._mcp_gateway.intercept_tool_response = MagicMock( + return_value=MagicMock( + allowed=scan_allowed, + content=scan_content if scan_content is not None else response_text, + threats=threats or [], + action="allowed" if scan_allowed else "blocked", + ) + ) + return proxy diff --git a/tests/unit/test_benchmarks.py b/tests/unit/test_benchmarks.py index db323a5e..74662865 100644 --- a/tests/unit/test_benchmarks.py +++ b/tests/unit/test_benchmarks.py @@ -6,8 +6,30 @@ import pytest +from tests.unit.conftest import wire_mock_gateway -def test_benchmark_runs_smoke(tmp_path): + +@pytest.fixture +def benchmark_gateway_seam(monkeypatch): + """Rewire the benchmark proxy onto the current gateway seam. + + benchmarks._make_proxy still mocks the removed `_mcp_gateway.call_tool` + coroutine; wrap it so the proxy gets the intercept_tool_call / + _forward_to_upstream / intercept_tool_response mocks instead. + """ + import cmcp_runtime.benchmarks as benchmarks + + original = benchmarks._make_proxy + + def _patched(bundle, catalog): + proxy, evaluator = original(bundle, catalog) + wire_mock_gateway(proxy, response_text='{"result": "benchmark-ok"}') + return proxy, evaluator + + monkeypatch.setattr(benchmarks, "_make_proxy", _patched) + + +def test_benchmark_runs_smoke(tmp_path, benchmark_gateway_seam): """Smoke test: benchmark produces valid JSON output with all required keys.""" import asyncio @@ -29,7 +51,7 @@ def test_benchmark_runs_smoke(tmp_path): assert stats["p99"] >= stats["p95"] >= stats["p50"] -def test_benchmark_writes_output_file(tmp_path): +def test_benchmark_writes_output_file(tmp_path, benchmark_gateway_seam): """Benchmark writes a timestamped JSON file to the output directory.""" import asyncio diff --git a/tests/unit/test_break_glass.py b/tests/unit/test_break_glass.py index f5cb8cad..6f8e2fce 100644 --- a/tests/unit/test_break_glass.py +++ b/tests/unit/test_break_glass.py @@ -19,6 +19,7 @@ from cmcp_runtime.mcp.server import MCPServer from cmcp_runtime.policy.evaluator import PolicyDecision, PolicyEvaluator from cmcp_runtime.session.state import SessionState +from tests.unit.conftest import wire_mock_gateway # ── Helpers ─────────────────────────────────────────────────────────────────── @@ -333,8 +334,6 @@ def _make_real_proxy_for_break_glass(): ) ) - mock_agt_result = MagicMock(sensitivity_tags=[], injection_detected=False, modified_response=b"ok") - with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): proxy = CMCPProxy( @@ -344,8 +343,7 @@ def _make_real_proxy_for_break_glass(): audit_chain=chain, config=config, ) - proxy._mcp_gateway = MagicMock() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=mock_agt_result) + wire_mock_gateway(proxy, response_text="ok") return proxy, chain diff --git a/tests/unit/test_call_log_integration.py b/tests/unit/test_call_log_integration.py index 2c973d46..97b6d449 100644 --- a/tests/unit/test_call_log_integration.py +++ b/tests/unit/test_call_log_integration.py @@ -2,7 +2,7 @@ from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest @@ -18,6 +18,7 @@ from cmcp_runtime.policy.evaluator import PolicyDecision, PolicyEvaluator from cmcp_runtime.session.call_log import CallLog from cmcp_runtime.session.state import SessionState +from tests.unit.conftest import wire_mock_gateway def _make_entry(tool_name: str = "test.tool") -> CatalogEntry: @@ -81,10 +82,7 @@ def _make_proxy(catalog=None, evaluator=None, call_log=None): with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): proxy = CMCPProxy(cat, ev, session, chain, cfg, call_log=call_log) - proxy._mcp_gateway = MagicMock() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=[], injection_detected=False - )) + wire_mock_gateway(proxy) return proxy, session, chain diff --git a/tests/unit/test_egress_policy.py b/tests/unit/test_egress_policy.py index c00d3bbb..21587da2 100644 --- a/tests/unit/test_egress_policy.py +++ b/tests/unit/test_egress_policy.py @@ -2,7 +2,7 @@ from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest @@ -18,6 +18,7 @@ from cmcp_runtime.policy.bundle import PolicyBundle, PolicyManifest from cmcp_runtime.policy.evaluator import PolicyDecision, PolicyEvaluator from cmcp_runtime.session.state import SessionState +from tests.unit.conftest import wire_mock_gateway # ── Helpers ─────────────────────────────────────────────────────────────────── @@ -117,10 +118,7 @@ def _side_effect(context: dict) -> PolicyDecision: with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): proxy = CMCPProxy(catalog, evaluator, session, chain, cfg) - proxy._mcp_gateway = MagicMock() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=[], injection_detected=False, modified_response=None - )) + wire_mock_gateway(proxy) return proxy, session, chain, evaluator @@ -292,10 +290,7 @@ def _make_evaluator_for(session_sensitivity: str) -> MagicMock: with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): high_proxy = CMCPProxy(catalog, high_ev, high_session, high_chain, cfg) - high_proxy._mcp_gateway = MagicMock() - high_proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=[], injection_detected=False, modified_response=None - )) + wire_mock_gateway(high_proxy) high_result = await high_proxy.call_tool("c1", "test.tool", {}) assert high_result.allowed is False @@ -308,10 +303,7 @@ def _make_evaluator_for(session_sensitivity: str) -> MagicMock: with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): low_proxy = CMCPProxy(catalog, low_ev, low_session, low_chain, cfg) - low_proxy._mcp_gateway = MagicMock() - low_proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=[], injection_detected=False, modified_response=None - )) + wire_mock_gateway(low_proxy) low_result = await low_proxy.call_tool("c1", "test.tool", {}) assert low_result.allowed is True @@ -352,10 +344,7 @@ def _egress_aware(context: dict) -> PolicyDecision: with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): proxy = CMCPProxy(catalog, evaluator, session, chain, cfg) - proxy._mcp_gateway = MagicMock() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=[], injection_detected=False, modified_response=None - )) + wire_mock_gateway(proxy) # Before reset — high sensitivity should be blocked result_before = await proxy.call_tool("c1", "test.tool", {}) diff --git a/tests/unit/test_mcp_proxy.py b/tests/unit/test_mcp_proxy.py index a3240622..7d283136 100644 --- a/tests/unit/test_mcp_proxy.py +++ b/tests/unit/test_mcp_proxy.py @@ -2,7 +2,7 @@ from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest @@ -17,6 +17,7 @@ from cmcp_runtime.errors import PolicyDeny from cmcp_runtime.policy.evaluator import PolicyDecision, PolicyEvaluator from cmcp_runtime.session.state import SessionState +from tests.unit.conftest import wire_mock_gateway def _make_entry(tool_name: str = "test.tool") -> CatalogEntry: @@ -89,10 +90,7 @@ def _make_proxy(catalog=None, evaluator=None, mode=EnforcementMode.ENFORCING): with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): proxy = CMCPProxy(cat, ev, session, chain, cfg, attestation_platform="software-only") - proxy._mcp_gateway = MagicMock() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=[], injection_detected=False - )) + wire_mock_gateway(proxy) return proxy, session, chain @@ -152,10 +150,6 @@ async def test_proxy_updates_session_state_on_allow(): catalog = ToolCatalog(entries={"test.tool": entry}, catalog_hash="sha256:" + "1" * 64) proxy, session, _ = _make_proxy(catalog=catalog) - proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=["pii"], injection_detected=False - )) - assert session.max_sensitivity == "public" await proxy.call_tool("c1", "test.tool", {}) assert session.max_sensitivity == "pii" @@ -227,10 +221,7 @@ async def test_cedar_context_includes_attestation_platform(): _make_catalog(), evaluator, session, chain, cfg, attestation_platform="amd-sev-snp", ) - proxy._mcp_gateway = MagicMock() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=[], injection_detected=False - )) + wire_mock_gateway(proxy) await proxy.call_tool("c1", "test.tool", {}) ctx = evaluator.evaluate.call_args[0][0] @@ -343,21 +334,22 @@ async def test_audit_entry_detail_includes_injection_scanner_and_pattern(): include injection_scanner and matched_pattern so a verifier can reconstruct why the request was denied.""" proxy, _, chain = _make_proxy() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=[], - injection_detected=True, - injection_scanner="agt_mcp", - matched_pattern="test_pattern", - injection_pattern=None, - )) + # Allowed-but-detected path: scanner reports threats but does not block. + wire_mock_gateway( + proxy, + scan_allowed=True, + threats=[{"category": "prompt_injection", "description": "x"}], + ) await proxy.call_tool("c1", "test.tool", {"q": "hello"}) tool_entries = [e for e in chain.entries if e.entry_type == "tool_call"] entry = tool_entries[-1] assert entry.detail is not None, "detail must be set when injection is detected" - assert entry.detail["injection_scanner"] == "agt_mcp" - assert entry.detail["matched_pattern"] == "test_pattern" + assert entry.detail["injection_scanner"] == "agt_response_scanner" + assert entry.detail["matched_pattern"] == "prompt_injection" + # INJECT-007 threshold no longer applies to the AGT response scanner + assert "injection_threshold" not in entry.detail @pytest.mark.asyncio @@ -375,19 +367,55 @@ async def test_audit_entry_detail_is_none_when_no_injection(): @pytest.mark.asyncio async def test_audit_entry_detail_falls_back_to_unknown_when_fields_absent(): - """INJECT-003 — when injection_detected=True but scanner/pattern attrs are absent, - detail values must fall back to 'unknown' rather than crashing.""" + """INJECT-003 — when a scanner threat dict is missing its 'category' key, the + matched_pattern must fall back to 'unknown' rather than crashing.""" proxy, _, chain = _make_proxy() - agt_mock = MagicMock(spec=[]) # no attributes beyond spec - agt_mock.sensitivity_tags = [] - agt_mock.injection_detected = True - # injection_scanner and matched_pattern intentionally absent - proxy._mcp_gateway.call_tool = AsyncMock(return_value=agt_mock) + # Threat dict intentionally missing the "category" key + wire_mock_gateway(proxy, scan_allowed=True, threats=[{"description": "x"}]) await proxy.call_tool("c1", "test.tool", {}) tool_entries = [e for e in chain.entries if e.entry_type == "tool_call"] entry = tool_entries[-1] assert entry.detail is not None - assert entry.detail["injection_scanner"] == "unknown" + assert entry.detail["injection_scanner"] == "agt_response_scanner" assert entry.detail["matched_pattern"] == "unknown" + + +# ── Cedar-safe context coercion ──────────────────────────────────────────────── + + +def test_cedar_safe_coerces_floats_and_drops_none(): + from cmcp_runtime.mcp.proxy import _cedar_safe + + out = _cedar_safe({ + "score": 72.3, + "labs": {"glucose": 9.2, "note": None}, + "tags": ["a", 1, 2.5, None], + "count": 3, + "active": True, + }) + assert out == { + "score": "72.3", + "labs": {"glucose": "9.2"}, + "tags": ["a", 1, "2.5"], + "count": 3, + "active": True, + } + + +@pytest.mark.asyncio +async def test_float_arguments_do_not_fail_policy_evaluation(): + """A float in tool arguments must not crash Cedar and deny the call.""" + proxy, _, _ = _make_proxy() + captured = {} + original = proxy._policy.evaluate + + def capture(ctx): + captured.update(ctx) + return original(ctx) + + proxy._policy.evaluate = capture + result = await proxy.call_tool("c1", "test.tool", {"risk_score": 72.3}) + assert result.allowed is True + assert captured["arguments"] == {"risk_score": "72.3"} diff --git a/tests/unit/test_mcp_server_auth.py b/tests/unit/test_mcp_server_auth.py index dd8cefa5..c34fab02 100644 --- a/tests/unit/test_mcp_server_auth.py +++ b/tests/unit/test_mcp_server_auth.py @@ -400,6 +400,7 @@ def test_deny_response_does_not_include_internal_reason(): audit_entry_hash=None, would_have_denied=False, latency_us=0, + advice=None, )) with patch("cmcp_runtime.mcp.server.StatelessKernel"): server = MCPServer(proxy) diff --git a/tests/unit/test_mid_session_failure.py b/tests/unit/test_mid_session_failure.py index 9b74ee54..cc34122a 100644 --- a/tests/unit/test_mid_session_failure.py +++ b/tests/unit/test_mid_session_failure.py @@ -10,7 +10,7 @@ from __future__ import annotations from datetime import UTC, datetime, timedelta -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest @@ -24,6 +24,7 @@ from cmcp_runtime.config import AttestationConfig, Config, EnforcementMode from cmcp_runtime.policy.evaluator import PolicyDecision, PolicyEvaluator from cmcp_runtime.session.state import SessionState +from tests.unit.conftest import wire_mock_gateway # ── Helpers ─────────────────────────────────────────────────────────────────── @@ -95,10 +96,7 @@ def _make_proxy( attestation_validity_seconds=attestation_validity_seconds, catalog_hash=catalog_hash, ) - proxy._mcp_gateway = MagicMock() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=[], injection_detected=False - )) + wire_mock_gateway(proxy) return proxy, session, chain diff --git a/tests/unit/test_policy_advice.py b/tests/unit/test_policy_advice.py new file mode 100644 index 00000000..4b54c5cf --- /dev/null +++ b/tests/unit/test_policy_advice.py @@ -0,0 +1,128 @@ +"""Tests for Cedar annotation-based advice on policy denies.""" + +from __future__ import annotations + +import pytest + +from cmcp_runtime.config import AttestationConfig, Config, EnforcementMode +from cmcp_runtime.errors import PolicyDeny +from cmcp_runtime.policy.annotations import parse_policy_annotations +from cmcp_runtime.policy.bundle import PolicyBundle, PolicyManifest +from cmcp_runtime.policy.evaluator import PolicyEvaluator + +HITL_POLICY = """ +// Rule 1: permit everything by default. +permit (principal, action, resource); + +// Rule 2: HITL escalation for high-risk patients (EU AI Act Art. 14). +@id("hitl-high-risk") +@reason("human-review-required") +@regulation("eu-ai-act-art-14") +@reviewer_role("attending-physician") +forbid (principal, action == Action::"Ehr.treatmentPlanWriter", resource) +when { context.arguments has patient_risk_category + && context.arguments.patient_risk_category == "high" }; +""" + + +def _bundle(policy_content: str) -> PolicyBundle: + return PolicyBundle( + manifest=PolicyManifest( + version="1.0.0", + authored_at="2026-06-10T00:00:00Z", + author_identity="test", + commit_sha="abc", + ), + policy_files={"allow.cedar": policy_content}, + schema_content='{"cMCP": {}}', + bundle_hash="sha256:" + "0" * 64, + ) + + +def _config(mode: EnforcementMode) -> Config: + return Config(attestation=AttestationConfig(enforcement_mode=mode)) + + +def _context(risk: str) -> dict: + return { + "tool_name": "ehr.treatment_plan_writer", + "arguments": {"patient_risk_category": risk}, + "session_max_sensitivity": "confidential", + "workflow_id": "clinical-decision-support", + } + + +# ── parse_policy_annotations ────────────────────────────────────────────────── + + +def test_parse_annotations_maps_implicit_ids(): + annotations = parse_policy_annotations(HITL_POLICY) + assert "policy0" not in annotations # unannotated permit + assert annotations["policy1"] == { + "id": "hitl-high-risk", + "reason": "human-review-required", + "regulation": "eu-ai-act-art-14", + "reviewer_role": "attending-physician", + } + + +def test_parse_annotations_ignores_commented_policies(): + text = """ +// permit (principal, action, resource); +@reason("real") +forbid (principal, action, resource); +""" + annotations = parse_policy_annotations(text) + assert annotations == {"policy0": {"reason": "real"}} + + +def test_parse_annotations_handles_escaped_quotes(): + text = '@note("say \\"hello\\"") forbid (principal, action, resource);' + assert parse_policy_annotations(text) == {"policy0": {"note": 'say "hello"'}} + + +def test_parse_annotations_empty_for_unannotated_bundle(): + assert parse_policy_annotations("permit (principal, action, resource);") == {} + + +# ── PolicyEvaluator advice flow ─────────────────────────────────────────────── + + +def test_enforcing_deny_carries_advice(): + evaluator = PolicyEvaluator( + bundle=_bundle(HITL_POLICY), config=_config(EnforcementMode.ENFORCING) + ) + with pytest.raises(PolicyDeny) as exc_info: + evaluator.evaluate(_context("high")) + assert exc_info.value.advice["reason"] == "human-review-required" + assert exc_info.value.advice["regulation"] == "eu-ai-act-art-14" + assert exc_info.value.advice["reviewer_role"] == "attending-physician" + + +def test_advisory_deny_carries_advice(): + evaluator = PolicyEvaluator( + bundle=_bundle(HITL_POLICY), config=_config(EnforcementMode.ADVISORY) + ) + decision = evaluator.evaluate(_context("high")) + assert decision.allowed is True + assert decision.would_have_denied is True + assert decision.advice["reason"] == "human-review-required" + + +def test_allow_has_no_advice(): + evaluator = PolicyEvaluator( + bundle=_bundle(HITL_POLICY), config=_config(EnforcementMode.ENFORCING) + ) + decision = evaluator.evaluate(_context("standard")) + assert decision.allowed is True + assert decision.advice == {} + + +def test_deny_without_annotations_has_empty_advice(): + deny_all = 'forbid (principal, action, resource);' + evaluator = PolicyEvaluator( + bundle=_bundle(deny_all), config=_config(EnforcementMode.ENFORCING) + ) + with pytest.raises(PolicyDeny) as exc_info: + evaluator.evaluate(_context("standard")) + assert exc_info.value.advice == {} diff --git a/tests/unit/test_session_close.py b/tests/unit/test_session_close.py new file mode 100644 index 00000000..c8d99ab7 --- /dev/null +++ b/tests/unit/test_session_close.py @@ -0,0 +1,105 @@ +"""Tests for POST /sessions/{id}/close — claim issuance and session rotation.""" + +from __future__ import annotations + +from datetime import UTC, datetime +from unittest.mock import MagicMock, patch + +import pytest +from starlette.testclient import TestClient + +from cmcp_runtime.audit.keys import SigningKey +from cmcp_runtime.cli import build_server +from cmcp_runtime.config import AttestationConfig, Config +from cmcp_runtime.policy.bundle import PolicyStore +from cmcp_runtime.startup import RuntimeContext + + +@pytest.fixture +def server(): + config = Config(attestation=AttestationConfig(), dev_mode=True) + + attestation_report = MagicMock() + attestation_report.provider = "software-only" + attestation_report.attestation_generated_at = datetime.now(UTC) + attestation_report.attestation_validity_seconds = 86400 + attestation_report.measurement = "0" * 64 + attestation_report.report_data = "0" * 64 + attestation_report.measurement_note = None + attestation_report.raw_evidence = None + + bundle = MagicMock() + bundle.bundle_hash = "sha256:" + "0" * 64 + bundle.policy_files = {"allow.cedar": "permit (principal, action, resource);"} + bundle.manifest = MagicMock() + bundle.manifest.version = "test-v1" + policy_store = MagicMock(spec=PolicyStore) + policy_store.bundle = bundle + + catalog = MagicMock() + catalog.entries = {} + catalog.catalog_hash = "sha256:" + "1" * 64 + catalog.exceptions = [] + + ctx = RuntimeContext( + config=config, + tee_provider=MagicMock(), + attestation_report=attestation_report, + signing_key=SigningKey(), + policy_bundle=policy_store, + catalog=catalog, + ) + return build_server(ctx) + + +def test_close_returns_signed_claim_and_rotates(server): + client = TestClient(server.app) + old_session_id = server._session.session_id + + resp = client.post(f"/sessions/{old_session_id}/close") + assert resp.status_code == 200 + claim = resp.json() + assert claim["gateway"]["session_id"] == old_session_id + assert claim["signature"] # signed claim + + # Session rotated: new id, proxy rebound. + assert server._session.session_id != old_session_id + assert server._proxy._session.session_id == server._session.session_id + assert server._audit_chain is server._proxy._audit + + +def test_closed_claim_retrievable_via_trace_claim_endpoint(server): + client = TestClient(server.app) + session_id = server._session.session_id + client.post(f"/sessions/{session_id}/close") + + resp = client.get(f"/sessions/{session_id}/trace-claim") + assert resp.status_code == 200 + assert resp.json()["gateway"]["session_id"] == session_id + + +def test_close_unknown_session_404(server): + client = TestClient(server.app) + resp = client.post("/sessions/not-a-real-session/close") + assert resp.status_code == 404 + + +def test_close_twice_404_on_second(server): + client = TestClient(server.app) + session_id = server._session.session_id + assert client.post(f"/sessions/{session_id}/close").status_code == 200 + assert client.post(f"/sessions/{session_id}/close").status_code == 404 + + +def test_audit_export_serves_closed_session(server): + client = TestClient(server.app) + session_id = server._session.session_id + client.post(f"/sessions/{session_id}/close") + + resp = client.get(f"/audit/export?session_id={session_id}") + assert resp.status_code == 200 + bundle = resp.json() + assert bundle["session_id"] == session_id + entry_types = [e["entry_type"] for e in bundle["entries"]] + assert "session_start" in entry_types + assert "session_end" in entry_types diff --git a/tests/unit/test_session_reset.py b/tests/unit/test_session_reset.py index 306d8239..6c8faa90 100644 --- a/tests/unit/test_session_reset.py +++ b/tests/unit/test_session_reset.py @@ -9,7 +9,7 @@ from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch from starlette.testclient import TestClient @@ -24,6 +24,7 @@ from cmcp_runtime.mcp.server import MCPServer from cmcp_runtime.policy.evaluator import PolicyDecision, PolicyEvaluator from cmcp_runtime.session.state import SessionState +from tests.unit.conftest import wire_mock_gateway # ── Helpers ─────────────────────────────────────────────────────────────────── @@ -82,10 +83,7 @@ def _make_server(session_id: str = "sess-reset-001"): with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): proxy = CMCPProxy(cat, ev, session, chain, cfg) - proxy._mcp_gateway = MagicMock() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=[], injection_detected=False - )) + wire_mock_gateway(proxy) server = MCPServer(proxy, session=session, audit_chain=chain) return server, session, chain diff --git a/tests/unit/test_upstream_forwarding.py b/tests/unit/test_upstream_forwarding.py new file mode 100644 index 00000000..000b21cc --- /dev/null +++ b/tests/unit/test_upstream_forwarding.py @@ -0,0 +1,125 @@ +"""Tests for CMCPProxy._forward_to_upstream against a real local HTTP server.""" + +from __future__ import annotations + +import json +import threading +from http.server import BaseHTTPRequestHandler, HTTPServer +from unittest.mock import MagicMock, patch + +import pytest + +from cmcp_runtime.audit.chain import AuditChain +from cmcp_runtime.catalog.loader import ( + ApprovedDefinition, + CatalogEntry, + ServerIdentity, + ToolCatalog, +) +from cmcp_runtime.config import AttestationConfig, Config, EnforcementMode +from cmcp_runtime.errors import UpstreamToolError, UpstreamUnavailable +from cmcp_runtime.session.state import SessionState + + +class _MockMCPHandler(BaseHTTPRequestHandler): + """Serves canned JSON-RPC responses; behavior keyed on tool name.""" + + def do_POST(self): # noqa: N802 + length = int(self.headers.get("Content-Length", 0)) + request = json.loads(self.rfile.read(length)) + tool = request["params"]["name"] + if tool == "test.fail": + body = { + "jsonrpc": "2.0", + "id": request["id"], + "error": {"code": -32000, "message": "tool exploded"}, + } + else: + body = { + "jsonrpc": "2.0", + "id": request["id"], + "result": { + "content": [{"type": "text", "text": f"echo:{tool}"}] + }, + } + payload = json.dumps(body).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, *args): # silence test output + pass + + +@pytest.fixture(scope="module") +def mock_server(): + server = HTTPServer(("127.0.0.1", 0), _MockMCPHandler) + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + yield f"http://127.0.0.1:{server.server_port}/mcp" + server.shutdown() + + +def _make_proxy(server_url: str): + from cmcp_runtime.mcp.proxy import CMCPProxy + + entry = CatalogEntry( + tool_name="test.echo", + server=ServerIdentity( + display_name="Mock", + url=server_url, + tls_fingerprint="SHA256:" + "A" * 43 + "=", + spiffe_id=None, + transport="http-sse", + rotation_mode="key-pinned", + ), + approved_definition=ApprovedDefinition( + description="echo", input_schema={"type": "object"}, output_schema=None + ), + definition_hash="sha256:" + "0" * 64, + compliance_domain="public", + requires_baa=False, + sensitivity_level="public", + added_at="2026-06-10T00:00:00Z", + approved_by="test", + ) + catalog = ToolCatalog( + entries={"test.echo": entry, "test.fail": entry}, + catalog_hash="sha256:" + "1" * 64, + ) + config = Config(attestation=AttestationConfig(enforcement_mode=EnforcementMode.ENFORCING)) + session = SessionState(session_id="fwd-test") + chain = AuditChain("fwd-test") + with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ + patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): + proxy = CMCPProxy( + catalog=catalog, + policy_evaluator=MagicMock(), + session=session, + audit_chain=chain, + config=config, + ) + return proxy, entry + + +@pytest.mark.asyncio +async def test_forward_returns_text_content(mock_server): + proxy, entry = _make_proxy(mock_server) + text = await proxy._forward_to_upstream("c1", entry, "test.echo", {"message": "hi"}) + assert text == "echo:test.echo" + + +@pytest.mark.asyncio +async def test_forward_raises_on_jsonrpc_error(mock_server): + proxy, entry = _make_proxy(mock_server) + with pytest.raises(UpstreamToolError, match="tool exploded"): + await proxy._forward_to_upstream("c1", entry, "test.fail", {}) + + +@pytest.mark.asyncio +async def test_forward_raises_when_unreachable(): + proxy, entry = _make_proxy("http://127.0.0.1:1/mcp") + with pytest.raises(UpstreamUnavailable): + await proxy._forward_to_upstream("c1", entry, "test.echo", {}) diff --git a/tests/unit/test_workflow_scope.py b/tests/unit/test_workflow_scope.py index 73162242..f6f508e3 100644 --- a/tests/unit/test_workflow_scope.py +++ b/tests/unit/test_workflow_scope.py @@ -2,7 +2,7 @@ from __future__ import annotations -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock, patch import pytest @@ -16,6 +16,7 @@ from cmcp_runtime.config import AttestationConfig, Config, EnforcementMode from cmcp_runtime.policy.evaluator import PolicyDecision, PolicyEvaluator from cmcp_runtime.session.state import SessionState +from tests.unit.conftest import wire_mock_gateway # ── Helpers ─────────────────────────────────────────────────────────────────── @@ -79,10 +80,7 @@ def _make_proxy(evaluator=None): with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): proxy = CMCPProxy(cat, ev, session, chain, cfg) - proxy._mcp_gateway = MagicMock() - proxy._mcp_gateway.call_tool = AsyncMock(return_value=MagicMock( - sensitivity_tags=[], injection_detected=False - )) + wire_mock_gateway(proxy) return proxy, ev, chain