fix: bound MCP request cancellation and recovery

tests/test_mcp_memory_ergonomics_helpers.py

@@ -463,9 +463,9 @@ def test_tee_and_bulk_create_by_path_update_by_identity(mod):
         "content_text": payload.get("content_text", existing["content_text"]),
     }
     mod.index_node = lambda node: None
-    mod.acquire_project_lock = lambda args: {"id": "lock-1", "lease_token": "lease-1"}
+    mod.acquire_project_lock = lambda args, **kwargs: {"id": "lock-1", "lease_token": "lease-1"}
     mod.lock_headers = lambda lock: {"X-Lock": lock["id"]}
-    mod.release_project_lock = lambda args: None
+    mod.release_project_lock = lambda args, **kwargs: None
     mod.request_json = lambda method, path, payload=None, **kwargs: writes.append((method, path, payload, ())) or {
         **existing,
         "content_text": payload.get("content_text", existing["content_text"]) if payload else existing["content_text"],

tests/test_mcp_memory_node_encryption.py

@@ -393,9 +393,9 @@ def run():
     assert patch_body["encrypted_content_envelope"]["aad_subject"] == "node_id:node-1"
 
     bulk_writes = []
-    mod.acquire_project_lock = lambda args: {"id": "lock-1", "lease_token": "lease-1"}
+    mod.acquire_project_lock = lambda args, **kwargs: {"id": "lock-1", "lease_token": "lease-1"}
     mod.lock_headers = lambda lock: {"X-Lock": lock["id"]}
-    mod.release_project_lock = lambda args: None
+    mod.release_project_lock = lambda args, **kwargs: None
     mod.try_get_node_by_path = lambda path, allow_agent=False: None
     mod.request_json = lambda method, path, payload=None, **kwargs: bulk_writes.append((method, path, payload, kwargs)) or {
         "id": "node-bulk",

tests/test_mcp_supervisor.py

@@ -75,6 +75,10 @@ class SupervisorClient:
         worker_path: pathlib.Path,
         project_root: pathlib.Path | None = None,
         watch_paths: tuple[pathlib.Path, ...] = (),
+        response_timeout: int = 5,
+        cancel_grace: int = 5,
+        max_in_flight: int = 32,
+        write_timeout: int = 5,
     ):
         environment = os.environ.copy()
         environment["NAPSEER_MCP_WORKER_PATH"] = str(worker_path)
@@ -82,7 +86,10 @@ class SupervisorClient:
             environment["NAPSEER_MCP_WORKER_WATCH_PATHS"] = os.pathsep.join(
                 str(path) for path in watch_paths
             )
-        environment["NAPSEER_MCP_RESPONSE_TIMEOUT_SECONDS"] = "5"
+        environment["NAPSEER_MCP_RESPONSE_TIMEOUT_SECONDS"] = str(response_timeout)
+        environment["NAPSEER_MCP_CANCEL_GRACE_SECONDS"] = str(cancel_grace)
+        environment["NAPSEER_MCP_MAX_IN_FLIGHT_REQUESTS"] = str(max_in_flight)
+        environment["NAPSEER_MCP_WRITE_TIMEOUT_SECONDS"] = str(write_timeout)
         environment["NAPSEER_TELEMETRY"] = "0"
         if project_root is not None:
             environment["NAPSEER_PROJECT_ROOT"] = str(project_root)
@@ -92,6 +99,7 @@ class SupervisorClient:
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             env=environment,
+            bufsize=0,
         )
 
     def request(self, request_id: int, method: str, params=None):
@@ -321,3 +329,201 @@ def test_worker_crash_returns_error_without_closing_transport(tmp_path):
         assert client.process.poll() is None
     finally:
         client.close()
