napseer/nap-cli
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: bound MCP request cancellation and recovery

Browse Source
This commit is contained in:
HumbleTurtle
2026-09-05 02:20:58 -03:00
parent f8594bd3d8
commit e8c6a5109e
7 changed files with 593 additions and 68 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+18
View File
@@ -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
resources/scripts/napseer_mcp_server.py
+100 -35
View File
@@ -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
resources/scripts/napseer_mcp_supervisor.py
+105 -28
View File
@@ -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)
tests/test_mcp_memory_ergonomics_helpers.py
+2 -2
View File
@@ -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
+2 -2
View File
@@ -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
+207 -1
View File
@@ -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
+159
View File
@@ -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)