diff --git a/README.md b/README.md index 166b7ef..67d0ed1 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,24 @@ replaceable runtime, set `NAPSEER_MCP_WORKER_WATCH_PATHS` to the after any watched runtime is atomically replaced, preserving the same bounded, non-replaying request lifecycle used for direct workers. +MCP admits at most 32 unfinished requests by default, including cancelled work +that is still cleaning up. `-32097` means a request was rejected before +execution; retry after capacity is available. `-32098` means the result +is uncertain: inspect current state before retrying a mutation. Cancellation +prevents queued mutations from starting and attempts to release acquired locks. +It cannot undo a mutation already accepted by the API. + +The supervisor waits up to 60 seconds for a response, then requests cancellation. +After 10 seconds without a cleanup acknowledgement, it terminates the worker and +fails affected requests without replay. Worker input writes have a five-second +deadline. Override these limits with `NAPSEER_MCP_RESPONSE_TIMEOUT_SECONDS`, +`NAPSEER_MCP_CANCEL_GRACE_SECONDS`, `NAPSEER_MCP_WRITE_TIMEOUT_SECONDS`, and +`NAPSEER_MCP_MAX_IN_FLIGHT_REQUESTS`. New source loads only after unfinished work +has drained. `nap_contract` exposes the runtime's request lifecycle contract. +These process recovery guarantees require the supervisor; the direct worker is +for isolated protocol debugging. An already-running older supervisor needs a +client reconnect after updating. + The public operator surface is intentionally small: ```text diff --git a/resources/scripts/napseer_mcp_server.py b/resources/scripts/napseer_mcp_server.py index 1cf23df..51c2663 100755 --- a/resources/scripts/napseer_mcp_server.py +++ b/resources/scripts/napseer_mcp_server.py @@ -586,6 +586,9 @@ MCP_MUTATION_LOCK = threading.RLock() MCP_STDOUT_LOCK = threading.Lock() MCP_REQUEST_LOCAL = threading.local() MCP_MAX_WORKERS = max(2, int(os.environ.get("NAPSEER_MCP_MAX_WORKERS", "8"))) +MCP_MAX_IN_FLIGHT_REQUESTS = max(2, int(os.environ.get("NAPSEER_MCP_MAX_IN_FLIGHT_REQUESTS", "32"))) +MCP_SUPERVISED = os.environ.get("NAPSEER_MCP_SUPERVISED") == "1" +MCP_REQUEST_FINISHED = "notifications/napseer/requestFinished" MCP_BATCH_MAX_ACTIONS = 32 MCP_BATCH_MAX_PARALLEL = max( 1, @@ -6334,8 +6337,11 @@ def raise_if_mcp_request_cancelled(): raise SafeToolError("request_cancelled", "The MCP request was cancelled.") -def request_json(method, path, payload=None, token_required=True, retry_auth=True, extra_headers=None): - raise_if_mcp_request_cancelled() +def request_json(method, path, payload=None, token_required=True, retry_auth=True, extra_headers=None, + *, cancellation_sensitive=True, preserve_response_on_cancel=False, timeout_seconds=None): + if cancellation_sensitive: + raise_if_mcp_request_cancelled() + timeout = HTTP_TIMEOUT_SECONDS if timeout_seconds is None else timeout_seconds if token_required: # `nap project attach` may run in a separate process while this MCP # server stays alive. Reload the project auth file before building @@ -6353,19 +6359,21 @@ def request_json(method, path, payload=None, token_required=True, retry_auth=Tru request_token = TOKEN request = urllib.request.Request(f"{BASE_URL}{path}", data=body, headers=headers, method=method) try: - with api_urlopen(request, timeout=HTTP_TIMEOUT_SECONDS) as response: + with api_urlopen(request, timeout=timeout) as response: text = response.read().decode("utf-8") - raise_if_mcp_request_cancelled() + if cancellation_sensitive and not preserve_response_on_cancel: + raise_if_mcp_request_cancelled() return json.loads(text) if text else {} except (TimeoutError, socket.timeout) as exc: raise SafeToolError( "request_timeout", - f"Napseer request timed out after {HTTP_TIMEOUT_SECONDS}s.", + f"Napseer request timed out after {timeout}s.", ) from exc except urllib.error.HTTPError as exc: error_body = exc.read(65536).decode("utf-8", errors="replace") exc.close() - raise_if_mcp_request_cancelled() + if cancellation_sensitive: + raise_if_mcp_request_cancelled() if exc.code == 401 and token_required and retry_auth: try: # The attach flow may have completed after this request was @@ -6380,11 +6388,16 @@ def request_json(method, path, payload=None, token_required=True, retry_auth=Tru token_required=True, retry_auth=False, extra_headers=extra_headers, + cancellation_sensitive=cancellation_sensitive, + preserve_response_on_cancel=preserve_response_on_cancel, + timeout_seconds=timeout_seconds, ) renew_auth(stale_token=request_token) except Exception as renew_exc: raise safe_auth_refresh_error() from renew_exc - return request_json(method, path, payload, token_required=True, retry_auth=False, extra_headers=extra_headers) + return request_json(method, path, payload, token_required=True, retry_auth=False, + extra_headers=extra_headers, cancellation_sensitive=cancellation_sensitive, + preserve_response_on_cancel=preserve_response_on_cancel, timeout_seconds=timeout_seconds) raise safe_http_error( exc, operation=f"{method} request", @@ -6447,9 +6460,9 @@ def lock_headers(lock): } -def acquire_project_lock(args=None): +def acquire_project_lock(args=None, *, project_id=None): args = args or {} - project_id = resolve_project_id(args) + project_id = project_id or resolve_project_id(args) payload = { "scope_type": args.get("scope_type", "project"), "scope_key": args.get("scope_key", "*"), @@ -6457,7 +6470,16 @@ def acquire_project_lock(args=None): "holder_label": args.get("holder_label") or AUTH.get("worker_name") or "local MCP", "ttl_seconds": int(args.get("ttl_seconds", 600)), } - return request_json("POST", f"/v1/projects/{project_id}/locks", payload) + # Retain the lease even if cancellation arrives while the server grants it; + # otherwise the caller loses the only information needed to release it. + lock = request_json("POST", f"/v1/projects/{project_id}/locks", payload, + preserve_response_on_cancel=True) + try: + raise_if_mcp_request_cancelled() + except SafeToolError: + cleanup_project_lock(project_id, lock) + raise + return lock def renew_project_lock(args): @@ -6474,8 +6496,8 @@ def renew_project_lock(args): ) -def release_project_lock(args): - project_id = resolve_project_id(args) +def release_project_lock(args, *, cleanup=False): + project_id = args["project_id"] if cleanup else resolve_project_id(args) lock_id = args.get("lock_id") lease_token = args.get("lease_token") if not lock_id or not lease_token: @@ -6485,6 +6507,9 @@ def release_project_lock(args): f"/v1/projects/{project_id}/locks/{lock_id}", payload=None, extra_headers={"X-Napseer-Lock-Token": lease_token}, + cancellation_sensitive=not cleanup, + retry_auth=not cleanup, + timeout_seconds=5 if cleanup else None, ) @@ -6493,21 +6518,31 @@ def list_project_locks(args): return request_json("GET", f"/v1/projects/{project_id}/locks") -def with_project_lock(args, operation): - lock = acquire_project_lock(args) +def with_project_lock(args, operation, *, project_id=None): + project_id = project_id or resolve_project_id(args) + lock = acquire_project_lock(args, project_id=project_id) try: + raise_if_mcp_request_cancelled() return operation(lock_headers(lock)) finally: - try: - release_project_lock({"lock_id": lock["id"], "lease_token": lock["lease_token"], "project_id": args.get("project_id")}) - except Exception: - pass + cleanup_project_lock(project_id, lock) + + +def cleanup_project_lock(project_id, lock): + try: + release_project_lock({"lock_id": lock["id"], "lease_token": lock["lease_token"], + "project_id": project_id}, cleanup=True) + except Exception: + # Cleanup cannot erase the original outcome. The server lease still + # bounds recovery if transport/auth prevents this single release attempt. + pass def request_project_write(method, path, payload, project_id, purpose, scope_type="project", scope_key="*"): return with_project_lock( {"project_id": project_id, "scope_type": scope_type, "scope_key": scope_key, "purpose": purpose}, lambda headers: request_json(method, path, payload, extra_headers=headers), + project_id=project_id, ) @@ -13424,7 +13459,7 @@ def move_folder(args): planned.append({"from": node.get("full_path"), "to_folder": next_folder}) links_rewritten = sum(len(item["references_to_rewrite"]) for item in reference_plan) return {"ok": True, "changed": False, "dry_run": True, "moved": len(planned), "from_folder": from_folder, "to_folder": to_folder, "items": planned, "reference_plan": reference_plan, "links_rewritten": links_rewritten} - lock = acquire_project_lock({"scope_type": "project", "scope_key": "*", "purpose": f"move folder {from_folder} to {to_folder}"}) + lock = acquire_project_lock({"scope_type": "project", "scope_key": "*", "purpose": f"move folder {from_folder} to {to_folder}"}, project_id=project_id) updated = [] try: headers = lock_headers(lock) @@ -13436,10 +13471,7 @@ def move_folder(args): index_node(saved) updated.append(node_write_reference(saved)) finally: - try: - release_project_lock({"lock_id": lock["id"], "lease_token": lock["lease_token"]}) - except Exception: - pass + cleanup_project_lock(project_id, lock) links_rewritten = sum(len(item["references_to_rewrite"]) for item in reference_plan) return {"ok": True, "changed": bool(updated), "dry_run": False, "moved": len(updated), "from_folder": from_folder, "to_folder": to_folder, "paths": [item.get("path") for item in updated if item.get("path")], "items": updated, "reference_plan": reference_plan, "links_rewritten": links_rewritten} @@ -13452,7 +13484,7 @@ def bulk_upsert_nodes(args): if args.get("dry_run"): paths = [normalize_node_path(item.get("path")) for item in nodes if item.get("path")] return {"ok": True, "changed": False, "dry_run": True, "count": len(paths), "paths": paths, "items": []} - lock = acquire_project_lock({"scope_type": "project", "scope_key": "*", "purpose": args.get("purpose", "bulk upsert nodes")}) + lock = acquire_project_lock({"scope_type": "project", "scope_key": "*", "purpose": args.get("purpose", "bulk upsert nodes")}, project_id=project_id) results = [] try: headers = lock_headers(lock) @@ -13478,10 +13510,7 @@ def bulk_upsert_nodes(args): node = decrypt_node_for_return(saved, project_id=project_id) results.append({"created": created, **node_write_reference(node)}) finally: - try: - release_project_lock({"lock_id": lock["id"], "lease_token": lock["lease_token"]}) - except Exception: - pass + cleanup_project_lock(project_id, lock) return {"ok": True, "changed": bool(results), "dry_run": False, "count": len(results), "paths": [item.get("path") for item in results if item.get("path")], "items": results} @@ -16842,6 +16871,18 @@ def contract_profile(args=None): }, "core": sorted(CORE_TOOL_NAMES), } + profile["request_lifecycle"] = { + "max_unfinished_requests": MCP_MAX_IN_FLIGHT_REQUESTS, + "max_workers": MCP_MAX_WORKERS, + "supervised": MCP_SUPERVISED, + "cancellation": "cooperative; unfinished work remains admitted until cleanup completes", + "queued_mutations": "cancelled work does not start", + "replay": "never automatic; inspect authoritative state after an uncertain outcome", + "rejected_error_code": -32097, + "uncertain_error_code": -32098, + "cleanup": "one lock release attempt against the original project, without auth replay", + "recovery_owner": "stdio supervisor bounds response waits, cancellation grace, and worker replacement", + } return profile @@ -17918,6 +17959,7 @@ def call_tool(name, args): started_at = time.time() tool_name = telemetry_safe_text(name, 96) try: + raise_if_mcp_request_cancelled() if name == "nap_batch": # Batch validates its complete graph first, overlaps only reads, # and acquires the mutation lock around each write barrier. @@ -17929,6 +17971,7 @@ def call_tool(name, args): # auth/index state. Keep their execution order explicit while # allowing independent reads to proceed concurrently. with MCP_MUTATION_LOCK: + raise_if_mcp_request_cancelled() result = call_tool_impl(name, args) send_telemetry_event_async( "mcp_tool_call", @@ -17954,6 +17997,10 @@ def call_tool(name, args): def handle(message): + if "id" not in message: + # Only requests may execute tools. Unknown/extension notifications are + # one-way and must never invoke a mutation without an acknowledgement. + return None request_id = message.get("id") method = message.get("method") if method == "initialize": @@ -18001,12 +18048,14 @@ class McpRequestState: class McpRequestRuntime: """Bounded concurrent request dispatcher for the stdio worker.""" - def __init__(self, handler=handle, sender=send, max_workers=MCP_MAX_WORKERS): + def __init__(self, handler=handle, sender=send, max_workers=MCP_MAX_WORKERS, + max_in_flight=MCP_MAX_IN_FLIGHT_REQUESTS): self.handler = handler self.sender = sender self._lock = threading.Lock() self._requests = {} self._closed = False + self._max_in_flight = max(2, int(max_in_flight)) self._executor = concurrent.futures.ThreadPoolExecutor( max_workers=max(2, int(max_workers)), thread_name_prefix="napseer-mcp-tool", @@ -18019,20 +18068,31 @@ class McpRequestRuntime: def _execute(self, key, state): MCP_REQUEST_LOCAL.cancelled = state.cancelled try: + raise_if_mcp_request_cancelled() response = self.handler(state.message) except Exception: response = rpc_error(state.message.get("id"), -32603, "Internal error") finally: with self._lock: + cancelled = state.cancelled.is_set() if self._requests.get(key) is state: self._requests.pop(key, None) try: del MCP_REQUEST_LOCAL.cancelled except AttributeError: pass - if response is not None and not state.cancelled.is_set(): + if cancelled: + self._notify_finished(state) + elif response is not None: self.sender(response) + def _notify_finished(self, state): + if MCP_SUPERVISED: + # Private worker-to-supervisor acknowledgement, after handler + # cleanup. It is consumed locally and never reaches the MCP client. + self.sender({"jsonrpc": "2.0", "method": MCP_REQUEST_FINISHED, + "params": {"requestId": state.message.get("id")}}) + def cancel(self, request_id): key = self._key(request_id) with self._lock: @@ -18040,10 +18100,10 @@ class McpRequestRuntime: if state is None or state.message.get("method") == "initialize": return False state.cancelled.set() - future = state.future - if future is not None and future.cancel(): - self._requests.pop(key, None) - return True + # Keep queued work admitted until the executor consumes it. Future + # cancellation alone leaves its work item in an unbounded queue, + # allowing a cancel/submit loop to defeat the admission limit. + return True def dispatch(self, message): if not isinstance(message, dict): @@ -18059,7 +18119,8 @@ class McpRequestRuntime: # handled before the next request is accepted. Any accidental # handler return value is discarded because notifications never # receive responses. - self.handler(message) + if message.get("method") == "notifications/initialized": + self.handler(message) return None key = self._key(message.get("id")) @@ -18070,6 +18131,10 @@ class McpRequestRuntime: if key in self._requests: self.sender(rpc_error(message.get("id"), -32600, "Duplicate in-flight request id")) return None + if len(self._requests) >= self._max_in_flight: + self.sender(rpc_error(message.get("id"), -32097, + "Napseer MCP is busy; this request was not started. Retry later.")) + return None self._requests[key] = state state.future = self._executor.submit(self._execute, key, state) return state.future diff --git a/resources/scripts/napseer_mcp_supervisor.py b/resources/scripts/napseer_mcp_supervisor.py index 30c932f..adebfb9 100644 --- a/resources/scripts/napseer_mcp_supervisor.py +++ b/resources/scripts/napseer_mcp_supervisor.py @@ -8,8 +8,8 @@ session. This supervisor owns those client-facing pipes while treating - valid requests are forwarded without head-of-line blocking and responses are correlated by JSON-RPC request id; - notifications never create client-facing responses; -- cancellation releases the matching pending request and forwards the - notification to the worker; +- cancellation stops waiting for a result but retains unfinished work until + the worker acknowledges cleanup, with bounded grace before termination; - a worker exit fails every affected in-flight request once and never replays it automatically; - a replaced worker, or an explicitly configured runtime dependency behind a @@ -26,10 +26,12 @@ import json import os import pathlib import queue +import select import signal import subprocess import sys import threading +import time from dataclasses import dataclass, field from typing import BinaryIO, Callable @@ -54,6 +56,10 @@ MAX_IN_FLIGHT_REQUESTS = max( int(os.environ.get("NAPSEER_MCP_MAX_IN_FLIGHT_REQUESTS", "32")), ) WORKER_RESTARTED_ERROR = -32098 +REQUEST_REJECTED_ERROR = -32097 +CANCEL_GRACE_SECONDS = max(1, int(os.environ.get("NAPSEER_MCP_CANCEL_GRACE_SECONDS", "10"))) +WRITE_TIMEOUT_SECONDS = max(1, int(os.environ.get("NAPSEER_MCP_WRITE_TIMEOUT_SECONDS", "5"))) +REQUEST_FINISHED = "notifications/napseer/requestFinished" _CANCELLED = object() @@ -61,6 +67,10 @@ class WorkerUnavailable(RuntimeError): """The replaceable MCP worker could not complete a request.""" +class RequestRejected(WorkerUnavailable): + """The request was not sent to the worker and cannot have executed.""" + + def source_identity(path: pathlib.Path) -> tuple[int, int, int]: stat = path.stat() return stat.st_ino, stat.st_size, stat.st_mtime_ns @@ -112,16 +122,16 @@ def cancelled_request_id(message: bytes): return True, params.get("requestId") -def worker_error(message: bytes) -> bytes: +def worker_error(message: bytes, *, rejected=False) -> bytes: payload = { "jsonrpc": "2.0", "id": request_id(message), "error": { - "code": WORKER_RESTARTED_ERROR, + "code": REQUEST_REJECTED_ERROR if rejected else WORKER_RESTARTED_ERROR, "message": ( "Napseer MCP could not complete this request. The outcome may be " "uncertain; inspect state before retrying a mutation." - ), + ) if not rejected else "Napseer MCP is busy or this request id is still active. This request was not started; retry later.", }, } return json.dumps(payload, separators=(",", ":")).encode("utf-8") + b"\n" @@ -133,6 +143,8 @@ class PendingRequest: key: str generation: int response: queue.Queue = field(default_factory=lambda: queue.Queue(maxsize=1)) + cancelled: bool = False + cancel_timer: threading.Timer | None = None class Worker: @@ -156,7 +168,7 @@ class Worker: def _discard_stderr(self, stream: BinaryIO) -> None: try: - for _line in stream: + while stream.read(4096): pass finally: stream.close() @@ -180,6 +192,8 @@ class Worker: if pending.generation != generation: continue self._pending.pop(key, None) + if pending.cancel_timer is not None: + pending.cancel_timer.cancel() failed.append(pending) for pending in failed: self._signal(pending, error) @@ -189,18 +203,32 @@ class Worker: if payload is None: return False normalized = response if response.endswith(b"\n") else response + b"\n" + if payload.get("method") == REQUEST_FINISHED: + params = payload.get("params") or {} + if isinstance(params, dict) and "requestId" in params: + key = request_key(params["requestId"]) + with self._pending_lock: + pending = self._pending.get(key) + if pending is not None and pending.generation == generation and pending.cancelled: + self._pending.pop(key, None) + if pending.cancel_timer is not None: + pending.cancel_timer.cancel() + return True if "id" in payload: key = request_key(payload.get("id")) with self._pending_lock: pending = self._pending.get(key) if pending is not None and pending.generation == generation: self._pending.pop(key, None) + if pending.cancel_timer is not None: + pending.cancel_timer.cancel() else: pending = None if pending is not None: - self._signal(pending, normalized) + if not pending.cancelled: + self._signal(pending, normalized) return True - if "method" in payload and self.unsolicited is not None: + if generation == self.generation and "method" in payload and self.unsolicited is not None: self.unsolicited(normalized) return True @@ -240,12 +268,15 @@ class Worker: if process is None: return if process.poll() is None: - process.terminate() + try: + process.terminate() + except ProcessLookupError: + pass try: process.wait(timeout=2) except subprocess.TimeoutExpired: process.kill() - process.wait() + process.wait(timeout=2) for stream in (process.stdin, process.stdout): if stream is not None: try: @@ -265,6 +296,7 @@ class Worker: stdout=subprocess.PIPE, stderr=subprocess.PIPE, bufsize=0, + env={**os.environ, "NAPSEER_MCP_SUPERVISED": "1"}, ) except (OSError, subprocess.SubprocessError) as exc: raise WorkerUnavailable("unable to start Napseer MCP worker") from exc @@ -275,6 +307,7 @@ class Worker: self.generation += 1 generation = self.generation self.process = process + os.set_blocking(process.stdin.fileno(), False) self.identity = identity threading.Thread( target=self._discard_stderr, @@ -309,14 +342,30 @@ class Worker: self.start() def _write(self, process: subprocess.Popen[bytes], message: bytes) -> None: - with self._stdin_lock: + deadline = time.monotonic() + WRITE_TIMEOUT_SECONDS + if not self._stdin_lock.acquire(timeout=WRITE_TIMEOUT_SECONDS): + raise WorkerUnavailable("Napseer MCP worker input timed out") + try: if self.process is not process or process.poll() is not None or process.stdin is None: raise WorkerUnavailable("Napseer MCP worker is not running") try: - process.stdin.write(message if message.endswith(b"\n") else message + b"\n") - process.stdin.flush() - except (BrokenPipeError, OSError) as exc: + remaining = memoryview(message if message.endswith(b"\n") else message + b"\n") + while remaining: + timeout = deadline - time.monotonic() + if timeout <= 0 or self.process is not process: + raise WorkerUnavailable("Napseer MCP worker input timed out") + fd = process.stdin.fileno() + if not select.select([], [fd], [], timeout)[1]: + raise WorkerUnavailable("Napseer MCP worker input timed out") + try: + count = os.write(fd, remaining) + remaining = remaining[count:] + except BlockingIOError: + continue + except (BrokenPipeError, OSError, ValueError) as exc: raise WorkerUnavailable("Napseer MCP worker input closed") from exc + finally: + self._stdin_lock.release() def submit(self, message: bytes) -> PendingRequest | None: self.ensure_current() @@ -326,16 +375,20 @@ class Worker: if process is None: raise WorkerUnavailable("Napseer MCP worker is not running") if is_notification(message): - self._write(process, message) + try: + self._write(process, message) + except WorkerUnavailable: + self.stop() + raise return None key = request_key(request_id(message)) pending = PendingRequest(message=message, key=key, generation=generation) with self._pending_lock: if key in self._pending: - raise WorkerUnavailable("duplicate in-flight JSON-RPC request id") + raise RequestRejected("duplicate in-flight JSON-RPC request id") if len(self._pending) >= MAX_IN_FLIGHT_REQUESTS: - raise WorkerUnavailable("too many in-flight Napseer MCP requests") + raise RequestRejected("too many unfinished Napseer MCP requests") self._pending[key] = pending try: self._write(process, message) @@ -343,26 +396,54 @@ class Worker: with self._pending_lock: if self._pending.get(key) is pending: self._pending.pop(key, None) + # A partial write has an uncertain outcome and corrupts framing. + # Stop this generation rather than appending another request to it. + self.stop() raise return pending - def cancel(self, value) -> bool: + def cancel(self, value, *, expected: PendingRequest | None = None) -> bool: key = request_key(value) with self._pending_lock: pending = self._pending.get(key) + if expected is not None and pending is not expected: + return False if pending is None or request_method(pending.message) == "initialize": return False - self._pending.pop(key, None) + if pending.cancelled: + return True + pending.cancelled = True + pending.cancel_timer = threading.Timer(CANCEL_GRACE_SECONDS, self._expire_cancelled, args=(pending,)) + pending.cancel_timer.daemon = True + timer = pending.cancel_timer self._signal(pending, _CANCELLED) + timer.start() + process = self.process + if process is not None and pending.generation == self.generation: + try: + self._write(process, json.dumps({"jsonrpc": "2.0", "method": "notifications/cancelled", + "params": {"requestId": value}}).encode("utf-8")) + except WorkerUnavailable: + self._expire_cancelled(pending) return True + def _expire_cancelled(self, pending: PendingRequest) -> None: + with self._lifecycle_lock: + with self._pending_lock: + unfinished = self._pending.get(pending.key) is pending + if unfinished and pending.generation == self.generation: + self._stop_locked() + def wait(self, pending: PendingRequest) -> bytes | None: try: result = pending.response.get(timeout=RESPONSE_TIMEOUT_SECONDS) except queue.Empty as exc: - with self._pending_lock: - if self._pending.get(pending.key) is pending: - self._pending.pop(pending.key, None) + if request_method(pending.message) == "initialize": + # Initialize is not cancellable, but a stuck initialization + # must not permanently consume the worker generation. + self._expire_cancelled(pending) + else: + self.cancel(request_id(pending.message), expected=pending) raise WorkerUnavailable("Napseer MCP worker response timed out") from exc if result is _CANCELLED: return None @@ -416,16 +497,12 @@ def supervise() -> int: cancelled, value = cancelled_request_id(message) if cancelled: worker.cancel(value) - try: - worker.submit(message) - except WorkerUnavailable: - pass continue try: pending = worker.submit(message) - except WorkerUnavailable: + except WorkerUnavailable as exc: if not is_notification(message): - emit(worker_error(message)) + emit(worker_error(message, rejected=isinstance(exc, RequestRejected))) continue if pending is not None: waiters.submit(finish, message, pending) diff --git a/tests/test_mcp_memory_ergonomics_helpers.py b/tests/test_mcp_memory_ergonomics_helpers.py index 3538265..0e8a31e 100644 --- a/tests/test_mcp_memory_ergonomics_helpers.py +++ b/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"], diff --git a/tests/test_mcp_memory_node_encryption.py b/tests/test_mcp_memory_node_encryption.py index ef88bdb..9fddeeb 100644 --- a/tests/test_mcp_memory_node_encryption.py +++ b/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", diff --git a/tests/test_mcp_supervisor.py b/tests/test_mcp_supervisor.py index 82f1cd9..28bcad3 100644 --- a/tests/test_mcp_supervisor.py +++ b/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() diff --git a/tests/test_mcp_worker_lifecycle.py b/tests/test_mcp_worker_lifecycle.py index 64a16e6..eb8a2ae 100644 --- a/tests/test_mcp_worker_lifecycle.py +++ b/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)