+
+
+def write_lifecycle_worker(path, version=1):
+    # The actual worker dispatcher owns cancellation/completion; the handler
+    # supplies deterministic slow and uncooperative work without credentials.
+    path.write_text(f'''import json, os, pathlib, runpy, sys, time
+sys.path.insert(0, {str(REAL_WORKER.parent)!r})
+mod = runpy.run_path({str(REAL_WORKER)!r})
+root = pathlib.Path({str(path.parent)!r})
+def handle(message):
+    name = (message.get("params") or {{}}).get("name")
+    if name in ("slow", "stuck"):
+        (root / (str(message["id"]) + ".started")).write_text("started")
+        try:
+            deadline = time.monotonic() + 10
+            while time.monotonic() < deadline:
+                if name == "slow":
+                    mod["raise_if_mcp_request_cancelled"]()
+                time.sleep(0.01)
+        finally:
+            time.sleep(0.3)
+            (root / (str(message["id"]) + ".cleaned")).write_text("cleaned")
+    return mod["rpc_result"](message.get("id"), {{"pid": os.getpid(), "version": {version}}})
+runtime = mod["McpRequestRuntime"](handler=handle)
+try:
+    for line in sys.stdin:
+        runtime.dispatch(json.loads(line))
+finally:
+    runtime.close()
+''')
+
+
+def wait_for_path(path, timeout=3):
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if path.exists():
+            return
+        time.sleep(0.01)
+    raise AssertionError("worker did not reach the expected lifecycle boundary")
+
+
+def test_cancelled_work_counts_until_cleanup_before_source_replacement(tmp_path):
+    worker_path = tmp_path / "worker.py"
+    write_lifecycle_worker(worker_path)
+    client = SupervisorClient(worker_path, project_root=tmp_path)
+    try:
+        client.send(1, "tools/call", {"name": "slow"})
+        wait_for_path(tmp_path / "1.started")
+        client.notify("notifications/cancelled", {"requestId": 1})
+        write_lifecycle_worker(worker_path, version=2)
+        during_cleanup = client.request(2, "tools/list")
+        assert during_cleanup["result"]["version"] == 1
+        wait_for_path(tmp_path / "1.cleaned")
+        time.sleep(0.05)
+        after_cleanup = client.request(3, "tools/list")
+        assert after_cleanup["result"]["version"] == 2
+        assert after_cleanup["result"]["pid"] != during_cleanup["result"]["pid"]
+    finally:
+        client.close()
+
+
+def test_timeout_cancels_work_and_keeps_transport_usable(tmp_path):
+    worker_path = tmp_path / "worker.py"
+    write_lifecycle_worker(worker_path)
+    client = SupervisorClient(worker_path, project_root=tmp_path, response_timeout=1)
+    try:
+        before = client.request(1, "tools/list")
+        failed = client.request(2, "tools/call", {"name": "slow"})
+        assert failed["error"]["code"] == -32098
+        wait_for_path(tmp_path / "2.cleaned")
+        recovered = client.request(3, "tools/list")
+        assert recovered["result"]["pid"] == before["result"]["pid"]
+    finally:
+        client.close()
+
+
+def test_cancelled_work_cannot_bypass_admission_limit(tmp_path):
+    worker_path = tmp_path / "worker.py"
+    write_lifecycle_worker(worker_path)
+    client = SupervisorClient(worker_path, project_root=tmp_path, max_in_flight=2)
+    try:
+        client.send(1, "tools/call", {"name": "stuck"})
+        wait_for_path(tmp_path / "1.started")
+        client.notify("notifications/cancelled", {"requestId": 1})
+        client.send(2, "tools/call", {"name": "stuck"})
+        wait_for_path(tmp_path / "2.started")
+        client.send(3, "tools/call", {"name": "stuck"})
+        rejected = client.receive(timeout=1)
+        assert rejected["id"] == 3
+        assert rejected["error"]["code"] == -32097
+        assert not (tmp_path / "3.started").exists()
+    finally:
+        client.close()
+
+
+def test_uncooperative_timeout_restarts_worker_with_no_replay(tmp_path):
+    worker_path = tmp_path / "worker.py"
+    write_lifecycle_worker(worker_path)
+    client = SupervisorClient(worker_path, project_root=tmp_path,
+                              response_timeout=1, cancel_grace=1)
+    try:
+        before = client.request(1, "tools/list")
+        failed = client.request(2, "tools/call", {"name": "stuck"})
+        assert failed["error"]["code"] == -32098
+        time.sleep(1.3)
+        recovered = client.request(3, "tools/list")
+        assert recovered["result"]["pid"] != before["result"]["pid"]
+        assert not (tmp_path / "2.cleaned").exists()
+        assert client.process.poll() is None
+    finally:
+        client.close()
+
+
+def test_reused_id_is_rejected_until_cancelled_request_finishes(tmp_path):
+    worker_path = tmp_path / "worker.py"
+    write_lifecycle_worker(worker_path)
+    client = SupervisorClient(worker_path, project_root=tmp_path)
+    try:
+        client.send(1, "tools/call", {"name": "slow"})
+        wait_for_path(tmp_path / "1.started")
+        client.notify("notifications/cancelled", {"requestId": 1})
+        rejected = client.request(1, "tools/list")
+        assert rejected["error"]["code"] == -32097
+        wait_for_path(tmp_path / "1.cleaned")
+        time.sleep(0.05)
+        assert client.request(1, "tools/list")["result"]["version"] == 1
+    finally:
+        client.close()
+
+
+def test_blocked_worker_stdin_has_a_deadline_and_recovers(tmp_path):
+    worker_path = tmp_path / "worker.py"
+    worker_path.write_text("import time\ntime.sleep(60)\n")
+    client = SupervisorClient(worker_path, project_root=tmp_path, write_timeout=1)
+    try:
+        client.send(1, "tools/call", {"name": "large", "arguments": {"data": "x" * 262144}})
+        failed = client.receive(timeout=3)
+        assert failed["error"]["code"] == -32098
+        write_worker(worker_path, 2)
+        assert client.request(2, "tools/list")["result"]["version"] == 2
+    finally:
+        client.close()
+
+
+def test_eof_stops_worker_with_unfinished_requests(tmp_path):
+    worker_path = tmp_path / "worker.py"
+    write_lifecycle_worker(worker_path)
+    client = SupervisorClient(worker_path, project_root=tmp_path)
+    try:
+        worker_pid = client.request(1, "tools/list")["result"]["pid"]
+        client.send(2, "tools/call", {"name": "stuck"})
+        wait_for_path(tmp_path / "2.started")
+        client.close()
+        assert client.process.returncode == 0
+        try:
+            os.kill(worker_pid, 0)
+        except ProcessLookupError:
+            pass
+        else:
+            raise AssertionError("worker survived supervisor EOF")
+    finally:
+        client.close()
+
+
+def test_initialize_timeout_retires_uncancellable_worker(tmp_path):
+    worker_path = tmp_path / "worker.py"
+    worker_path.write_text('''import json, os, sys, time
+for line in sys.stdin:
+    msg = json.loads(line)
+    if msg.get("method") == "initialize":
+        time.sleep(10)
+    print(json.dumps({"jsonrpc": "2.0", "id": msg["id"], "result": {"pid": os.getpid()}}), flush=True)
+''')
+    client = SupervisorClient(worker_path, project_root=tmp_path, response_timeout=1)
+    try:
+        before = client.request(1, "tools/list")
+        failed = client.request(2, "initialize")
+        assert failed["error"]["code"] == -32098
+        recovered = client.request(3, "tools/list")
+        assert recovered["result"]["pid"] != before["result"]["pid"]
+    finally:
+        client.close()
+
+
+def test_real_worker_contract_exposes_recovery_without_private_notifications(tmp_path, monkeypatch):
+    monkeypatch.setenv("NAPSEER_TOOL_PROFILES", "maintenance")
+    client = SupervisorClient(REAL_WORKER, project_root=tmp_path)
+    try:
+        assert "result" in client.request(1, "initialize")
+        response = client.request(2, "tools/call", {"name": "nap_contract", "arguments": {}})
+        contract = response["result"]["structuredContent"]["request_lifecycle"]
+        assert contract["max_unfinished_requests"] == 32
+        assert contract["supervised"] is True
+        assert contract["rejected_error_code"] == -32097
+        assert contract["uncertain_error_code"] == -32098
+        assert not (tmp_path / ".napseer").exists()
+    finally:
+        client.close()

