diff --git a/src/nagini_translation/mcp_server.py b/src/nagini_translation/mcp_server.py index f5b14639..23a1edde 100644 --- a/src/nagini_translation/mcp_server.py +++ b/src/nagini_translation/mcp_server.py @@ -23,7 +23,7 @@ import sys import tempfile -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import Future, ThreadPoolExecutor from typing import List, Optional from mcp.server.fastmcp import FastMCP @@ -33,12 +33,28 @@ mcp = FastMCP('nagini') -_service = None +_service_future: Optional[Future] = None # Multiple verifications can run at once; the service serializes only the fast # translation step internally. _executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix='nagini-verify') +def _get_service(): + """Wait for the service to finish booting. Blocks; call only from + `_executor` threads, never from the event loop (FastMCP runs even sync + tools on the loop).""" + assert _service_future is not None + return _service_future.result() + + +def _service_if_ready(): + """The booted service, or None while it is still booting or failed to boot.""" + if (_service_future is not None and _service_future.done() + and _service_future.exception() is None): + return _service_future.result() + return None + + async def _run(fn): return await asyncio.get_event_loop().run_in_executor(_executor, fn) @@ -78,7 +94,7 @@ class by `ClassName` to verify all its methods. Set `ignore_global` to skip of lines, so only request it when needed. """ selected = {method} if method else None - result = await _run(lambda: _service.verify( + result = await _run(lambda: _get_service().verify( path, selected=selected, counterexample=counterexample, base_dir=base_dir, ignore_global=ignore_global, viper_args=viper_args, include_viper=include_viper, job_token=job_token)) @@ -97,7 +113,7 @@ async def verify_method(path: str, method: str, counterexample: bool = False, (its bare name also matches), or a whole class by `ClassName`. The other parameters are as in `verify_file`. """ - result = await _run(lambda: _service.verify( + result = await _run(lambda: _get_service().verify( path, selected={method}, counterexample=counterexample, viper_args=viper_args, include_viper=include_viper, job_token=job_token)) @@ -120,7 +136,7 @@ async def verify_snippet(code: str, counterexample: bool = False, try: with open(tmp_path, 'w') as f: f.write(code) - result = await _run(lambda: _service.verify( + result = await _run(lambda: _get_service().verify( tmp_path, counterexample=counterexample, base_dir=tmp_dir, ignore_global=ignore_global, viper_args=viper_args, include_viper=include_viper, job_token=job_token)) @@ -130,7 +146,7 @@ async def verify_snippet(code: str, counterexample: bool = False, @mcp.tool() -def configure(options: dict) -> dict: +async def configure(options: dict) -> dict: """Change verification options for subsequent requests; returns the effective configuration. @@ -141,21 +157,25 @@ def configure(options: dict) -> dict: `sif`/`intBitopsSize`/`floatEncoding` reloads the Silver resources; already-running verifications are unaffected. """ - return _service.reconfigure(**options_to_kwargs(options)) + return await _run(lambda: _get_service().reconfigure(**options_to_kwargs(options))) @mcp.tool() def cancel(job_token: Optional[str] = None) -> dict: """Cancel verification: a specific run if `job_token` is given, else all.""" - _service.cancel(job_token=job_token) - return {'cancelled': True, 'jobToken': job_token} + service = _service_if_ready() + if service is not None: + service.cancel(job_token=job_token) + return {'cancelled': service is not None, 'jobToken': job_token} @mcp.tool() def flush_cache() -> dict: """Clear the ViperServer result cache.""" - _service.flush_cache() - return {'flushed': True} + service = _service_if_ready() + if service is not None: + service.flush_cache() + return {'flushed': service is not None} def main(): @@ -164,13 +184,18 @@ def main(): parser.add_argument('--log', default='WARNING') args = parser.parse_args() logging.basicConfig(level=getattr(logging, args.log.upper(), logging.WARNING)) - global _service - _service = make_service(args) + global _service_future + # The JVM-backed service boots in the background (see `main`) so the MCP handshake is answered + # immediately instead of after JVM startup, which can exceed common client startup timeouts. + _service_future = ThreadPoolExecutor( + max_workers=1, thread_name_prefix='nagini-boot').submit(make_service, args) try: mcp.run() finally: try: - _service.shutdown() + service = _service_if_ready() + if service is not None: + service.shutdown() except Exception: logging.exception('Error shutting down service.') # The stdio transport may already have closed these streams by the time diff --git a/tests/servers/test_mcp_server.py b/tests/servers/test_mcp_server.py index 20a9a2e9..38cc284f 100644 --- a/tests/servers/test_mcp_server.py +++ b/tests/servers/test_mcp_server.py @@ -12,6 +12,8 @@ import sys import tempfile +from concurrent.futures import Future + import pytest pytest.importorskip("mcp") @@ -23,6 +25,13 @@ "code", "source", "message"} +def _install_service(service): + """Install a booted service as if the background boot had completed.""" + future = Future() + future.set_result(service) + mcp_server._service_future = future + + def _tool_result_payload(result): """Extract the JSON dict a tool returned from an MCP CallToolResult.""" structured = getattr(result, "structuredContent", None) @@ -92,7 +101,7 @@ def test_stdio_transport_end_to_end_and_clean_shutdown(pass_file): def test_verify_file_reports_structured_failure(service, fail_file): - mcp_server._service = service + _install_service(service) result = asyncio.run(mcp_server.verify_file(fail_file)) assert result["success"] is False assert result["diagnostics"] @@ -103,14 +112,14 @@ def test_verify_file_reports_structured_failure(service, fail_file): def test_verify_file_success(service, pass_file): - mcp_server._service = service + _install_service(service) result = asyncio.run(mcp_server.verify_file(pass_file)) assert result["success"] is True assert result["diagnostics"] == [] def test_verify_method_restricts_to_one_method(service, fail_file): - mcp_server._service = service + _install_service(service) result = asyncio.run(mcp_server.verify_method(fail_file, "add")) assert "success" in result assert "diagnostics" in result @@ -118,7 +127,7 @@ def test_verify_method_restricts_to_one_method(service, fail_file): def test_verify_method_excludes_error_outside_selection(service, mixed_file): """Selecting a method must ignore verification errors in other methods.""" - mcp_server._service = service + _install_service(service) # Baseline: verifying the whole file surfaces the error in `failing`. full = asyncio.run(mcp_server.verify_file(mixed_file)) assert full["success"] is False @@ -134,25 +143,31 @@ def test_verify_method_excludes_error_outside_selection(service, mixed_file): def test_verify_snippet_inline_code(service, pass_src): - mcp_server._service = service + _install_service(service) result = asyncio.run(mcp_server.verify_snippet(pass_src)) assert result["success"] is True def test_verify_snippet_inline_failure(service, fail_src): - mcp_server._service = service + _install_service(service) result = asyncio.run(mcp_server.verify_snippet(fail_src)) assert result["success"] is False assert any("postcondition" in d["code"] for d in result["diagnostics"]) def test_flush_cache_and_cancel_are_callable(service): - mcp_server._service = service + _install_service(service) assert mcp_server.flush_cache() == {"flushed": True} cancelled = mcp_server.cancel() assert cancelled["cancelled"] is True +def test_flush_cache_and_cancel_report_noop_while_booting(): + mcp_server._service_future = Future() # boot still in flight + assert mcp_server.cancel()["cancelled"] is False + assert mcp_server.flush_cache()["flushed"] is False + + def test_options_to_kwargs_maps_and_filters(): from nagini_translation.service import options_to_kwargs kw = options_to_kwargs({"verifier": "carbon", "intBitopsSize": 16, @@ -162,10 +177,12 @@ def test_options_to_kwargs_maps_and_filters(): def test_configure_switches_backend_and_service_stays_usable(service, pass_src): - mcp_server._service = service + _install_service(service) try: - assert mcp_server.configure({"verifier": "carbon"})["verifier"] == "carbon" - assert mcp_server.configure({"verifier": "silicon"})["verifier"] == "silicon" + assert asyncio.run( + mcp_server.configure({"verifier": "carbon"}))["verifier"] == "carbon" + assert asyncio.run( + mcp_server.configure({"verifier": "silicon"}))["verifier"] == "silicon" finally: # Restore the shared session service for other tests. service.reconfigure(verifier_backend="silicon") @@ -175,8 +192,8 @@ def test_configure_switches_backend_and_service_stays_usable(service, pass_src): def test_configure_ignores_unknown_and_null_options(service): - mcp_server._service = service - effective = mcp_server.configure({"unknownKey": 1, "z3Path": None}) + _install_service(service) + effective = asyncio.run(mcp_server.configure({"unknownKey": 1, "z3Path": None})) # Returns the effective configuration without error; nothing changed. assert effective["verifier"] == "silicon" @@ -184,10 +201,10 @@ def test_configure_ignores_unknown_and_null_options(service): def test_configure_reload_option_then_verify(service, pass_file, fail_file): # A reload-triggering option via the configure tool, then verification still # produces correct results. - mcp_server._service = service + _install_service(service) original = service.current_options() try: - effective = mcp_server.configure({"intBitopsSize": 16}) + effective = asyncio.run(mcp_server.configure({"intBitopsSize": 16})) assert effective["intBitopsSize"] == 16 assert asyncio.run(mcp_server.verify_file(pass_file))["success"] is True assert asyncio.run(mcp_server.verify_file(fail_file))["success"] is False @@ -205,7 +222,7 @@ def test_configure_reload_option_then_verify(service, pass_file, fail_file): def test_verify_ignore_global_via_tool(service): - mcp_server._service = service + _install_service(service) # Top-level statements verified by default -> the false assert fails. assert asyncio.run(mcp_server.verify_snippet(_TOPLEVEL_ASSERT_SRC))["success"] is False # ignore_global param skips them -> verifies. @@ -214,7 +231,7 @@ def test_verify_ignore_global_via_tool(service): def test_verify_viper_args_and_include_viper_via_tool(service, pass_src): - mcp_server._service = service + _install_service(service) result = asyncio.run(mcp_server.verify_snippet( pass_src, viper_args=["--timeout=300"], include_viper=True)) assert result["success"] is True @@ -224,15 +241,15 @@ def test_verify_viper_args_and_include_viper_via_tool(service, pass_src): def test_configure_disable_branch_conditions_via_tool(service): - mcp_server._service = service + _install_service(service) original = service.current_options() try: - mcp_server.configure({"disableBranchConditions": False}) + asyncio.run(mcp_server.configure({"disableBranchConditions": False})) on = asyncio.run(mcp_server.verify_snippet(_BRANCH_ERROR_SRC)) assert on["success"] is False assert any(d["branchConditions"] for d in on["diagnostics"]) - mcp_server.configure({"disableBranchConditions": True}) + asyncio.run(mcp_server.configure({"disableBranchConditions": True})) off = asyncio.run(mcp_server.verify_snippet(_BRANCH_ERROR_SRC)) assert off["success"] is False assert all(not d["branchConditions"] for d in off["diagnostics"])