From 574ab75f41540fcc59c93ce63fa5046231cac598 Mon Sep 17 00:00:00 2001 From: Imran Siddique Date: Thu, 10 Sep 2026 16:13:49 -0700 Subject: [PATCH 1/2] fix(session): hold the session boundary a credentialed reset draws Found while reviewing the Monotonic Session Sensitivity State draft against the runtime. A reset is supposed to close one session and open a successor at the minimum level. Four things let the boundary leak. A response arriving during a reset raised the successor. The mutation lock serialised the reset and the response elevation but did not order them, so whichever coroutine acquired it second won: a response in flight when the reset landed was applied to the successor and recorded its pre-reset call_id as the raiser. update_from_inspection() now takes the reset_count observed at call entry and drops a response whose generation no longer matches. The discriminator is reset_count and not session_id because upgrade_attestation() rotates the identifier while deliberately continuing the same session, so a call in flight across an attestation upgrade must still apply. The previous concurrency test asserted only that max_sensitivity stayed a member of SENSITIVITY_ORDER, which every value satisfies. The reset route accepted the tool-invocation token. POST /sessions/{id}/reset is not reachable as an MCP tool, but it sat behind the same CMCP_BEARER_TOKEN as POST /mcp, so an agent host holding its own credential could clear accumulated sensitivity. The operator interface now takes CMCP_OPERATOR_TOKEN, required outside dev mode and refused if it equals the bearer token. Unset, those routes still fall back, so existing deployments keep working until they set it. The reset audit entry recorded the sensitivity transition and nothing that identified the boundary. detail now carries the closed session id, the successor id, the resulting reset counter and which credential was verified. detail is inside the canonical body, so those fields are hash-covered. The chain attributed every entry after a reset to the closed session, and the successor's identifier appeared nowhere in it. rotate_session_id() moves attribution after the boundary entry is written. The closed session's final value is now preserved as a distinct ClosedSessionRecord rather than overwritten, and the reset response returns it. Nine tests added. Two suite failures are pre-existing on the base: test_release_distribution_smoke (version 0.5.0 vs metadata 0.4.0) and test_tpm_chained_verify (assertion text drift). Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01CDzMWcn12dszchgrKfXzZY --- CHANGELOG.md | 48 +++++++++++ docs/configuration.md | 1 + src/cmcp_runtime/audit/chain.py | 12 +++ src/cmcp_runtime/cli.py | 1 + src/cmcp_runtime/config.py | 16 ++++ src/cmcp_runtime/mcp/proxy.py | 17 +++- src/cmcp_runtime/mcp/server.py | 67 +++++++++++++-- src/cmcp_runtime/session/state.py | 63 +++++++++++++- src/cmcp_runtime/startup.py | 13 +++ tests/unit/test_config.py | 28 ++++++ tests/unit/test_session.py | 50 +++++++++++ tests/unit/test_session_reset.py | 138 ++++++++++++++++++++++++++++++ 12 files changed, 444 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c59bf8a..95449f1e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,54 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Security + +- **A response arriving during an operator reset raised the successor session.** + The per-session mutation lock serialised a reset and a response elevation but + did not order them, so whichever coroutine acquired it second won. A response + in flight when the reset landed was applied to the successor, which had just + been initialised to `public`, and recorded the pre-reset `call_id` as the call + that raised it. The successor exists to start at the minimum level, so this + carried the closed session's sensitivity across the boundary the reset drew. + `update_from_inspection()` now takes the `reset_count` observed at call entry + and drops a response whose generation no longer matches, logging + `SESSION_RESET_RACE`. The discriminator is `reset_count` rather than + `session_id` because `upgrade_attestation()` rotates the identifier while + deliberately continuing the same session, so a call in flight across an + attestation upgrade must still apply. The previous concurrency test asserted + only that `max_sensitivity` remained a member of `SENSITIVITY_ORDER`, which + every value satisfies. + +- **The reset route accepted the tool-invocation token.** `POST + /sessions/{id}/reset` is not reachable as an MCP tool, but it sat behind the + same single `CMCP_BEARER_TOKEN` as `POST /mcp`, so an agent host holding its + own tool-invocation credential could clear accumulated session sensitivity. + The operator interface (session reset and catalog exception) now takes + `CMCP_OPERATOR_TOKEN`, which must differ from `CMCP_BEARER_TOKEN` and is + required outside `CMCP_DEV_MODE=1` (`OPERATOR_TOKEN_REQUIRED`). Where it is + unset those routes still fall back to the bearer token, so an existing + single-token deployment keeps working until it sets the new variable. + +### Changed + +- The reset audit entry now identifies the session boundary rather than only the + sensitivity transition: `detail` carries the closed session identifier, the + successor identifier, the resulting reset counter, and which credential was + verified. `detail` is inside the canonical body, so those fields are covered by + the entry hash. + +- **The audit chain no longer attributes post-reset entries to the closed + session.** `AuditChain.rotate_session_id()` moves attribution to the successor + after the boundary entry is written, so the reset entry belongs to the session + that reached the recorded value and later entries belong to the successor. + Previously every entry after a reset carried the closed session's identifier + and the successor's identifier appeared nowhere in the chain. + +- A reset now preserves the closed session's final state as a distinct + `ClosedSessionRecord` instead of overwriting it, and + `POST /sessions/{id}/reset` returns `closed_session_max_sensitivity` and + `reset_count`. + ## [0.5.0] - 2026-09-05 ### Security diff --git a/docs/configuration.md b/docs/configuration.md index f07f495f..a83d6e4f 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -135,6 +135,7 @@ Environment variables control secrets and mode flags that must not appear in con |----------|-------------|-----------| | `CMCP_DEV_MODE=1` | Enables software-only attestation. No hardware TEE required. TRACE Claims will show `partially_verified` status. Required when `provider` is `software-only`. | `attestation.provider` (forces software-only) | | `CMCP_BEARER_TOKEN` | Optional bearer token for runtime HTTP auth. If set, all requests to the runtime must include `Authorization: Bearer `. If unset, no bearer auth is enforced. This token is required for non-loopback binds. | none | +| `CMCP_OPERATOR_TOKEN` | Credential for the operator interface: `POST /sessions/{id}/reset` and `POST /catalog/exception`. Required outside `CMCP_DEV_MODE=1` (`OPERATOR_TOKEN_REQUIRED`), and must differ from `CMCP_BEARER_TOKEN`. When set, those two routes accept only this token and reject the tool-invocation token; when unset they fall back to `CMCP_BEARER_TOKEN`. A reset lowers accumulated session sensitivity, so an agent host holding only the tool-invocation token cannot clear the state that monotonicity exists to keep. | none | | `OPAQUE_ATTESTATION_URL` | Enables the OPAQUE Managed Runtime provider. Must be set to the OPAQUE attestation service URL. Required when `provider` is `opaque` or `auto` on OPAQUE infrastructure. | enables `opaque` provider detection | | `CMCP_POLICY_HASH` | SHA-256 hash of the approved policy bundle. Required in non-dev mode and checked by startup before Agent Manifest binding. The gateway fails closed at startup if this is unset and `CMCP_DEV_MODE` is not `1`. Format: `sha256:`. | none (startup policy integrity check) | | `CMCP_CATALOG_HASH` | SHA-256 hash of the approved `catalog.json`. Required in non-dev mode. The gateway fails closed at startup if this is unset and `CMCP_DEV_MODE` is not `1`. Format: `sha256:`. | none (additional startup check) | diff --git a/src/cmcp_runtime/audit/chain.py b/src/cmcp_runtime/audit/chain.py index 1be15c35..02886efc 100644 --- a/src/cmcp_runtime/audit/chain.py +++ b/src/cmcp_runtime/audit/chain.py @@ -266,6 +266,18 @@ def append( self._notify_sinks(entry) return entry + def rotate_session_id(self, new_session_id: str) -> None: + """Attribute subsequent entries to ``new_session_id``. + + A credentialed reset closes one session and opens a successor on the same + hash-linked chain. Without this, every entry after a reset carries the + closed session's identifier and the successor's identifier appears + nowhere, so the record cannot say which session an entry belongs to. + Entries already appended are unchanged: they are hashed and remain + attributed to the session that produced them. + """ + self._session_id = new_session_id + def add_sink(self, sink: Callable[[AuditEntry], None]) -> None: """Register a read-only observer of appended entries. See __init__.""" self._sinks.append(sink) diff --git a/src/cmcp_runtime/cli.py b/src/cmcp_runtime/cli.py index e1f91f62..db99c156 100644 --- a/src/cmcp_runtime/cli.py +++ b/src/cmcp_runtime/cli.py @@ -95,6 +95,7 @@ def build_server(ctx: RuntimeContext) -> MCPServer: audit_chain=audit_chain, session=session, bearer_token=ctx.config.bearer_token, + operator_token=ctx.config.operator_token, ) diff --git a/src/cmcp_runtime/config.py b/src/cmcp_runtime/config.py index dae935e6..7fd45976 100644 --- a/src/cmcp_runtime/config.py +++ b/src/cmcp_runtime/config.py @@ -134,6 +134,13 @@ class Config: audit_db_path: str = "audit.db" # AUDIT-001: durable audit chain storage dev_mode: bool = False bearer_token: str | None = None + #: Credential for the operator interface (session reset, catalog exception). + #: Held separately from ``bearer_token`` so that the credential authorizing a + #: reset is not the credential an agent host already holds to invoke tools. + #: A reset lowers accumulated session sensitivity, so an agent able to + #: present its own tool-invocation token to the reset route could clear the + #: state that monotonicity exists to keep. + operator_token: str | None = None #: AARM R6. A named conformance profile tightens defaults that stay #: permissive for developers. None is the default, and nothing changes. #: "aarm" requires an Agent Manifest binding, because R6 says every receipt @@ -446,6 +453,14 @@ def load_config(path: str) -> Config: dev_mode = DEV_MODE # TEE-002: use the frozen constant, never re-read from env bearer_token = os.environ.get("CMCP_BEARER_TOKEN") or None + operator_token = os.environ.get("CMCP_OPERATOR_TOKEN") or None + + if operator_token is not None and operator_token == bearer_token: + raise ConfigError( + "CMCP_OPERATOR_TOKEN must differ from CMCP_BEARER_TOKEN. The operator " + "credential authorizes a session-sensitivity reset and must not be " + "reachable by a holder of the tool-invocation credential." + ) default_listen_addr = ( "127.0.0.1:8443" @@ -533,4 +548,5 @@ def load_config(path: str) -> Config: conformance_profile=profile, dev_mode=dev_mode, bearer_token=bearer_token, + operator_token=operator_token, ) diff --git a/src/cmcp_runtime/mcp/proxy.py b/src/cmcp_runtime/mcp/proxy.py index 9cb58b1c..ce3c2ee2 100644 --- a/src/cmcp_runtime/mcp/proxy.py +++ b/src/cmcp_runtime/mcp/proxy.py @@ -92,6 +92,9 @@ class _CallFinalizationState: """Per-invocation facts needed for honest terminal finalization.""" failure_stage: str = "call_entry" + # Session generation observed at call entry, so a response landing after an + # operator reset is not applied to the successor session. + reset_count: int | None = None effect_boundary_state: _EffectBoundaryState = _EffectBoundaryState.PRE_TRANSPORT request_payload_hash: str | None = None response_payload_hash: str | None = None @@ -926,6 +929,10 @@ async def call_tool( ) -> CallResult: """Run one call and guarantee one terminal on failure or cancellation.""" finalization = _CallFinalizationState() + # The session generation this call was issued under. A reset arriving + # mid-call closes that session, and this response must not raise the + # successor. + finalization.reset_count = self._session.reset_count try: return await self._call_tool_impl( call_id, @@ -1370,6 +1377,7 @@ class above the tool's catalogued sensitivity_level. It can never lower else [entry.sensitivity_level] ), injection_detected=injection_detected, + for_reset_count=_finalization.reset_count, response_allowed=False, ) threat_categories = ",".join( @@ -1442,11 +1450,18 @@ class above the tool's catalogued sensitivity_level. It can never lower injection_threshold = None _finalization.failure_stage = "session_update" async with self._session.mutation_lock: - self._session.update_from_inspection( + applied = self._session.update_from_inspection( call_id=call_id, sensitivity_tags=response_sensitivity, injection_detected=injection_detected, response_allowed=True, + for_reset_count=_finalization.reset_count, + ) + if not applied: + logger.warning( + "SESSION_RESET_RACE: response for call_id=%s dropped from session " + "state; the session it was issued under was closed by a reset", + call_id, ) # Step 5: egress Cedar policy check diff --git a/src/cmcp_runtime/mcp/server.py b/src/cmcp_runtime/mcp/server.py index 6a188060..5a1b70b5 100644 --- a/src/cmcp_runtime/mcp/server.py +++ b/src/cmcp_runtime/mcp/server.py @@ -14,6 +14,7 @@ import json import logging import os +import re import time import uuid from collections import defaultdict @@ -32,7 +33,7 @@ if TYPE_CHECKING: from cmcp_runtime.audit.chain import AuditChain from cmcp_runtime.session.manager import SessionManager - from cmcp_runtime.session.state import SessionState + from cmcp_runtime.session.state import ClosedSessionRecord, SessionState logger = logging.getLogger(__name__) @@ -48,6 +49,12 @@ class StatelessKernel: # Endpoints exempt from bearer-token auth (Kubernetes liveness / readiness probes) _AUTH_EXEMPT_PATHS = {"/health", "/readyz"} +# The operator interface. These routes are not reachable as MCP tools and, when an +# operator token is configured, they do not accept the tool-invocation token: a +# reset lowers accumulated session sensitivity, so the credential that authorizes +# one must not be the credential an agent host already holds. +_OPERATOR_PATH_RE = re.compile(r"^/(?:sessions/[^/]+/reset|catalog/exception)$") + # DOS-001: default ceiling on a single request body. Overridable per # deployment via MCPServer(max_request_bytes=...). Named here rather than # left inline on the constructor so the argument-shape caps below can be @@ -264,15 +271,27 @@ async def dispatch(self, request: Request, call_next: Any) -> Response: class _BearerAuthMiddleware(BaseHTTPMiddleware): - """AUTH-001 (CRITICAL): validate Authorization: Bearer on all protected endpoints.""" + """AUTH-001 (CRITICAL): validate Authorization: Bearer on all protected endpoints. + + Operator routes are matched against ``_OPERATOR_PATH_RE`` and, when an + operator token is configured, accept only that token. Where none is + configured they fall back to the bearer token, which keeps existing + single-token deployments working; startup refuses that outside dev mode. + """ - def __init__(self, app: Any, *, bearer_token: str) -> None: + def __init__( + self, app: Any, *, bearer_token: str, operator_token: str | None = None + ) -> None: super().__init__(app) self._token = bearer_token + self._operator_token = operator_token async def dispatch(self, request: Request, call_next: Any) -> Response: if request.url.path in _AUTH_EXEMPT_PATHS: return await call_next(request) + expected = self._token + if self._operator_token is not None and _OPERATOR_PATH_RE.match(request.url.path): + expected = self._operator_token auth = request.headers.get("Authorization", "") prefix = "Bearer " if not auth.startswith(prefix): @@ -283,7 +302,7 @@ async def dispatch(self, request: Request, call_next: Any) -> Response: ) provided = auth[len(prefix):] # Constant-time compare to prevent timing oracle on the token - if not hmac.compare_digest(provided, self._token): + if not hmac.compare_digest(provided, expected): logger.warning("AUTH_FAILURE: invalid bearer token from %s", request.client) return JSONResponse( {"error": "unauthorized", "error_code": "INVALID_BEARER_TOKEN"}, @@ -308,6 +327,7 @@ def __init__( session_manager: SessionManager | None = None, audit_chain: AuditChain | None = None, bearer_token: str | None = None, + operator_token: str | None = None, session: SessionState | None = None, max_request_bytes: int = _DEFAULT_MAX_REQUEST_BYTES, ) -> None: @@ -316,6 +336,7 @@ def __init__( self._audit_chain = audit_chain self._session = session self._max_request_bytes = max_request_bytes + self._operator_token = operator_token self._audit = audit_chain # Chains of closed sessions, kept so /audit/export still serves them # after the live session rotates. @@ -329,10 +350,20 @@ def __init__( requests_per_minute=60, ) middleware = [rate_limit] + ( - [Middleware(_BearerAuthMiddleware, bearer_token=bearer_token)] + [ + Middleware( + _BearerAuthMiddleware, + bearer_token=bearer_token, + operator_token=operator_token, + ) + ] if bearer_token is not None else [] ) + # Final state of sessions closed by a credentialed reset, kept so the + # value a closed session reached survives the successor starting at the + # minimum level. + self._closed_sessions: dict[str, ClosedSessionRecord] = {} # AUTH-004: session cleanup interval configurable via env var (default 60s) self._cleanup_interval_s: int = int( os.environ.get("CMCP_SESSION_CLEANUP_INTERVAL_SECONDS", "60") @@ -935,16 +966,29 @@ async def _session_reset(self, request: Request) -> Response: return JSONResponse( {"error": f"session_id={session_id} not found"}, status_code=404 ) + # The middleware has already authenticated the operator credential on this + # route; record which credential was verified so the entry says so. + credential = ( + "operator_token" if self._operator_token is not None else "bearer_token" + ) # AUTH-002: lock guards against a concurrent tool-call coroutine modifying sensitivity. async with self._session.mutation_lock: # Capture the pre-reset sensitivity: reset() drops it back to # "public", and the elevated value the session held at reset time # is exactly the forensic detail the audit entry must preserve. sensitivity_before = self._session.max_sensitivity + closed = self._session.snapshot_for_close( + reason="operator reset via API", + authorized_by=credential, + ) old_id, new_id = self._session.reset( reason="operator reset via API", - authorized_by="api", + authorized_by=credential, ) + reset_count = self._session.reset_count + self._closed_sessions[closed.session_id] = closed + # Written while the chain still names the closed session, so the entry + # recording the boundary belongs to the session that reached that value. self._audit_chain.append( "session_reset", call_id=None, @@ -952,10 +996,21 @@ async def _session_reset(self, request: Request) -> Response: policy_decision="n/a", session_sensitivity_before=sensitivity_before, session_sensitivity_after=self._session.max_sensitivity, + detail={ + "closed_session_id": old_id, + "successor_session_id": new_id, + "reset_count": reset_count, + "credential_verified": credential, + "reason": "operator reset via API", + }, ) + # Entries after the boundary belong to the successor. + self._audit_chain.rotate_session_id(new_id) return JSONResponse({ "old_session_id": old_id, "new_session_id": new_id, + "closed_session_max_sensitivity": closed.max_sensitivity, + "reset_count": reset_count, "status": "reset", "attestation_stale": False, }) diff --git a/src/cmcp_runtime/session/state.py b/src/cmcp_runtime/session/state.py index b1e7006e..44aaa51a 100644 --- a/src/cmcp_runtime/session/state.py +++ b/src/cmcp_runtime/session/state.py @@ -78,6 +78,26 @@ class InjectionEvent: timestamp: str +@dataclass(frozen=True) +class ClosedSessionRecord: + """The final state of a session closed by a credentialed reset. + + Held apart from the successor's live state so that the accumulated value the + closed session reached is preserved rather than overwritten. The successor + starts at the minimum level, and this is the only place its predecessor's + final value survives outside the audit chain. + """ + + session_id: str + max_sensitivity: str + sensitivity_raised_at: str | None + sensitivity_raised_by_call: str | None + reset_count: int + closed_at: str + reason: str + authorized_by: str + + @dataclass class SessionState: """ @@ -88,9 +108,17 @@ class SessionState: only way to lower sensitivity. update_from_inspection() is the ONLY place where session sensitivity state - is updated. It is called by InspectionPipeline after all inspection stages + is updated. It is called by the proxy response path after all inspection stages complete, including for denied responses (a denied high-sensitivity response still raises session sensitivity because the agent knows the call was attempted). + + A response is only allowed to raise the session it was issued under. Callers + pass the ``reset_count`` observed when the call started and a response that + lands after a reset is dropped rather than applied to the successor. The + discriminator is ``reset_count`` and not ``session_id`` because + upgrade_attestation() rotates ``session_id`` while deliberately continuing + the same session at its current sensitivity, so a call in flight across an + attestation upgrade must still be applied. """ session_id: str @@ -125,12 +153,23 @@ def update_from_inspection( sensitivity_tags: list[str], injection_detected: bool, response_allowed: bool, # noqa: ARG002 (logged for future use) - ) -> None: + *, + for_reset_count: int | None = None, + ) -> bool: """ Update session state from an inspection result. - Called by InspectionPipeline after all stages complete. + Called by the proxy response path after all stages complete. Returns True + if the state was updated, False if the response belonged to a session that + has since been closed by a reset and was therefore dropped. + + ``for_reset_count`` is the reset counter observed when the call started. + When it does not match the current counter the response is evidence about + a closed session and must not raise the successor, whose whole purpose is + to start at the minimum level. """ + if for_reset_count is not None and for_reset_count != self.reset_count: + return False for tag in sensitivity_tags: new_max = _max_sensitivity(self.max_sensitivity, tag, self.sensitivity_order) if new_max != self.max_sensitivity: @@ -145,6 +184,24 @@ def update_from_inspection( timestamp=datetime.now(tz=UTC).isoformat(), ) ) + return True + + def snapshot_for_close(self, *, reason: str, authorized_by: str) -> ClosedSessionRecord: + """Capture this session's final state before a reset opens a successor. + + Call inside the mutation lock, immediately before reset(), so the value + recorded is the one the session held at the ordered session boundary. + """ + return ClosedSessionRecord( + session_id=self.session_id, + max_sensitivity=self.max_sensitivity, + sensitivity_raised_at=self.sensitivity_raised_at, + sensitivity_raised_by_call=self.sensitivity_raised_by_call, + reset_count=self.reset_count, + closed_at=datetime.now(tz=UTC).isoformat(), + reason=reason, + authorized_by=authorized_by, + ) def reset(self, *, reason: str, authorized_by: str) -> tuple[str, str]: """ diff --git a/src/cmcp_runtime/startup.py b/src/cmcp_runtime/startup.py index 7ca8ec7f..5989bb1f 100644 --- a/src/cmcp_runtime/startup.py +++ b/src/cmcp_runtime/startup.py @@ -441,6 +441,19 @@ def run_startup(config_path: str) -> RuntimeContext: ) sys.exit(1) + # A session reset lowers accumulated session sensitivity. Requiring a + # separate credential for it keeps the reset out of reach of a holder of the + # tool-invocation token, which is the whole point of the monotonic state. + if config.operator_token is None and not config.dev_mode: + _fatal( + "OPERATOR_TOKEN_REQUIRED", + "CMCP_OPERATOR_TOKEN env var is not set. " + "Set it to a secret token, distinct from CMCP_BEARER_TOKEN, that " + "operators must present to the session reset and catalog exception " + "routes. Set CMCP_DEV_MODE=1 only in development.", + ) + sys.exit(1) + # Step 4: policy bundle policy_expected_hash = os.environ.get("CMCP_POLICY_HASH") if policy_expected_hash is None and not config.dev_mode: diff --git a/tests/unit/test_config.py b/tests/unit/test_config.py index a99794be..a5b283a5 100644 --- a/tests/unit/test_config.py +++ b/tests/unit/test_config.py @@ -450,3 +450,31 @@ def test_compliance_domains_non_mapping_raises(config_file): path = config_file("sensitivity:\n compliance_domains: not_a_mapping\n") with pytest.raises(ConfigError, match="mapping"): load_config(path) + + +# ── OPQ_P0006: the operator credential must be distinct ─────────────────────── + + +def test_operator_token_is_loaded(config_file, monkeypatch): + import cmcp_runtime.config as config_module + + monkeypatch.setattr(config_module, "DEV_MODE", False) + monkeypatch.setenv("CMCP_BEARER_TOKEN", "tool-token") + monkeypatch.setenv("CMCP_OPERATOR_TOKEN", "operator-token") + + cfg = load_config(config_file("")) + + assert cfg.bearer_token == "tool-token" + assert cfg.operator_token == "operator-token" + + +def test_operator_token_equal_to_bearer_token_is_refused(config_file, monkeypatch): + """Reusing the tool-invocation token as the operator token defeats the separation.""" + import cmcp_runtime.config as config_module + + monkeypatch.setattr(config_module, "DEV_MODE", False) + monkeypatch.setenv("CMCP_BEARER_TOKEN", "same-token") + monkeypatch.setenv("CMCP_OPERATOR_TOKEN", "same-token") + + with pytest.raises(ConfigError, match="must differ from CMCP_BEARER_TOKEN"): + load_config(config_file("")) diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index 38654d29..ac8c3998 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -226,3 +226,53 @@ def test_upgrade_attestation_does_not_increment_reset_count(): state = SessionState(session_id="s1") state.upgrade_attestation() assert state.reset_count == 0 + + +@pytest.mark.asyncio +async def test_response_from_closed_session_does_not_raise_successor(): + """A response issued before an operator reset must not elevate the successor. + + The successor session exists to start at the minimum level. A response that + was in flight when the reset closed the previous session is evidence about + that closed session, so it is dropped rather than applied. + """ + state = SessionState(session_id="s-pre") + state.update_from_inspection("call-A", ["pii"], False, True) + generation = state.reset_count + + async def _responder() -> bool: + async with state.mutation_lock: + return state.update_from_inspection( + "call-B", ["confidential"], False, True, for_reset_count=generation + ) + + async def _reset() -> None: + async with state.mutation_lock: + state.reset(reason="operator reset via API", authorized_by="test") + + # Hold the lock so both queue behind us, then let the reset win it. + await state.mutation_lock.acquire() + resetter = asyncio.create_task(_reset()) + await asyncio.sleep(0) + responder = asyncio.create_task(_responder()) + await asyncio.sleep(0) + state.mutation_lock.release() + await resetter + applied = await responder + + assert applied is False + assert state.max_sensitivity == "public" + assert state.sensitivity_raised_by_call is None + + +@pytest.mark.asyncio +async def test_response_across_attestation_upgrade_still_raises(): + """upgrade_attestation() continues the session, so an in-flight response applies.""" + state = SessionState(session_id="s1") + generation = state.reset_count + state.upgrade_attestation() + applied = state.update_from_inspection( + "call-C", ["pii"], False, True, for_reset_count=generation + ) + assert applied is True + assert state.max_sensitivity == "pii" diff --git a/tests/unit/test_session_reset.py b/tests/unit/test_session_reset.py index 6c8faa90..55194184 100644 --- a/tests/unit/test_session_reset.py +++ b/tests/unit/test_session_reset.py @@ -224,3 +224,141 @@ def test_reset_without_session_configured_returns_501(): client = TestClient(server.app, raise_server_exceptions=True) resp = client.post("/sessions/bare-sess/reset") assert resp.status_code == 501 + + +# ── OPQ_P0006: the reset credential is not the tool-invocation credential ───── + + +def _make_token_server(session_id: str = "sess-tok-001", *, operator_token: str | None): + """Server with a tool-invocation bearer token and an optional operator token.""" + from cmcp_runtime.mcp.proxy import CMCPProxy + + cfg = Config() + cfg.attestation = AttestationConfig(enforcement_mode=EnforcementMode.ENFORCING) + session = SessionState(session_id=session_id) + chain = AuditChain(session_id) + with patch("cmcp_runtime.mcp.proxy.MCPGateway"), \ + patch("cmcp_runtime.mcp.proxy.MCPResponseScanner"): + proxy = CMCPProxy(_make_catalog(), _make_evaluator(), session, chain, cfg) + wire_mock_gateway(proxy) + server = MCPServer( + proxy, + session=session, + audit_chain=chain, + bearer_token="tool-token", + operator_token=operator_token, + ) + return server, session, chain + + +def test_reset_rejects_the_tool_invocation_token(): + """The credential that invokes tools must not authorize a sensitivity reset.""" + server, session, _ = _make_token_server(operator_token="operator-token") + client = TestClient(server.app, raise_server_exceptions=True) + + resp = client.post( + f"/sessions/{session.session_id}/reset", + headers={"Authorization": "Bearer tool-token"}, + ) + assert resp.status_code == 401 + assert resp.json()["error_code"] == "INVALID_BEARER_TOKEN" + + +def test_reset_accepts_the_operator_token(): + server, session, _ = _make_token_server(operator_token="operator-token") + original_id = session.session_id + client = TestClient(server.app, raise_server_exceptions=True) + + resp = client.post( + f"/sessions/{original_id}/reset", + headers={"Authorization": "Bearer operator-token"}, + ) + assert resp.status_code == 200 + assert resp.json()["old_session_id"] == original_id + + +def test_tool_endpoint_rejects_the_operator_token(): + """The separation runs both ways: the operator credential is not a tool credential.""" + server, _, _ = _make_token_server(operator_token="operator-token") + client = TestClient(server.app, raise_server_exceptions=True) + + resp = client.get("/tools/list", headers={"Authorization": "Bearer operator-token"}) + assert resp.status_code == 401 + + +def test_reset_falls_back_to_bearer_token_when_no_operator_token(): + """Single-token deployments keep working; startup refuses them outside dev mode.""" + server, session, _ = _make_token_server(operator_token=None) + client = TestClient(server.app, raise_server_exceptions=True) + + resp = client.post( + f"/sessions/{session.session_id}/reset", + headers={"Authorization": "Bearer tool-token"}, + ) + assert resp.status_code == 200 + + +# ── OPQ_P0006: the reset record carries the session boundary ────────────────── + + +def test_reset_audit_entry_identifies_both_sessions_and_the_credential(): + server, session, chain = _make_token_server(operator_token="operator-token") + session.update_from_inspection("call-A", ["pii"], False, True) + original_id = session.session_id + + client = TestClient(server.app, raise_server_exceptions=True) + resp = client.post( + f"/sessions/{original_id}/reset", + headers={"Authorization": "Bearer operator-token"}, + ) + new_id = resp.json()["new_session_id"] + + entry = next(e for e in chain.entries if e.entry_type == "session_reset") + assert entry.session_id == original_id + assert entry.session_sensitivity_before == "pii" + assert entry.session_sensitivity_after == "public" + assert entry.detail["closed_session_id"] == original_id + assert entry.detail["successor_session_id"] == new_id + assert entry.detail["reset_count"] == 1 + assert entry.detail["credential_verified"] == "operator_token" + assert entry.prev_entry_hash + # detail is inside the canonical body, so these fields are hash-covered + assert entry.entry_hash == entry.compute_hash() + + +def test_entries_after_a_reset_are_attributed_to_the_successor(): + """Before this, every later entry carried the closed session's identifier.""" + server, session, chain = _make_token_server(operator_token="operator-token") + original_id = session.session_id + + client = TestClient(server.app, raise_server_exceptions=True) + new_id = client.post( + f"/sessions/{original_id}/reset", + headers={"Authorization": "Bearer operator-token"}, + ).json()["new_session_id"] + + later = chain.append("session_start", policy_decision="n/a") + assert later.session_id == new_id + reset_entry = next(e for e in chain.entries if e.entry_type == "session_reset") + assert reset_entry.session_id == original_id + assert chain.verify_chain() + + +def test_closed_session_final_value_is_preserved_apart_from_the_successor(): + server, session, _ = _make_token_server(operator_token="operator-token") + session.update_from_inspection("call-A", ["pii"], False, True) + original_id = session.session_id + + client = TestClient(server.app, raise_server_exceptions=True) + resp = client.post( + f"/sessions/{original_id}/reset", + headers={"Authorization": "Bearer operator-token"}, + ) + + assert resp.json()["closed_session_max_sensitivity"] == "pii" + assert session.max_sensitivity == "public" + closed = server._closed_sessions[original_id] + assert closed.session_id == original_id + assert closed.max_sensitivity == "pii" + assert closed.sensitivity_raised_by_call == "call-A" + assert closed.authorized_by == "operator_token" From b5da3595c0ab09c6fa8c87dd3caa7be2ff42cca9 Mon Sep 17 00:00:00 2001 From: Imran Siddique Date: Thu, 10 Sep 2026 16:47:39 -0700 Subject: [PATCH 2/2] feat(session): shared, persistent home for the accumulated sensitivity value The accumulated session-sensitivity value lived in the gateway process, which left two gaps in what the ratchet actually guarantees. It did not survive a restart. A session that had reached hipaa_phi came back at public while the session identifier the agent host holds was still live, so the sensitivity accumulated over the whole session was gone and enforcement resumed from the minimum level. It was not shared. Where several instances serve one agent session, each held its own value, so the ratchet held per instance rather than per session. An agent that read sensitive data through one instance and egressed through another was evaluated by an instance that never saw the read. session/store.py holds both answers behind one protocol. InMemorySessionStateStore is the default and is the previous behaviour exactly. SqliteSessionStateStore serialises the read-modify-write with BEGIN IMMEDIATE, which takes SQLite's RESERVED lock: two processes cannot hold it at once, so the critical section spans instances on a shared volume. An asyncio.Lock cannot do this, being invisible to every other instance, and a test proves the difference against a real second process holding the lock. The audit chain already takes its durability from SQLite in WAL mode, so turning this on adds no infrastructure. apply_inspection and apply_reset are now the write path. Each holds the store's exclusive section, reads the shared value back, folds the response into it and writes it, so the greater-of comparison is made against what every instance shares rather than against one instance's copy. The gateway hydrates at call entry, before the pre-call policy evaluation reads the value and before the generation the call belongs to is captured, so an instance joining a session another instance opened does not permit what the session already forbids. A reset also advances the closed session's generation in the store while keeping the value it reached. The successor is written under its own identifier, so without that an instance still holding the old identifier would read a generation matching the one it captured and raise a session already closed. Only the value that must be shared is stored: injection events, drift lists and the kill switch stay per instance, because copying them would present one instance's observations as another's. session_state_path is unset by default. 10 tests added, 1681 passing. The two suite failures are pre-existing on the base: test_release_distribution_smoke (runtime 0.5.0 against metadata 0.4.0) and test_tpm_chained_verify (assertion text drift). Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01CDzMWcn12dszchgrKfXzZY --- CHANGELOG.md | 17 ++ docs/configuration.md | 1 + src/cmcp_runtime/config.py | 10 + src/cmcp_runtime/mcp/proxy.py | 51 ++--- src/cmcp_runtime/mcp/server.py | 26 +-- src/cmcp_runtime/session/manager.py | 6 +- src/cmcp_runtime/session/state.py | 135 ++++++++++++ src/cmcp_runtime/session/store.py | 271 +++++++++++++++++++++++++ src/cmcp_runtime/startup.py | 4 + tests/unit/test_session_state_store.py | 209 +++++++++++++++++++ 10 files changed, 691 insertions(+), 39 deletions(-) create mode 100644 src/cmcp_runtime/session/store.py create mode 100644 tests/unit/test_session_state_store.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 95449f1e..f47ec0ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 unset those routes still fall back to the bearer token, so an existing single-token deployment keeps working until it sets the new variable. +### Added + +- **The accumulated session-sensitivity value can now live in a shared, + persistent store** (`session_state_path`). Without one it is held in the + gateway process, so it is lost on restart while the session identifier the + agent host holds is still live, and where several instances serve one session + the ratchet holds per instance rather than per session: an agent that reads + sensitive data through one instance and egresses through another is evaluated + by an instance that never saw the read. `SqliteSessionStateStore` serialises + the read-modify-write with `BEGIN IMMEDIATE`, which takes SQLite's RESERVED + lock and so spans processes; an `asyncio.Lock` cannot, being invisible to every + other instance. A gateway now hydrates the session's stored value at call entry, + before the pre-call policy evaluation reads it. A reset also advances the closed + session's generation in the store, so an instance still holding the old + identifier stops applying responses to it. Unset is the default and preserves + the previous single-instance behaviour exactly. + ### Changed - The reset audit entry now identifies the session boundary rather than only the diff --git a/docs/configuration.md b/docs/configuration.md index a83d6e4f..a7be0b69 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -121,6 +121,7 @@ chain and TRACE Claim; this setting controls whether the session also fails clos | `policy_bundle_path` | string | `policies/` | Path to the Cedar policy bundle directory. Must contain `.cedar` files and a `manifest.json`. Path traversal (`..` components) is rejected. | | `catalog_path` | string | `catalog.json` | Path to the JSON tool catalog. Path traversal (`..` components) is rejected. | | `listen_addr` | string | `0.0.0.0:8443` | Address and port the gateway binds to. Default is `127.0.0.1:8443` in tokenless `CMCP_DEV_MODE=1`, otherwise `0.0.0.0:8443`. Tokenless dev mode requires loopback (e.g., `127.0.0.1:8443`, `localhost:8443`, `[::1]:8443`). Wildcard, LAN, public, and non-loopback hostname binds require `CMCP_BEARER_TOKEN`. | +| `session_state_path` | string or unset | unset | Path to the shared session-state database. Unset keeps the accumulated session-sensitivity value in the gateway process: correct for a single instance, but the value is lost on restart and each instance ratchets separately. Set it to a path on a volume every instance shares to make the ratchet hold per session across instances and survive a restart. Writes are serialised with SQLite `BEGIN IMMEDIATE`, which is a cross-process lock, so no external coordinator is needed. Path traversal (`..` components) is rejected. | | `max_response_size_bytes` | integer | `2097152` | Maximum tool response size in bytes (2MB). Must be a positive integer. Responses exceeding this limit are rejected before inspection. | | `policy_reload_interval_seconds` | integer | `0` | Interval in seconds between automatic Cedar bundle reloads. `0` disables automatic reload. Above `0` requires a pinned `CMCP_POLICY_SIGNING_KEY`: **an interval alongside only a pinned `CMCP_POLICY_HASH` aborts startup** (`POLICY_RELOAD_PINNED_HASH`): the reload re-validates against that pinned hash, so a bundle that actually changed would always be rejected and the old policy would stay in force. A hash pins one artifact and so cannot authorise a bundle that changed; a signing key approves any bundle the authority signs, which is what reload needs. Both pins together is the supported production shape. See [Policy Hot-Reload](spec/policy-hot-reload.md). | | `conformance_profile` | string or unset | unset | A named profile that tightens defaults which stay permissive for developers. Only `aarm` is defined: it requires an Agent Manifest binding, because AARM R6 says every receipt MUST be bound to an agent identity while the developer default leaves binding optional. With it set, the gateway refuses to start unless `agent_manifest.path` and `agent_manifest.trust_anchor_path` are configured (`CONFORMANCE_PROFILE_UNSATISFIED`). An unrecognised name is a config error rather than a profile that enforces nothing. | diff --git a/src/cmcp_runtime/config.py b/src/cmcp_runtime/config.py index 7fd45976..1bc9a345 100644 --- a/src/cmcp_runtime/config.py +++ b/src/cmcp_runtime/config.py @@ -132,6 +132,12 @@ class Config: max_response_size_bytes: int = 2 * 1024 * 1024 # 2MB policy_reload_interval_seconds: int = 0 # 0 = disabled (POLICY-001) audit_db_path: str = "audit.db" # AUDIT-001: durable audit chain storage + #: Path to the shared session-state database. Unset keeps the accumulated + #: session-sensitivity value in the gateway process, which is correct for a + #: single instance and loses the value on restart. Set it to a path on a + #: volume every instance shares to make the ratchet hold per session across + #: instances and survive a restart. See ``session/store.py``. + session_state_path: str | None = None dev_mode: bool = False bearer_token: str | None = None #: Credential for the operator interface (session reset, catalog exception). @@ -480,6 +486,9 @@ def load_config(path: str) -> Config: policy_bundle_path = raw.get("policy_bundle_path", "policy/") catalog_path = raw.get("catalog_path", "catalog.json") audit_db_path = raw.get("audit_db_path", "audit.db") + session_state_path = raw.get("session_state_path") or None + if session_state_path is not None: + _check_no_traversal("session_state_path", session_state_path) _check_no_traversal("policy_bundle_path", policy_bundle_path) _check_no_traversal("catalog_path", catalog_path) _check_no_traversal("audit_db_path", audit_db_path) @@ -549,4 +558,5 @@ def load_config(path: str) -> Config: dev_mode=dev_mode, bearer_token=bearer_token, operator_token=operator_token, + session_state_path=session_state_path, ) diff --git a/src/cmcp_runtime/mcp/proxy.py b/src/cmcp_runtime/mcp/proxy.py index ce3c2ee2..fd90eba0 100644 --- a/src/cmcp_runtime/mcp/proxy.py +++ b/src/cmcp_runtime/mcp/proxy.py @@ -929,9 +929,16 @@ async def call_tool( ) -> CallResult: """Run one call and guarantee one terminal on failure or cancellation.""" finalization = _CallFinalizationState() - # The session generation this call was issued under. A reset arriving - # mid-call closes that session, and this response must not raise the - # successor. + # Adopt the session's shared value before anything evaluates this call. + # An instance joining a session another instance opened, or one that has + # restarted, would otherwise evaluate the first call against its own + # empty copy and permit what the session's accumulated value forbids. + # No-op when no shared store is configured. + await self._session.hydrate() + # The session generation this call was issued under, read after hydration + # so a reset performed on another instance is already visible. A reset + # arriving mid-call closes that session, and this response must not raise + # the successor. finalization.reset_count = self._session.reset_count try: return await self._call_tool_impl( @@ -1368,18 +1375,17 @@ class above the tool's catalogued sensitivity_level. It can never lower ) 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, declared_data_class] - if declared_data_class is not None - else [entry.sensitivity_level] - ), - injection_detected=injection_detected, - for_reset_count=_finalization.reset_count, - response_allowed=False, - ) + await self._session.apply_inspection( + call_id=call_id, + sensitivity_tags=( + [entry.sensitivity_level, declared_data_class] + if declared_data_class is not None + else [entry.sensitivity_level] + ), + injection_detected=injection_detected, + for_reset_count=_finalization.reset_count, + response_allowed=False, + ) threat_categories = ",".join( sorted({str(t.get("category", "unknown")) for t in scan.threats}) ) @@ -1449,14 +1455,13 @@ class above the tool's catalogued sensitivity_level. It can never lower ) injection_threshold = None _finalization.failure_stage = "session_update" - async with self._session.mutation_lock: - applied = self._session.update_from_inspection( - call_id=call_id, - sensitivity_tags=response_sensitivity, - injection_detected=injection_detected, - response_allowed=True, - for_reset_count=_finalization.reset_count, - ) + applied = await self._session.apply_inspection( + call_id=call_id, + sensitivity_tags=response_sensitivity, + injection_detected=injection_detected, + response_allowed=True, + for_reset_count=_finalization.reset_count, + ) if not applied: logger.warning( "SESSION_RESET_RACE: response for call_id=%s dropped from session " diff --git a/src/cmcp_runtime/mcp/server.py b/src/cmcp_runtime/mcp/server.py index 5a1b70b5..aa574070 100644 --- a/src/cmcp_runtime/mcp/server.py +++ b/src/cmcp_runtime/mcp/server.py @@ -971,21 +971,17 @@ async def _session_reset(self, request: Request) -> Response: credential = ( "operator_token" if self._operator_token is not None else "bearer_token" ) - # AUTH-002: lock guards against a concurrent tool-call coroutine modifying sensitivity. - async with self._session.mutation_lock: - # Capture the pre-reset sensitivity: reset() drops it back to - # "public", and the elevated value the session held at reset time - # is exactly the forensic detail the audit entry must preserve. - sensitivity_before = self._session.max_sensitivity - closed = self._session.snapshot_for_close( - reason="operator reset via API", - authorized_by=credential, - ) - old_id, new_id = self._session.reset( - reason="operator reset via API", - authorized_by=credential, - ) - reset_count = self._session.reset_count + # AUTH-002: apply_reset serialises against concurrent tool-call coroutines, + # and against other gateway instances where a shared store is configured. + # The pre-reset sensitivity is captured inside that section, because the + # elevated value the session held at the boundary is exactly the forensic + # detail the audit entry must preserve. + old_id, new_id, closed = await self._session.apply_reset( + reason="operator reset via API", + authorized_by=credential, + ) + sensitivity_before = closed.max_sensitivity + reset_count = self._session.reset_count self._closed_sessions[closed.session_id] = closed # Written while the chain still names the closed session, so the entry # recording the boundary belongs to the session that reached that value. diff --git a/src/cmcp_runtime/session/manager.py b/src/cmcp_runtime/session/manager.py index ffd0100d..6d8b8db1 100644 --- a/src/cmcp_runtime/session/manager.py +++ b/src/cmcp_runtime/session/manager.py @@ -116,7 +116,11 @@ def create_session(self) -> tuple[SessionState, AuditChain]: ) session_id = str(uuid4()) - state = SessionState(session_id=session_id, sensitivity_order=self._sensitivity_order) + state = SessionState( + session_id=session_id, + sensitivity_order=self._sensitivity_order, + state_store=getattr(self._ctx, "session_state_store", None), + ) chain = AuditChain( session_id=session_id, store=self._ctx.audit_store, diff --git a/src/cmcp_runtime/session/state.py b/src/cmcp_runtime/session/state.py index 44aaa51a..84d3f902 100644 --- a/src/cmcp_runtime/session/state.py +++ b/src/cmcp_runtime/session/state.py @@ -7,6 +7,8 @@ from datetime import UTC, datetime from uuid import uuid4 +from cmcp_runtime.session.store import SessionStateStore, StoredSensitivity + # Sensitivity level ordering: monotonically increasing only. # hipaa_phi, mnpi, trade_secret are all at level 3 (equal highest). SENSITIVITY_ORDER: dict[str, int] = { @@ -146,6 +148,13 @@ class SessionState: ) # AUTH-002: guards concurrent mutations from tool-call coroutines and session-reset requests mutation_lock: asyncio.Lock = field(default_factory=asyncio.Lock, init=False, repr=False, compare=False) + #: Where the accumulated value is held. None keeps the value in this object + #: alone, which is the single-instance default. A shared store makes the + #: ratchet hold per session across gateway instances rather than per + #: instance, and makes it survive a restart. See ``session/store.py``. + state_store: SessionStateStore | None = field( + default=None, repr=False, compare=False + ) def update_from_inspection( self, @@ -203,6 +212,132 @@ def snapshot_for_close(self, *, reason: str, authorized_by: str) -> ClosedSessio authorized_by=authorized_by, ) + async def apply_inspection( + self, + call_id: str, + sensitivity_tags: list[str], + injection_detected: bool, + response_allowed: bool, + *, + for_reset_count: int | None = None, + ) -> bool: + """Serialise and apply one inspection result to the accumulated value. + + This is the only write path for tool-call processing. With no store + configured it is the previous behaviour: take the in-process lock and + mutate. With a store it holds the store's exclusive section, reads the + current value back, folds this response into it and writes it, so the + greater-of comparison is made against the value every instance shares + rather than against this instance's copy. + """ + if self.state_store is None: + async with self.mutation_lock: + return self.update_from_inspection( + call_id, + sensitivity_tags, + injection_detected, + response_allowed, + for_reset_count=for_reset_count, + ) + async with self.state_store.exclusive(self.session_id): + self._adopt(self.state_store.load(self.session_id)) + applied = self.update_from_inspection( + call_id, + sensitivity_tags, + injection_detected, + response_allowed, + for_reset_count=for_reset_count, + ) + if applied: + self.state_store.save(self.session_id, self._stored()) + return applied + + async def apply_reset( + self, *, reason: str, authorized_by: str + ) -> tuple[str, str, ClosedSessionRecord]: + """Serialise and apply a credentialed reset, closing the session. + + Returns the closed identifier, the successor identifier, and the record + preserving what the closed session reached. The snapshot is taken inside + the exclusive section so the preserved value is the one held at the + ordered session boundary, and the successor is written to the store so + that no instance keeps serving the closed session's value. + """ + if self.state_store is None: + async with self.mutation_lock: + closed = self.snapshot_for_close( + reason=reason, authorized_by=authorized_by + ) + old_id, new_id = self.reset(reason=reason, authorized_by=authorized_by) + return old_id, new_id, closed + async with self.state_store.exclusive(self.session_id): + self._adopt(self.state_store.load(self.session_id)) + closed = self.snapshot_for_close(reason=reason, authorized_by=authorized_by) + old_id, new_id = self.reset(reason=reason, authorized_by=authorized_by) + self.state_store.record_closed(old_id, closed) + # Advance the closed session's generation in the store, keeping the + # value it reached. Another instance may still be holding the old + # identifier with a response in flight; without this it would read a + # generation matching the one it captured and raise a session that is + # already closed. The successor is written under its own identifier, + # so bumping the old row is the only way that instance finds out. + self.state_store.save( + old_id, + StoredSensitivity( + max_sensitivity=closed.max_sensitivity, + sensitivity_raised_at=closed.sensitivity_raised_at, + sensitivity_raised_by_call=closed.sensitivity_raised_by_call, + reset_count=self.reset_count, + ), + ) + self.state_store.save(new_id, self._stored()) + return old_id, new_id, closed + + async def hydrate(self) -> bool: + """Adopt this session's stored value, if the store holds one. + + Called at startup so a gateway that restarts, or an instance joining a + session another instance opened, enforces against what the session + already accumulated instead of starting the ratchet again at the minimum + level. Returns True when a stored value was adopted. + """ + if self.state_store is None: + return False + async with self.state_store.exclusive(self.session_id): + stored = self.state_store.load(self.session_id) + if stored is None: + self.state_store.save(self.session_id, self._stored()) + return False + self._adopt(stored) + return True + + def _stored(self) -> StoredSensitivity: + return StoredSensitivity( + max_sensitivity=self.max_sensitivity, + sensitivity_raised_at=self.sensitivity_raised_at, + sensitivity_raised_by_call=self.sensitivity_raised_by_call, + reset_count=self.reset_count, + ) + + def _adopt(self, stored: StoredSensitivity | None) -> None: + """Take the store's value as this instance's own. + + Only ever raises: the stored value is the greater of what any instance + has seen, and a local value above it would mean this instance observed + something it has not yet written. The reset counter is taken whole, + because a reset performed on another instance closed this session there + and this instance must not keep applying responses to it. + """ + if stored is None: + return + self.max_sensitivity = _max_sensitivity( + self.max_sensitivity, stored.max_sensitivity, self.sensitivity_order + ) + if self.max_sensitivity == stored.max_sensitivity: + self.sensitivity_raised_at = stored.sensitivity_raised_at + self.sensitivity_raised_by_call = stored.sensitivity_raised_by_call + self.reset_count = max(self.reset_count, stored.reset_count) + def reset(self, *, reason: str, authorized_by: str) -> tuple[str, str]: """ Reset session sensitivity to 'public'. Returns (previous_session_id, new_session_id). diff --git a/src/cmcp_runtime/session/store.py b/src/cmcp_runtime/session/store.py new file mode 100644 index 00000000..2d2447bb --- /dev/null +++ b/src/cmcp_runtime/session/store.py @@ -0,0 +1,271 @@ +""" +Persistent, shared session-sensitivity state. + +The accumulated session-sensitivity value is the one piece of session state that +enforcement depends on: a gateway that loses it stops being able to say what the +session already handled, and a gateway that keeps it privately cannot enforce +across more than one instance. Two consequences follow, and this module holds +both. + +**Persistence.** An in-memory value does not survive a restart. A session that +had reached ``restricted`` comes back at ``public`` after a process bounce, and +the sensitivity that was accumulated over the whole session is gone while the +session identifier the agent host holds is still live. + +**Sharing.** Where several gateway instances serve one agent session, each +holding its own value means the ratchet holds per instance and not per session. +An agent that reads sensitive data through one instance and then egresses +through another is evaluated by an instance that never saw the read. + +Both are answered by keeping the value in one store and serialising the +read-modify-write against it. The mechanism has to serialise *across processes*, +not merely across coroutines in one of them, so an ``asyncio.Lock`` is not +sufficient: it is invisible to every other instance. + +``SqliteSessionStateStore`` uses ``BEGIN IMMEDIATE``, which takes SQLite's +database-level RESERVED lock. Two processes cannot hold it at once, so the +critical section spans instances on a shared filesystem volume. This follows the +audit chain, which already takes its durability from SQLite in WAL mode, so a +deployment gains no new infrastructure by turning it on. + +``InMemorySessionStateStore`` is the default and preserves the single-instance +behaviour exactly: no file, no cross-process guarantee, and the same +``asyncio.Lock`` semantics as before. +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import sqlite3 +import threading +from collections.abc import AsyncIterator +from contextlib import AbstractAsyncContextManager, asynccontextmanager +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Protocol, runtime_checkable + +logger = logging.getLogger(__name__) + +_CREATE_TABLES = """ +CREATE TABLE IF NOT EXISTS session_state ( + session_id TEXT PRIMARY KEY, + max_sensitivity TEXT NOT NULL, + sensitivity_raised_at TEXT, + sensitivity_raised_by_call TEXT, + reset_count INTEGER NOT NULL DEFAULT 0, + updated_at TEXT NOT NULL DEFAULT (datetime('now')) +); + +CREATE TABLE IF NOT EXISTS closed_sessions ( + session_id TEXT PRIMARY KEY, + payload TEXT NOT NULL, + closed_at TEXT NOT NULL +); +""" + + +@dataclass(frozen=True) +class StoredSensitivity: + """The subset of session state that is shared and must survive a restart. + + Deliberately narrow. Injection events, drift lists and the kill switch are + per-instance observations, and copying them between instances would present + one instance's observations as another's. + """ + + max_sensitivity: str + sensitivity_raised_at: str | None + sensitivity_raised_by_call: str | None + reset_count: int + + +@runtime_checkable +class SessionStateStore(Protocol): + """Where the accumulated session-sensitivity value lives. + + ``exclusive`` must serialise every operation that modifies the value for one + session, across all instances sharing the store. Implementations that cannot + do that across processes must say so in their docstring, because a caller + cannot tell the difference from the outside until it is enforcing wrongly. + """ + + def exclusive(self, session_id: str) -> AbstractAsyncContextManager[None]: + """Hold the session's write lock for the duration of the context.""" + ... + + def load(self, session_id: str) -> StoredSensitivity | None: + """Return the stored value, or None when the session is not yet stored.""" + ... + + def save(self, session_id: str, value: StoredSensitivity) -> None: + """Write the value. Called inside ``exclusive``.""" + ... + + def record_closed(self, session_id: str, record: object) -> None: + """Preserve the final state of a session closed by a reset.""" + ... + + def is_shared(self) -> bool: + """True when the store serialises across gateway instances.""" + ... + + +class InMemorySessionStateStore: + """Single-instance store. The default, and the previous behaviour exactly. + + Serialises with an ``asyncio.Lock``, which is confined to one event loop in + one process: it is not visible to another gateway instance and does not + survive a restart. ``is_shared()`` returns False so a caller can tell. + """ + + def __init__(self) -> None: + self._values: dict[str, StoredSensitivity] = {} + self._closed: dict[str, object] = {} + self._locks: dict[str, asyncio.Lock] = {} + + @asynccontextmanager + async def exclusive(self, session_id: str) -> AsyncIterator[None]: + lock = self._locks.setdefault(session_id, asyncio.Lock()) + async with lock: + yield + + def load(self, session_id: str) -> StoredSensitivity | None: + return self._values.get(session_id) + + def save(self, session_id: str, value: StoredSensitivity) -> None: + self._values[session_id] = value + + def record_closed(self, session_id: str, record: object) -> None: + self._closed[session_id] = record + + def closed(self, session_id: str) -> object | None: + return self._closed.get(session_id) + + def is_shared(self) -> bool: + return False + + +class SqliteSessionStateStore: + """Shared, persistent store serialised across gateway instances. + + ``exclusive`` opens a ``BEGIN IMMEDIATE`` transaction, which takes SQLite's + RESERVED lock on the database. Only one connection can hold it, so the + critical section covers every instance pointed at the same file, and the + read-modify-write of the accumulated value cannot interleave with another + instance's. ``busy_timeout`` makes a contending instance wait rather than + fail, and the transaction commits on exit so the value is on disk before the + lock is released. + + The lock is database-wide rather than per session. For the write path this + is a correct-but-coarse choice: sensitivity writes are short, and a + per-session lock in SQLite would need its own lease table with expiry, which + is a failure mode (a crashed holder blocking a session until its lease ages + out) in exchange for concurrency this workload does not need. + """ + + def __init__(self, db_path: Path, *, busy_timeout_ms: int = 5000) -> None: + self._db_path = db_path + self._conn = sqlite3.connect( + str(db_path), check_same_thread=False, isolation_level=None + ) + self._thread_lock = threading.Lock() + self._conn.execute("PRAGMA journal_mode=WAL") + self._conn.execute("PRAGMA synchronous=FULL") + self._conn.execute(f"PRAGMA busy_timeout={int(busy_timeout_ms)}") + self._conn.executescript(_CREATE_TABLES) + logger.info("Session state store opened: path=%s", db_path) + + @asynccontextmanager + async def exclusive(self, session_id: str) -> AsyncIterator[None]: + # Taken in a worker thread: sqlite3 blocks, and blocking the event loop + # while another instance holds the lock would stall every other call. + await asyncio.to_thread(self._begin_immediate) + try: + yield + except BaseException: + await asyncio.to_thread(self._rollback) + raise + else: + await asyncio.to_thread(self._commit) + + def _begin_immediate(self) -> None: + self._thread_lock.acquire() + try: + self._conn.execute("BEGIN IMMEDIATE") + except BaseException: + self._thread_lock.release() + raise + + def _commit(self) -> None: + try: + self._conn.execute("COMMIT") + finally: + self._thread_lock.release() + + def _rollback(self) -> None: + try: + self._conn.execute("ROLLBACK") + finally: + self._thread_lock.release() + + def load(self, session_id: str) -> StoredSensitivity | None: + cur = self._conn.execute( + "SELECT max_sensitivity, sensitivity_raised_at, " + "sensitivity_raised_by_call, reset_count " + "FROM session_state WHERE session_id = ?", + (session_id,), + ) + row = cur.fetchone() + if row is None: + return None + return StoredSensitivity( + max_sensitivity=row[0], + sensitivity_raised_at=row[1], + sensitivity_raised_by_call=row[2], + reset_count=int(row[3]), + ) + + def save(self, session_id: str, value: StoredSensitivity) -> None: + self._conn.execute( + "INSERT INTO session_state " + "(session_id, max_sensitivity, sensitivity_raised_at, " + " sensitivity_raised_by_call, reset_count, updated_at) " + "VALUES (?, ?, ?, ?, ?, datetime('now')) " + "ON CONFLICT(session_id) DO UPDATE SET " + " max_sensitivity=excluded.max_sensitivity, " + " sensitivity_raised_at=excluded.sensitivity_raised_at, " + " sensitivity_raised_by_call=excluded.sensitivity_raised_by_call, " + " reset_count=excluded.reset_count, " + " updated_at=excluded.updated_at", + ( + session_id, + value.max_sensitivity, + value.sensitivity_raised_at, + value.sensitivity_raised_by_call, + value.reset_count, + ), + ) + + def record_closed(self, session_id: str, record: object) -> None: + payload = json.dumps(asdict(record), sort_keys=True) # type: ignore[call-overload] + self._conn.execute( + "INSERT INTO closed_sessions (session_id, payload, closed_at) " + "VALUES (?, ?, datetime('now')) " + "ON CONFLICT(session_id) DO UPDATE SET payload=excluded.payload", + (session_id, payload), + ) + + def closed(self, session_id: str) -> dict[str, object] | None: + cur = self._conn.execute( + "SELECT payload FROM closed_sessions WHERE session_id = ?", (session_id,) + ) + row = cur.fetchone() + return None if row is None else json.loads(row[0]) + + def is_shared(self) -> bool: + return True + + def close(self) -> None: + self._conn.close() diff --git a/src/cmcp_runtime/startup.py b/src/cmcp_runtime/startup.py index 5989bb1f..5a45ef30 100644 --- a/src/cmcp_runtime/startup.py +++ b/src/cmcp_runtime/startup.py @@ -32,6 +32,7 @@ PolicyHashMismatch, ) from cmcp_runtime.policy.bundle import PolicyStore, load_policy_bundle +from cmcp_runtime.session.store import SqliteSessionStateStore from cmcp_runtime.tee.base import AttestationReport, TEEProvider from cmcp_runtime.tee.detect import detect_provider from cmcp_runtime.tee.measurement import ( @@ -80,6 +81,9 @@ class RuntimeContext: # and still enforces via digest comparison. catalog_scanner: CatalogScanner | None = None audit_store: SqliteAuditStore | None = None + #: Shared, persistent home for the accumulated session-sensitivity value. + #: None means the value lives in this process only. + session_state_store: SqliteSessionStateStore | None = None spiffe: SpiffeClientResult | None = None nras_appraisal: AppraisalResult | None = None agent_manifest: AgentManifestBinding | None = None diff --git a/tests/unit/test_session_state_store.py b/tests/unit/test_session_state_store.py new file mode 100644 index 00000000..1e6f0a48 --- /dev/null +++ b/tests/unit/test_session_state_store.py @@ -0,0 +1,209 @@ +"""Shared, persistent session-sensitivity state. + +Covers the two properties that separate a shared store from the in-process +default: the accumulated value survives a restart, and the read-modify-write is +serialised across gateway instances rather than only across coroutines in one. +""" + +from __future__ import annotations + +import multiprocessing as mp +import sqlite3 +import time +from pathlib import Path + +import pytest + +from cmcp_runtime.session.state import SENSITIVITY_ORDER, SessionState +from cmcp_runtime.session.store import ( + InMemorySessionStateStore, + SessionStateStore, + SqliteSessionStateStore, + StoredSensitivity, +) + +# ── Protocol conformance ────────────────────────────────────────────────────── + + +def test_both_stores_satisfy_the_protocol(tmp_path: Path): + assert isinstance(InMemorySessionStateStore(), SessionStateStore) + assert isinstance(SqliteSessionStateStore(tmp_path / "s.db"), SessionStateStore) + + +def test_only_the_sqlite_store_claims_to_be_shared(tmp_path: Path): + """A caller must be able to tell, because it cannot from the outside.""" + assert InMemorySessionStateStore().is_shared() is False + assert SqliteSessionStateStore(tmp_path / "s.db").is_shared() is True + + +# ── Persistence across a restart ────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_accumulated_value_survives_a_restart(tmp_path: Path): + """The session identifier outlives the process; the value must too.""" + db = tmp_path / "session_state.db" + + store = SqliteSessionStateStore(db) + state = SessionState(session_id="s-restart", state_store=store) + await state.hydrate() + await state.apply_inspection("call-A", ["hipaa_phi"], False, True) + assert state.max_sensitivity == "hipaa_phi" + store.close() + + # A new process would open a new store over the same file. + revived_store = SqliteSessionStateStore(db) + revived = SessionState(session_id="s-restart", state_store=revived_store) + assert revived.max_sensitivity == "public" + adopted = await revived.hydrate() + + assert adopted is True + assert revived.max_sensitivity == "hipaa_phi" + assert revived.sensitivity_raised_by_call == "call-A" + + +@pytest.mark.asyncio +async def test_hydrate_seeds_a_session_the_store_has_not_seen(tmp_path: Path): + store = SqliteSessionStateStore(tmp_path / "s.db") + state = SessionState(session_id="s-new", state_store=store) + + assert await state.hydrate() is False + assert store.load("s-new") == StoredSensitivity("public", None, None, 0) + + +# ── Sharing across instances ────────────────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_a_second_instance_enforces_what_the_first_accumulated(tmp_path: Path): + """The ratchet must hold per session, not per gateway instance. + + Without a shared store, an agent that reads sensitive data through instance A + and egresses through instance B is evaluated by an instance that never saw + the read. + """ + store = SqliteSessionStateStore(tmp_path / "s.db") + a = SessionState(session_id="s-shared", state_store=store) + b = SessionState(session_id="s-shared", state_store=store) + + await a.apply_inspection("call-read", ["hipaa_phi"], False, True) + await b.hydrate() + + assert b.max_sensitivity == "hipaa_phi" + + +@pytest.mark.asyncio +async def test_the_stored_value_only_rises(tmp_path: Path): + """An instance holding a lower value must not lower the shared one.""" + store = SqliteSessionStateStore(tmp_path / "s.db") + high = SessionState(session_id="s-mono", state_store=store) + await high.apply_inspection("call-1", ["hipaa_phi"], False, True) + + low = SessionState(session_id="s-mono", state_store=store) + await low.apply_inspection("call-2", ["public"], False, True) + + assert store.load("s-mono").max_sensitivity == "hipaa_phi" + assert low.max_sensitivity == "hipaa_phi" + + +@pytest.mark.asyncio +async def test_a_reset_on_one_instance_closes_the_session_for_another(tmp_path: Path): + """A response on instance B must not raise the successor A opened.""" + store = SqliteSessionStateStore(tmp_path / "s.db") + a = SessionState(session_id="s-reset", state_store=store) + b = SessionState(session_id="s-reset", state_store=store) + await a.hydrate() + await b.hydrate() + await a.apply_inspection("call-A", ["pii"], False, True) + + generation = b.reset_count # B's view, taken before the reset + old_id, new_id, closed = await a.apply_reset(reason="op", authorized_by="operator") + + applied = await b.apply_inspection( + "call-B", ["confidential"], False, True, for_reset_count=generation + ) + + assert applied is False + assert closed.session_id == old_id + assert closed.max_sensitivity == "pii" + assert store.closed(old_id)["max_sensitivity"] == "pii" + assert store.load(new_id).max_sensitivity == "public" + + +# ── Cross-process serialisation ─────────────────────────────────────────────── + + +def _hold_reserved_lock(db_path: str, held_for_s: float, ready, done) -> None: + """Another gateway instance holding the store's write lock.""" + conn = sqlite3.connect(db_path, isolation_level=None) + conn.execute("PRAGMA busy_timeout=5000") + conn.execute("BEGIN IMMEDIATE") + ready.set() + time.sleep(held_for_s) + conn.execute("COMMIT") + conn.close() + done.set() + + +@pytest.mark.asyncio +async def test_the_write_lock_serialises_across_processes(tmp_path: Path): + """An asyncio.Lock cannot do this: it is invisible to another process. + + A real second process takes the store's RESERVED lock. This process must wait + for it rather than interleave, which is what makes the serialisation span + gateway instances instead of only coroutines. + """ + db = tmp_path / "s.db" + store = SqliteSessionStateStore(db, busy_timeout_ms=10_000) + state = SessionState(session_id="s-xproc", state_store=store) + await state.hydrate() + + ctx = mp.get_context("spawn") + ready, done = ctx.Event(), ctx.Event() + holder = ctx.Process( + target=_hold_reserved_lock, args=(str(db), 1.5, ready, done), daemon=True + ) + holder.start() + try: + assert ready.wait(timeout=15), "helper process never took the lock" + assert not done.is_set() + + started = time.monotonic() + await state.apply_inspection("call-x", ["pii"], False, True) + waited = time.monotonic() - started + + # The write could only land once the other process committed. + assert done.is_set() + assert waited > 0.5, f"did not wait for the other instance (waited {waited:.2f}s)" + assert store.load("s-xproc").max_sensitivity == "pii" + finally: + holder.join(timeout=15) + + +# ── The in-process default is unchanged ─────────────────────────────────────── + + +@pytest.mark.asyncio +async def test_no_store_configured_keeps_the_previous_behaviour(): + state = SessionState(session_id="s-plain") + assert state.state_store is None + + assert await state.hydrate() is False + assert await state.apply_inspection("call-A", ["pii"], False, True) is True + assert state.max_sensitivity == "pii" + + _, _, closed = await state.apply_reset(reason="op", authorized_by="test") + assert closed.max_sensitivity == "pii" + assert state.max_sensitivity == "public" + assert state.max_sensitivity in SENSITIVITY_ORDER + + +@pytest.mark.asyncio +async def test_in_memory_store_serialises_but_does_not_persist(): + store = InMemorySessionStateStore() + state = SessionState(session_id="s-mem", state_store=store) + await state.apply_inspection("call-A", ["pii"], False, True) + + assert store.load("s-mem").max_sensitivity == "pii" + # A different store instance stands for a restart: nothing carried over. + assert InMemorySessionStateStore().load("s-mem") is None