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
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)