tests/test_mcp_worker_lifecycle.py

@@ -4,6 +4,8 @@ import sys
 import threading
 import time
 
+import pytest
+
 
 def load_module():
     script_path = (
@@ -242,3 +244,160 @@ def test_auth_renewal_reuses_token_refreshed_by_another_request(monkeypatch):
         "refresh_expires_at": "later",
         "project_id": "project-1",
     }
+
+
+def test_tool_call_notification_cannot_execute_a_mutation(monkeypatch):
+    mod = load_module()
+    calls = []
+    monkeypatch.setattr(mod, "call_tool", lambda *args: calls.append(args))
+    assert mod.handle({"jsonrpc": "2.0", "method": "tools/call", "params": {
+        "name": "nap_create_node", "arguments": {"name": "unacknowledged"},
+    }}) is None
+    assert calls == []
+
+
+def test_worker_admission_is_bounded_and_recovers():
+    mod = load_module()
+    release = threading.Event()
+    emitted = []
+    executed = []
+
+    def handler(message):
+        executed.append(message["id"])
+        release.wait(2)
+        return mod.rpc_result(message["id"], {"ok": True})
+
+    runtime = mod.McpRequestRuntime(handler=handler, sender=emitted.append,
+                                    max_workers=2, max_in_flight=2)
+    try:
+        for request_id in (1, 2, 3):
+            runtime.dispatch({"jsonrpc": "2.0", "id": request_id, "method": "tools/call"})
+        wait_until(lambda: any(item["id"] == 3 for item in emitted))
+        rejection = next(item for item in emitted if item["id"] == 3)
+        assert rejection["error"]["code"] == -32097
+        assert 3 not in executed
+        release.set()
+        wait_until(lambda: len(emitted) == 3)
+        runtime.dispatch({"jsonrpc": "2.0", "id": 4, "method": "tools/call"})
+        wait_until(lambda: any(item["id"] == 4 for item in emitted))
+    finally:
+        release.set()
+        runtime.close()
+
+
+def test_cancelled_mutation_waiting_for_lock_never_starts(monkeypatch):
+    mod = load_module()
+    cancelled = threading.Event()
+    waiting = threading.Event()
+    calls = []
+    errors = []
+    monkeypatch.setattr(mod, "call_tool_impl", lambda *args: calls.append(args))
+    monkeypatch.setattr(mod, "send_telemetry_event_async", lambda *a, **kw: None)
+
+    def run():
+        mod.MCP_REQUEST_LOCAL.cancelled = cancelled
+        waiting.set()
+        try:
+            mod.call_tool("nap_create_node", {})
+        except mod.SafeToolError as exc:
+            errors.append(exc)
+
+    with mod.MCP_MUTATION_LOCK:
+        thread = threading.Thread(target=run)
+        thread.start()
+        assert waiting.wait(1)
+        cancelled.set()
+    thread.join(2)
+    assert not thread.is_alive()
+    assert calls == []
+    assert len(errors) == 1
+
+
+def test_cancelled_queue_entries_remain_bounded_until_consumed():
+    mod = load_module()
+    release = threading.Event()
+    executed = []
+    emitted = []
+
+    def handler(message):
+        executed.append(message["id"])
+        release.wait(2)
+        return mod.rpc_result(message["id"], {"ok": True})
+
+    runtime = mod.McpRequestRuntime(handler=handler, sender=emitted.append,
+                                    max_workers=2, max_in_flight=3)
+    try:
+        for request_id in (1, 2, 3):
+            runtime.dispatch({"id": request_id, "method": "tools/call"})
+        wait_until(lambda: len(executed) == 2)
+        runtime.cancel(3)
+        runtime.dispatch({"id": 4, "method": "tools/call"})
+        wait_until(lambda: any(item["id"] == 4 for item in emitted))
+        assert emitted[0]["error"]["code"] == -32097
+        release.set()
+        wait_until(lambda: len(emitted) == 3)
+        assert executed == [1, 2]
+    finally:
+        release.set()
+        runtime.close()
+
+
+@pytest.mark.parametrize("cancel_at", ["acquire", "operation"])
+def test_cancelled_write_releases_its_original_project_lock(tmp_path, monkeypatch, cancel_at):
+    import http.server
+    import json
+
+    monkeypatch.setenv("NAPSEER_PROJECT_ROOT", str(tmp_path))
+    mod = load_module()
+    cancelled = threading.Event()
+    requests = []
+    mutation_ran = []
+
+    class Handler(http.server.BaseHTTPRequestHandler):
+        def log_message(self, *args):
+            pass
+
+        def do_POST(self):
+            self.rfile.read(int(self.headers.get("Content-Length", 0)))
+            requests.append(("POST", self.path))
+            if cancel_at == "acquire":
+                cancelled.set()
+            body = json.dumps({"id": "lock-test", "lease_token": "synthetic-lease"}).encode()
+            self.send_response(200)
+            self.send_header("Content-Length", str(len(body)))
+            self.end_headers()
+            self.wfile.write(body)
+
+        def do_DELETE(self):
+            requests.append(("DELETE", self.path))
+            self.send_response(204)
+            self.end_headers()
+
+    server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler)
+    thread = threading.Thread(target=server.serve_forever, daemon=True)
+    thread.start()
+    mod.BASE_URL = f"http://127.0.0.1:{server.server_port}"
+    mod.TOKEN = "synthetic-access"
+    mod.DEFAULT_PROJECT_ID = "original-project"
+    monkeypatch.setattr(mod, "refresh_public_auth_state", lambda: None)
+    mod.MCP_REQUEST_LOCAL.cancelled = cancelled
+
+    def operation(_headers):
+        mutation_ran.append(True)
+        mod.DEFAULT_PROJECT_ID = "different-project"
+        cancelled.set()
+        mod.raise_if_mcp_request_cancelled()
+
+    try:
+        with pytest.raises(mod.SafeToolError):
+            mod.with_project_lock({}, operation)
+        assert requests == [
+            ("POST", "/v1/projects/original-project/locks"),
+            ("DELETE", "/v1/projects/original-project/locks/lock-test"),
+        ]
+        assert mutation_ran == ([] if cancel_at == "acquire" else [True])
+    finally:
+        del mod.MCP_REQUEST_LOCAL.cancelled
+        server.shutdown()
+        server.server_close()
+        thread.join(2)