diff --git a/server.py b/server.py index 3f53357..1c81b06 100644 --- a/server.py +++ b/server.py @@ -23,6 +23,7 @@ import time from datetime import datetime, timezone +import anyio import uvicorn from mcp.server.fastmcp import FastMCP from mcp.server.transport_security import TransportSecuritySettings @@ -72,6 +73,13 @@ class LoggingRegistrar: the tool layer having to know that logging exists. functools.wraps keeps the signature and type hints intact, which is what FastMCP builds its JSON schema from. + + The wrapper is also what keeps a slow call from stalling the server. + FastMCP runs a plain function straight on the event loop, so one call + waiting on AutoCAD (minutes, behind a dialog) would hold every other + request, /health included. The tool runs on a worker thread instead and + the loop stays free; the COM bridge still puts the AutoCAD work itself on + its one thread, one call at a time. """ def __init__(self, mcp: FastMCP) -> None: @@ -83,10 +91,12 @@ def tool(self, *args, **kwargs): def register(fn): @functools.wraps(fn) - def logged(*call_args, **call_kwargs): + async def logged(*call_args, **call_kwargs): started = time.monotonic() try: - result = fn(*call_args, **call_kwargs) + result = await anyio.to_thread.run_sync( + functools.partial(fn, *call_args, **call_kwargs) + ) except Exception as exc: # noqa: BLE001 - re-raised below _record( fn.__name__, call_kwargs, diff --git a/tests/test_assembled_server.py b/tests/test_assembled_server.py index 9f74c2b..2377a28 100644 --- a/tests/test_assembled_server.py +++ b/tests/test_assembled_server.py @@ -13,6 +13,7 @@ import os import tempfile +import threading import unittest from pathlib import Path from unittest import mock @@ -23,6 +24,7 @@ from starlette.testclient import TestClient # noqa: E402 +import acad # noqa: E402 import config # noqa: E402 import server # noqa: E402 @@ -128,6 +130,47 @@ def test_bearer_token_is_checked_in_front_of_the_transport(self): self.assertEqual(refused.status_code, 401) self.assertEqual(accepted.status_code, 200, accepted.text) + def test_a_stuck_tool_call_does_not_stall_the_server(self): + """A call waiting on AutoCAD (behind a dialog, say) can sit there for + minutes. It must not hold the event loop: /health has to answer + meanwhile, and so do other clients.""" + entered, release = threading.Event(), threading.Event() + + def stuck_run(_work, _timeout): + entered.set() + release.wait(30) + return "ABC" + + call = rpc("tools/call", { + "name": "get_entity_info", "arguments": {"handle": "ABC"}}, id_=3) + replies = {} + caller = threading.Thread( + target=lambda: replies.setdefault("tool", self.post(call)), + daemon=True) + prober = threading.Thread( + target=lambda: replies.setdefault("health", self.client.get( + f"/{SECRET}/health", headers=TUNNEL)), + daemon=True) + + with mock.patch.object(acad.bridge, "run", side_effect=stuck_run): + caller.start() + try: + self.assertTrue(entered.wait(10), "the tool call never started") + prober.start() + prober.join(5) + self.assertFalse( + prober.is_alive(), + "/health did not answer while a tool call was waiting") + finally: + release.set() + caller.join(10) + if prober.ident: + prober.join(10) + + self.assertEqual(replies["health"].status_code, 200) + self.assertEqual(replies["tool"].status_code, 200, replies["tool"].text) + self.assertFalse(replies["tool"].json()["result"].get("isError")) + if __name__ == "__main__": unittest.main()