530 lines
19 KiB
Python
530 lines
19 KiB
Python
import json
|
|
import os
|
|
import pathlib
|
|
import select
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
|
|
|
|
ROOT = pathlib.Path(__file__).resolve().parents[1]
|
|
SUPERVISOR = ROOT / "resources" / "scripts" / "napseer_mcp_supervisor.py"
|
|
REAL_WORKER = ROOT / "resources" / "scripts" / "napseer_mcp_server.py"
|
|
|
|
|
|
def write_worker(path: pathlib.Path, version: int) -> None:
|
|
path.write_text(
|
|
f"""\
|
|
import json
|
|
import os
|
|
import sys
|
|
import threading
|
|
import time
|
|
|
|
VERSION = {version}
|
|
WRITE_LOCK = threading.Lock()
|
|
CANCELLED = set()
|
|
|
|
|
|
def emit(payload):
|
|
with WRITE_LOCK:
|
|
print(json.dumps(payload), flush=True)
|
|
|
|
|
|
def handle(message):
|
|
method = message.get("method")
|
|
if method == "notifications/initialized":
|
|
return
|
|
if method == "notifications/cancelled":
|
|
CANCELLED.add((message.get("params") or {{}}).get("requestId"))
|
|
return
|
|
if method == "notifications/noisy":
|
|
emit({{"jsonrpc": "2.0", "id": None, "error": {{"code": -32601, "message": "noise"}}}})
|
|
return
|
|
if "id" not in message:
|
|
return
|
|
if method == "tools/call":
|
|
name = (message.get("params") or {{}}).get("name")
|
|
if name == "slow":
|
|
time.sleep(1)
|
|
if name == "crash":
|
|
time.sleep(0.1)
|
|
os._exit(17)
|
|
result = {{"pid": os.getpid(), "version": VERSION, "name": name}}
|
|
else:
|
|
result = {{"version": VERSION}}
|
|
if message.get("id") in CANCELLED:
|
|
return
|
|
emit({{"jsonrpc": "2.0", "id": message.get("id"), "result": result}})
|
|
|
|
for line in sys.stdin:
|
|
message = json.loads(line)
|
|
if "id" not in message:
|
|
handle(message)
|
|
continue
|
|
threading.Thread(target=handle, args=(message,), daemon=True).start()
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
class SupervisorClient:
|
|
def __init__(
|
|
self,
|
|
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)
|
|
if watch_paths:
|
|
environment["NAPSEER_MCP_WORKER_WATCH_PATHS"] = os.pathsep.join(
|
|
str(path) for path in watch_paths
|
|
)
|
|
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)
|
|
self.process = subprocess.Popen(
|
|
[sys.executable, str(SUPERVISOR)],
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
env=environment,
|
|
bufsize=0,
|
|
)
|
|
|
|
def request(self, request_id: int, method: str, params=None):
|
|
self.send(request_id, method, params)
|
|
return self.receive()
|
|
|
|
def send(self, request_id: int, method: str, params=None):
|
|
assert self.process.stdin is not None
|
|
message = {"jsonrpc": "2.0", "id": request_id, "method": method}
|
|
if params is not None:
|
|
message["params"] = params
|
|
self.process.stdin.write(json.dumps(message).encode("utf-8") + b"\n")
|
|
self.process.stdin.flush()
|
|
|
|
def notify(self, method: str, params=None):
|
|
assert self.process.stdin is not None
|
|
message = {"jsonrpc": "2.0", "method": method}
|
|
if params is not None:
|
|
message["params"] = params
|
|
self.process.stdin.write(json.dumps(message).encode("utf-8") + b"\n")
|
|
self.process.stdin.flush()
|
|
|
|
def receive(self, timeout=8):
|
|
assert self.process.stdout is not None
|
|
readable, _, _ = select.select([self.process.stdout], [], [], timeout)
|
|
assert readable, "supervisor did not return a response"
|
|
return json.loads(self.process.stdout.readline())
|
|
|
|
def close(self):
|
|
if self.process.poll() is None:
|
|
assert self.process.stdin is not None
|
|
self.process.stdin.close()
|
|
try:
|
|
self.process.wait(timeout=5)
|
|
except subprocess.TimeoutExpired:
|
|
self.process.terminate()
|
|
self.process.wait(timeout=5)
|
|
|
|
|
|
def test_restarts_worker_killed_between_requests(tmp_path):
|
|
worker_path = tmp_path / "worker.py"
|
|
write_worker(worker_path, 1)
|
|
client = SupervisorClient(worker_path)
|
|
try:
|
|
first = client.request(1, "tools/call", {"name": "pid", "arguments": {}})
|
|
os.kill(first["result"]["pid"], signal.SIGTERM)
|
|
time.sleep(0.1)
|
|
|
|
second = client.request(2, "tools/list")
|
|
|
|
assert second["result"]["version"] == 1
|
|
assert client.process.poll() is None
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
def test_concurrent_requests_are_correlated_without_head_of_line_blocking(tmp_path):
|
|
worker_path = tmp_path / "worker.py"
|
|
write_worker(worker_path, 1)
|
|
client = SupervisorClient(worker_path)
|
|
try:
|
|
client.send(1, "tools/call", {"name": "slow", "arguments": {}})
|
|
time.sleep(0.05)
|
|
client.send(2, "tools/call", {"name": "fast", "arguments": {}})
|
|
|
|
first = client.receive(timeout=0.75)
|
|
second = client.receive(timeout=2)
|
|
|
|
assert first["id"] == 2
|
|
assert first["result"]["name"] == "fast"
|
|
assert second["id"] == 1
|
|
assert second["result"]["name"] == "slow"
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
def test_notification_response_cannot_poison_the_next_request(tmp_path):
|
|
worker_path = tmp_path / "worker.py"
|
|
write_worker(worker_path, 1)
|
|
client = SupervisorClient(worker_path)
|
|
try:
|
|
client.notify("notifications/noisy")
|
|
time.sleep(0.05)
|
|
|
|
response = client.request(7, "tools/list")
|
|
|
|
assert response["id"] == 7
|
|
assert response["result"]["version"] == 1
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
def test_cancellation_drops_the_cancelled_and_late_worker_response(tmp_path):
|
|
worker_path = tmp_path / "worker.py"
|
|
write_worker(worker_path, 1)
|
|
client = SupervisorClient(worker_path)
|
|
try:
|
|
client.send(1, "tools/call", {"name": "slow", "arguments": {}})
|
|
time.sleep(0.05)
|
|
client.notify(
|
|
"notifications/cancelled",
|
|
{"requestId": 1, "reason": "test cancellation"},
|
|
)
|
|
client.send(2, "tools/call", {"name": "fast", "arguments": {}})
|
|
|
|
response = client.receive(timeout=0.75)
|
|
assert response["id"] == 2
|
|
|
|
time.sleep(1.05)
|
|
follow_up = client.request(3, "tools/list")
|
|
assert follow_up["id"] == 3
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
def test_worker_crash_fails_every_in_flight_request_once_then_recovers(tmp_path):
|
|
worker_path = tmp_path / "worker.py"
|
|
write_worker(worker_path, 1)
|
|
client = SupervisorClient(worker_path)
|
|
try:
|
|
client.send(1, "tools/call", {"name": "slow", "arguments": {}})
|
|
client.send(2, "tools/call", {"name": "crash", "arguments": {}})
|
|
|
|
failed = [client.receive(timeout=2), client.receive(timeout=2)]
|
|
|
|
assert {item["id"] for item in failed} == {1, 2}
|
|
assert all(item["error"]["code"] == -32098 for item in failed)
|
|
|
|
recovered = client.request(3, "tools/list")
|
|
assert recovered["id"] == 3
|
|
assert recovered["result"]["version"] == 1
|
|
assert client.process.poll() is None
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
def test_real_supervisor_worker_protocol_lifecycle(tmp_path):
|
|
client = SupervisorClient(REAL_WORKER, project_root=tmp_path)
|
|
try:
|
|
initialized = client.request(
|
|
1,
|
|
"initialize",
|
|
{
|
|
"protocolVersion": "2025-11-25",
|
|
"capabilities": {},
|
|
"clientInfo": {"name": "protocol-test", "version": "1"},
|
|
},
|
|
)
|
|
client.notify("notifications/initialized", {})
|
|
client.send(2, "tools/list")
|
|
client.send(3, "tools/call", {"name": "nap_whoami", "arguments": {}})
|
|
|
|
responses = {item["id"]: item for item in (client.receive(), client.receive())}
|
|
|
|
assert initialized["result"]["protocolVersion"] == "2025-11-25"
|
|
tool_names = {tool["name"] for tool in responses[2]["result"]["tools"]}
|
|
assert "nap_whoami" in tool_names
|
|
assert responses[3]["result"]["isError"] is False
|
|
|
|
client.notify("notifications/unknown", {"proof": "one-way"})
|
|
follow_up = client.request(
|
|
4,
|
|
"tools/call",
|
|
{"name": "nap_whoami", "arguments": {}},
|
|
)
|
|
assert follow_up["id"] == 4
|
|
assert follow_up["result"]["isError"] is False
|
|
finally:
|
|
client.close()
|
|
|
|
assert client.process.returncode == 0
|
|
assert not (tmp_path / ".napseer").exists()
|
|
|
|
|
|
def test_reloads_replaced_worker_between_requests(tmp_path):
|
|
worker_path = tmp_path / "worker.py"
|
|
write_worker(worker_path, 1)
|
|
client = SupervisorClient(worker_path)
|
|
try:
|
|
assert client.request(1, "initialize")["result"]["version"] == 1
|
|
time.sleep(0.01)
|
|
write_worker(worker_path, 2)
|
|
os.utime(worker_path, None)
|
|
|
|
assert client.request(2, "tools/list")["result"]["version"] == 2
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
def test_reloads_stable_launcher_when_watched_runtime_is_replaced(tmp_path):
|
|
runtime_path = tmp_path / "runtime.py"
|
|
launcher_path = tmp_path / "launcher.py"
|
|
write_worker(runtime_path, 1)
|
|
launcher_path.write_text(
|
|
"import runpy\n"
|
|
f"runpy.run_path({str(runtime_path)!r}, run_name='__main__')\n",
|
|
encoding="utf-8",
|
|
)
|
|
client = SupervisorClient(launcher_path, watch_paths=(runtime_path,))
|
|
try:
|
|
first = client.request(1, "initialize")
|
|
launcher_identity = launcher_path.stat().st_mtime_ns
|
|
|
|
time.sleep(0.01)
|
|
write_worker(runtime_path, 2)
|
|
os.utime(runtime_path, None)
|
|
|
|
second = client.request(2, "tools/list")
|
|
|
|
assert first["result"]["version"] == 1
|
|
assert second["result"]["version"] == 2
|
|
assert launcher_path.stat().st_mtime_ns == launcher_identity
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
def test_worker_crash_returns_error_without_closing_transport(tmp_path):
|
|
worker_path = tmp_path / "worker.py"
|
|
write_worker(worker_path, 1)
|
|
client = SupervisorClient(worker_path)
|
|
try:
|
|
failed = client.request(1, "tools/call", {"name": "crash", "arguments": {}})
|
|
recovered = client.request(2, "tools/list")
|
|
|
|
assert failed["error"]["code"] == -32098
|
|
assert recovered["result"]["version"] == 1
|
|
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()
|