From 972d3d00fc7a4dbecaa088feb7758a8c0caa8994 Mon Sep 17 00:00:00 2001
From: Imran Siddique <imran.siddique@opaque.co>
Date: Sat, 6 Jun 2026 15:36:04 -0700
Subject: [PATCH] fix(session): add asyncio.Lock to guard concurrent state
 mutations (AUTH-002)

Concurrent tool-call coroutines and session-reset requests could
interleave after an await boundary, corrupting max_sensitivity or
losing injection events. Added `mutation_lock: asyncio.Lock` to
SessionState and wrapped update_from_inspection() in proxy.py and
reset() in server.py with `async with self._session.mutation_lock`.

Closes #164

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
---
 src/cmcp_gateway/mcp/proxy.py     | 14 ++++++++------
 src/cmcp_gateway/mcp/server.py    | 10 ++++++----
 src/cmcp_gateway/session/state.py |  3 +++
 tests/unit/test_session.py        | 32 +++++++++++++++++++++++++++++++
 4 files changed, 49 insertions(+), 10 deletions(-)

diff --git a/src/cmcp_gateway/mcp/proxy.py b/src/cmcp_gateway/mcp/proxy.py
index 291d362a..ce214a95 100644
--- a/src/cmcp_gateway/mcp/proxy.py
+++ b/src/cmcp_gateway/mcp/proxy.py
@@ -384,14 +384,16 @@ async def call_tool(
             )
 
         # Step 4: session update from response sensitivity
+        # AUTH-002: lock protects against race with concurrent session reset requests.
         response_sensitivity = getattr(agt_result, "sensitivity_tags", [])
         injection_detected = getattr(agt_result, "injection_detected", False)
-        self._session.update_from_inspection(
-            call_id=call_id,
-            sensitivity_tags=response_sensitivity or [entry.sensitivity_level],
-            injection_detected=injection_detected,
-            response_allowed=True,
-        )
+        async with self._session.mutation_lock:
+            self._session.update_from_inspection(
+                call_id=call_id,
+                sensitivity_tags=response_sensitivity or [entry.sensitivity_level],
+                injection_detected=injection_detected,
+                response_allowed=True,
+            )
 
         # Step 5: egress Cedar policy check
         # Derive response bytes for size accounting and egress evaluation.
diff --git a/src/cmcp_gateway/mcp/server.py b/src/cmcp_gateway/mcp/server.py
index fc831b40..e9470255 100644
--- a/src/cmcp_gateway/mcp/server.py
+++ b/src/cmcp_gateway/mcp/server.py
@@ -373,10 +373,12 @@ async def _session_reset(self, request: Request) -> Response:
             return JSONResponse(
                 {"error": f"session_id={session_id} not found"}, status_code=404
             )
-        old_id, new_id = self._session.reset(
-            reason="operator reset via API",
-            authorized_by="api",
-        )
+        # AUTH-002: lock guards against a concurrent tool-call coroutine modifying sensitivity.
+        async with self._session.mutation_lock:
+            old_id, new_id = self._session.reset(
+                reason="operator reset via API",
+                authorized_by="api",
+            )
         self._audit_chain.append(
             "session_reset",
             call_id=None,
diff --git a/src/cmcp_gateway/session/state.py b/src/cmcp_gateway/session/state.py
index 5f705523..1b23cbc9 100644
--- a/src/cmcp_gateway/session/state.py
+++ b/src/cmcp_gateway/session/state.py
@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from dataclasses import dataclass, field
 from datetime import UTC, datetime
 from uuid import uuid4
@@ -55,6 +56,8 @@ class SessionState:
     suspicious_sequences: int = 0
     attestation_stale: bool = False
     catalog_drift: bool = False
+    # 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)
 
     def update_from_inspection(
         self,
diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py
index eefcf0be..298b44a2 100644
--- a/tests/unit/test_session.py
+++ b/tests/unit/test_session.py
@@ -2,6 +2,10 @@
 
 from __future__ import annotations
 
+import asyncio
+
+import pytest
+
 from cmcp_gateway.session.state import SENSITIVITY_ORDER, SessionState, _max_sensitivity
 
 # ── _max_sensitivity ──────────────────────────────────────────────────────────
@@ -107,3 +111,31 @@ def test_update_highest_tag_wins_per_update():
     state.update_from_inspection("c1", ["pii", "mnpi", "confidential"], False, True)
     assert state.max_sensitivity in ("mnpi", "hipaa_phi", "trade_secret")
     assert SENSITIVITY_ORDER[state.max_sensitivity] == 3
+
+
+# ── AUTH-002: asyncio.Lock guards concurrent mutations ────────────────────────
+
+def test_session_state_has_mutation_lock():
+    """AUTH-002 — SessionState must expose an asyncio.Lock for concurrent-mutation protection."""
+    import asyncio
+    state = SessionState(session_id="s-lock")
+    assert isinstance(state.mutation_lock, asyncio.Lock)
+
+
+@pytest.mark.asyncio
+async def test_concurrent_update_and_reset_do_not_corrupt_state():
+    """AUTH-002 — concurrent update_from_inspection and reset must not leave state inconsistent."""
+    state = SessionState(session_id="s-concurrent")
+
+    async def _update():
+        async with state.mutation_lock:
+            state.update_from_inspection("c1", ["pii"], False, True)
+
+    async def _reset():
+        async with state.mutation_lock:
+            state.reset(reason="concurrent reset", authorized_by="test")
+
+    # Run 10 interleaved updates and resets; state must be valid throughout.
+    tasks = [_update() for _ in range(5)] + [_reset() for _ in range(5)]
+    await asyncio.gather(*tasks)
+    assert state.max_sensitivity in SENSITIVITY_ORDER