From d77a1b7ce4eed53200ea448239cd0e2b6a6d6609 Mon Sep 17 00:00:00 2001 From: Imran Siddique Date: Sat, 6 Jun 2026 12:15:57 -0700 Subject: [PATCH] feat(session): implement session lifecycle and audit export (#60, #55) Add SessionManager (create/close/get) with signed TRACE Claim generation on session close, staleness detection for attestation reports, and a signed audit bundle export endpoint. Extend MCPServer with optional session_manager and audit_chain kwargs to serve GET /sessions/{id}/trace-claim and GET /audit/export with chain integrity verification before export. Co-Authored-By: Claude Sonnet 4.6 --- src/cmcp_gateway/mcp/server.py | 64 ++++++- src/cmcp_gateway/session/manager.py | 203 +++++++++++++++++++++++ tests/unit/test_session_manager.py | 247 ++++++++++++++++++++++++++++ 3 files changed, 512 insertions(+), 2 deletions(-) create mode 100644 src/cmcp_gateway/session/manager.py create mode 100644 tests/unit/test_session_manager.py diff --git a/src/cmcp_gateway/mcp/server.py b/src/cmcp_gateway/mcp/server.py index 37d56b5a..b59adc9f 100644 --- a/src/cmcp_gateway/mcp/server.py +++ b/src/cmcp_gateway/mcp/server.py @@ -13,7 +13,7 @@ import json import logging import uuid -from typing import Any +from typing import TYPE_CHECKING, Any from agent_os.stateless import StatelessKernel from starlette.applications import Starlette @@ -23,6 +23,10 @@ from cmcp_gateway.mcp.proxy import CMCPProxy +if TYPE_CHECKING: + from cmcp_gateway.audit.chain import AuditChain + from cmcp_gateway.session.manager import SessionManager + logger = logging.getLogger(__name__) @@ -34,13 +38,27 @@ class MCPServer: The proxy routes calls to upstream servers based on the attested catalog. """ - def __init__(self, proxy: CMCPProxy) -> None: + def __init__( + self, + proxy: CMCPProxy, + *, + session_manager: SessionManager | None = None, + audit_chain: AuditChain | None = None, + ) -> None: self._proxy = proxy + self._session_manager = session_manager + self._audit_chain = audit_chain self._kernel = StatelessKernel() self.app = Starlette(routes=[ Route("/mcp", self._handle_mcp, methods=["POST"]), Route("/health", self._health, methods=["GET"]), Route("/tools/list", self._list_tools, methods=["GET"]), + Route( + "/sessions/{session_id}/trace-claim", + self._get_trace_claim, + methods=["GET"], + ), + Route("/audit/export", self._audit_export, methods=["GET"]), ]) async def _handle_mcp(self, request: Request) -> Response: @@ -176,3 +194,45 @@ async def _list_tools(self, request: Request) -> Response: async def _health(self, request: Request) -> Response: return JSONResponse({"status": "ok"}) + + async def _get_trace_claim(self, request: Request) -> Response: + """GET /sessions/{session_id}/trace-claim — returns signed TRACE Claim for a closed session.""" + if self._session_manager is None: + return JSONResponse( + {"error": "session management not available"}, status_code=501 + ) + session_id: str = request.path_params["session_id"] + claim = self._session_manager.get_trace_claim(session_id) + if claim is None: + return JSONResponse( + {"error": f"trace claim not found for session_id={session_id}"}, + status_code=404, + ) + return JSONResponse(claim) + + async def _audit_export(self, request: Request) -> Response: + """GET /audit/export?session_id= — returns signed audit bundle.""" + if self._session_manager is None or self._audit_chain is None: + return JSONResponse( + {"error": "audit export not available"}, status_code=501 + ) + session_id: str | None = request.query_params.get("session_id") + if not session_id: + return JSONResponse( + {"error": "query parameter 'session_id' is required"}, + status_code=400, + ) + try: + bundle = self._session_manager.get_audit_bundle( + session_id, self._audit_chain + ) + except ValueError as exc: + logger.error( + "Audit chain integrity failure: session_id=%s error=%s", + session_id, + exc, + ) + return JSONResponse( + {"error": "audit chain integrity check failed"}, status_code=500 + ) + return JSONResponse(bundle) diff --git a/src/cmcp_gateway/session/manager.py b/src/cmcp_gateway/session/manager.py new file mode 100644 index 00000000..e8d478d4 --- /dev/null +++ b/src/cmcp_gateway/session/manager.py @@ -0,0 +1,203 @@ +"""Session lifecycle management — implements issues #60 and #55.""" + +from __future__ import annotations + +import base64 +import hashlib +import json +import logging +from dataclasses import asdict +from datetime import UTC, datetime +from typing import Any +from uuid import uuid4 + +from cmcp_gateway.audit.chain import AuditChain +from cmcp_gateway.audit.trace_claim import ( + AttestationReportInfo, + CallGraphSummary, + CallSummary, + PolicyBundleInfo, + ToolCatalogInfo, + generate_trace_claim, +) +from cmcp_gateway.session.state import SessionState +from cmcp_gateway.startup import GatewayContext + +logger = logging.getLogger(__name__) + + +class SessionManager: + """Creates, tracks, and closes agent sessions.""" + + def __init__(self, ctx: GatewayContext) -> None: + self._ctx = ctx + # Stores signed claim dicts keyed by session_id, populated on close. + self._closed_claims: dict[str, dict[str, Any]] = {} + + def create_session(self) -> tuple[SessionState, AuditChain]: + """Create a new session. Returns (state, chain).""" + session_id = str(uuid4()) + state = SessionState(session_id=session_id) + chain = AuditChain(session_id=session_id) + logger.info("Session created: session_id=%s", session_id) + return state, chain + + def close_session( + self, session_id: str, state: SessionState, chain: AuditChain + ) -> dict[str, Any]: + """ + Close a session: + 1. Append a session_end audit entry to the chain. + 2. Build the GatewayClaim from chain + state + ctx. + 3. Sign it with ctx.signing_key. + 4. Store the signed claim JSON, keyed by session_id. + 5. Return the signed claim dict. + """ + chain.append( + "session_end", + session_sensitivity_before=state.max_sensitivity, + session_sensitivity_after=state.max_sensitivity, + ) + + ctx = self._ctx + report = ctx.attestation_report + + # Convert AttestationReport (datetime) to AttestationReportInfo (str). + generated_at_str = report.attestation_generated_at.isoformat() + age_seconds = ( + datetime.now(UTC) - report.attestation_generated_at + ).total_seconds() + attestation_stale = age_seconds > report.attestation_validity_seconds + + attestation_info = AttestationReportInfo( + provider=report.provider, + measurement=report.measurement, + report_data=report.report_data, + attestation_generated_at=generated_at_str, + attestation_validity_seconds=report.attestation_validity_seconds, + measurement_note=report.measurement_note, + raw_evidence=( + base64.urlsafe_b64encode(report.raw_evidence).rstrip(b"=").decode() + if report.raw_evidence is not None + else None + ), + ) + + bundle = ctx.policy_bundle + policy_info = PolicyBundleInfo( + hash=bundle.bundle_hash, + enforcement_mode=str(ctx.config.attestation.enforcement_mode), + policy_version=bundle.manifest.version, + ) + + catalog = ctx.catalog + # Detect catalog exceptions — entries where catalog_exception=True. + catalog_exceptions: list[dict[str, str]] = [ + {"tool_name": name} + for name, entry in catalog.entries.items() + if entry.catalog_exception + ] + catalog_info = ToolCatalogInfo( + hash=catalog.catalog_hash, + drift_detected=False, + ) + + # Build call summary from chain entries. + entries = chain.entries + tool_calls = [e for e in entries if e.entry_type == "tool_call"] + tool_calls_total = len(tool_calls) + tool_calls_allowed = sum( + 1 for e in tool_calls if e.policy_decision == "allow" + ) + tool_calls_denied = sum( + 1 for e in tool_calls if e.policy_decision in ("deny", "advisory_deny") + ) + tool_calls_faulted = sum( + 1 for e in tool_calls if e.policy_decision == "fault" + ) + tools_invoked = sorted( + {e.tool_name for e in tool_calls if e.tool_name is not None} + ) + + # Identify compliance domains from catalog entries touched. + compliance_domains_touched = sorted( + { + catalog.entries[name].compliance_domain + for name in tools_invoked + if name in catalog.entries + } + ) + + call_summary = CallSummary( + tool_calls_total=tool_calls_total, + tool_calls_allowed=tool_calls_allowed, + tool_calls_denied=tool_calls_denied, + tool_calls_faulted=tool_calls_faulted, + tools_invoked=tools_invoked, + session_max_sensitivity=state.max_sensitivity, + call_graph_summary=CallGraphSummary( + compliance_domains_touched=compliance_domains_touched, + cross_boundary_events=[], + ), + ) + + claim = generate_trace_claim( + session_id=session_id, + signing_key=ctx.signing_key, + attestation_report=attestation_info, + policy_bundle=policy_info, + tool_catalog=catalog_info, + call_summary=call_summary, + audit_chain_root=chain.chain_root, + audit_chain_tip=chain.chain_tip, + audit_chain_length=chain.length, + attestation_stale=attestation_stale, + catalog_exceptions=catalog_exceptions, + do_sign=True, + ) + + claim_dict = claim.model_dump(exclude_none=True) + self._closed_claims[session_id] = claim_dict + logger.info("Session closed: session_id=%s", session_id) + return claim_dict + + def get_trace_claim(self, session_id: str) -> dict[str, Any] | None: + """Return the signed TRACE Claim for a closed session.""" + return self._closed_claims.get(session_id) + + def get_audit_bundle( + self, session_id: str, chain: AuditChain + ) -> dict[str, Any]: + """ + Build a signed audit bundle for export (issue #55): + { + "session_id": ..., + "entries": [list of entry dicts from chain], + "bundle_signature": base64url(sha256(canonical_json(entries)) signed with signing_key) + } + + Raises ValueError if the chain is broken (verify_chain() fails). + """ + if not chain.verify_chain(): + raise ValueError( + f"Audit chain integrity check failed for session_id={session_id}" + ) + + entries_dicts = [asdict(e) for e in chain.entries] + canonical = json.dumps( + entries_dicts, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode() + digest = hashlib.sha256(canonical).digest() + raw_sig = self._ctx.signing_key.sign(digest) + bundle_signature = ( + base64.urlsafe_b64encode(raw_sig).rstrip(b"=").decode() + ) + + return { + "session_id": session_id, + "entries": entries_dicts, + "bundle_signature": bundle_signature, + } diff --git a/tests/unit/test_session_manager.py b/tests/unit/test_session_manager.py new file mode 100644 index 00000000..06307922 --- /dev/null +++ b/tests/unit/test_session_manager.py @@ -0,0 +1,247 @@ +"""Unit tests for SessionManager (issues #60 and #55).""" + +from __future__ import annotations + +import base64 +import hashlib +import json +from datetime import UTC, datetime +from unittest.mock import MagicMock + +import pytest + +from cmcp_gateway.audit.chain import AuditChain +from cmcp_gateway.audit.keys import SigningKey +from cmcp_gateway.session.manager import SessionManager +from cmcp_gateway.session.state import SessionState + +# ── Fixtures ────────────────────────────────────────────────────────────────── + + +def _make_attestation_report(*, stale: bool = False) -> MagicMock: + """Return a mock AttestationReport with software-only provider values.""" + report = MagicMock() + report.provider = "software-only" + report.measurement = "DEVELOPMENT_ONLY_NOT_FOR_PRODUCTION" + report.report_data = "aa" * 32 + report.raw_evidence = None + report.measurement_note = "software-only mode — not hardware-backed" + report.attestation_validity_seconds = 86400 + if stale: + # Set generated_at far in the past so the report is expired. + report.attestation_generated_at = datetime(2020, 1, 1, tzinfo=UTC) + else: + report.attestation_generated_at = datetime.now(UTC) + return report + + +def _make_ctx(*, stale_attestation: bool = False) -> MagicMock: + """Return a fully-wired mock GatewayContext.""" + signing_key = SigningKey() + + policy_bundle = MagicMock() + policy_bundle.bundle_hash = "sha256:" + "a" * 64 + policy_bundle.manifest.version = "1.0.0" + + catalog_entry = MagicMock() + catalog_entry.compliance_domain = "external" + catalog_entry.catalog_exception = False + + catalog = MagicMock() + catalog.catalog_hash = "sha256:" + "b" * 64 + catalog.entries = {} + + config = MagicMock() + config.attestation.enforcement_mode = "enforcing" + + ctx = MagicMock() + ctx.signing_key = signing_key + ctx.attestation_report = _make_attestation_report(stale=stale_attestation) + ctx.policy_bundle = policy_bundle + ctx.catalog = catalog + ctx.config = config + return ctx + + +# ── create_session ───────────────────────────────────────────────────────────── + + +def test_create_session_returns_state_and_chain() -> None: + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + assert isinstance(state, SessionState) + assert isinstance(chain, AuditChain) + + +def test_create_session_ids_match() -> None: + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + assert state.session_id == chain.entries[0].session_id + + +def test_create_session_produces_unique_ids() -> None: + mgr = SessionManager(_make_ctx()) + state1, _ = mgr.create_session() + state2, _ = mgr.create_session() + assert state1.session_id != state2.session_id + + +def test_create_session_chain_has_session_start() -> None: + mgr = SessionManager(_make_ctx()) + _, chain = mgr.create_session() + assert chain.entries[0].entry_type == "session_start" + + +# ── close_session ───────────────────────────────────────────────────────────── + + +def test_close_session_produces_gateway_claim() -> None: + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + claim = mgr.close_session(state.session_id, state, chain) + assert claim["cmcp_version"] == "1.0" + assert "trace" in claim + assert "gateway" in claim + assert "signature" in claim + + +def test_close_session_claim_is_signed() -> None: + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + claim = mgr.close_session(state.session_id, state, chain) + assert len(claim["signature"]) > 0 + + +def test_close_session_appends_session_end_entry() -> None: + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + mgr.close_session(state.session_id, state, chain) + entry_types = [e.entry_type for e in chain.entries] + assert "session_end" in entry_types + + +def test_close_session_stores_claim_by_session_id() -> None: + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + claim = mgr.close_session(state.session_id, state, chain) + retrieved = mgr.get_trace_claim(state.session_id) + assert retrieved == claim + + +def test_close_session_claim_session_id_matches() -> None: + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + claim = mgr.close_session(state.session_id, state, chain) + assert claim["gateway"]["session_id"] == state.session_id + + +def test_close_session_attestation_stale_flag_false_when_fresh() -> None: + mgr = SessionManager(_make_ctx(stale_attestation=False)) + state, chain = mgr.create_session() + claim = mgr.close_session(state.session_id, state, chain) + assert claim["gateway"]["attestation_stale"] is False + + +def test_close_session_attestation_stale_flag_true_when_expired() -> None: + mgr = SessionManager(_make_ctx(stale_attestation=True)) + state, chain = mgr.create_session() + claim = mgr.close_session(state.session_id, state, chain) + assert claim["gateway"]["attestation_stale"] is True + + +def test_close_session_signature_verifiable() -> None: + """Signature on the claim must verify against the embedded JWK public key.""" + from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey + + from cmcp_gateway.audit.trace_claim import GatewayClaim, _to_dict, canonical_json + + ctx = _make_ctx() + mgr = SessionManager(ctx) + state, chain = mgr.create_session() + claim_dict = mgr.close_session(state.session_id, state, chain) + + # Re-validate through pydantic to get the proper model + claim = GatewayClaim.model_validate(claim_dict) + body = canonical_json(_to_dict(claim)) + sig_bytes = base64.urlsafe_b64decode(claim.signature + "==") + + pub = Ed25519PublicKey.from_public_bytes(ctx.signing_key.public_key_bytes) + pub.verify(sig_bytes, body) # raises InvalidSignature if wrong + + +# ── get_trace_claim ──────────────────────────────────────────────────────────── + + +def test_get_trace_claim_returns_none_for_unknown_session() -> None: + mgr = SessionManager(_make_ctx()) + assert mgr.get_trace_claim("nonexistent-session-id") is None + + +def test_get_trace_claim_returns_claim_for_closed_session() -> None: + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + mgr.close_session(state.session_id, state, chain) + result = mgr.get_trace_claim(state.session_id) + assert result is not None + assert result["cmcp_version"] == "1.0" + + +# ── get_audit_bundle ────────────────────────────────────────────────────────── + + +def test_audit_bundle_contains_required_keys() -> None: + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + bundle = mgr.get_audit_bundle(state.session_id, chain) + assert set(bundle.keys()) == {"session_id", "entries", "bundle_signature"} + + +def test_audit_bundle_session_id_matches() -> None: + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + bundle = mgr.get_audit_bundle(state.session_id, chain) + assert bundle["session_id"] == state.session_id + + +def test_audit_bundle_entries_count_matches_chain() -> None: + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + chain.append("tool_call", call_id="c1", tool_name="t", policy_decision="allow") + bundle = mgr.get_audit_bundle(state.session_id, chain) + assert len(bundle["entries"]) == chain.length + + +def test_audit_bundle_signature_is_valid() -> None: + """Bundle signature must verify: sign(sha256(canonical_json(entries))).""" + from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey + + ctx = _make_ctx() + mgr = SessionManager(ctx) + state, chain = mgr.create_session() + bundle = mgr.get_audit_bundle(state.session_id, chain) + + entries_dicts = bundle["entries"] + canonical = json.dumps( + entries_dicts, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=True, + ).encode() + digest = hashlib.sha256(canonical).digest() + + sig_bytes = base64.urlsafe_b64decode(bundle["bundle_signature"] + "==") + pub = Ed25519PublicKey.from_public_bytes(ctx.signing_key.public_key_bytes) + pub.verify(sig_bytes, digest) # raises InvalidSignature if wrong + + +def test_audit_bundle_broken_chain_raises_value_error() -> None: + """If the audit chain is tampered with, get_audit_bundle must raise ValueError.""" + mgr = SessionManager(_make_ctx()) + state, chain = mgr.create_session() + chain.append("tool_call", call_id="c1", tool_name="t", policy_decision="allow") + + # Tamper with the chain. + chain.entries[0].tool_name = "injected_tool" + + with pytest.raises(ValueError, match="integrity check failed"): + mgr.get_audit_bundle(state.session_id, chain)