diff --git a/zenith/src/zenith_harness/server.py b/zenith/src/zenith_harness/server.py index 19c37eb..f0f3c94 100644 --- a/zenith/src/zenith_harness/server.py +++ b/zenith/src/zenith_harness/server.py @@ -7,9 +7,10 @@ import argparse import asyncio +import fcntl import logging import os -from typing import Annotated, Any +from typing import Annotated, Any, Callable from fastmcp import Context, FastMCP from pydantic import Field @@ -99,6 +100,49 @@ def create_terminal_reviewer_server() -> FastMCP: # --------------------------------------------------------------------------- +def _run_project_locked( + controller: ProjectController, + project_id: str, + fn: Callable[..., Any], + *args: Any, +) -> Any: + """Run a mutating controller call while holding an OS advisory lock whose + lifetime is tied to THIS WORKER THREAD, not to the async MCP request. + + The tools hop into a thread via ``asyncio.to_thread``; a running thread's + future cannot be cancelled, so if the MCP request is aborted mid-wave (user + interrupt or client idle-abort) the thread keeps running to completion. The + in-process ``asyncio.Lock`` releases on that cancellation, which previously + let a retried call run concurrently and race on disk state + (task-state.json / attempts/*.json), stubbing still-in-flight validators and + failing gates on empty verdicts. + + An advisory ``flock`` taken here — inside the thread, released only when the + file handle closes as this function returns — spans the true wave lifetime. + A concurrent mutating call gets ``wave_in_progress`` instead of racing. + """ + runtime = controller.config.zenith_runtime_dir(project_id) + try: + runtime.mkdir(parents=True, exist_ok=True) + except OSError: + # Project may not exist yet or path is unwritable; let the real call + # raise its own not_found/validation ToolError rather than masking it. + return fn(*args) + lock_path = runtime / ".wave.lock" + with open(lock_path, "w") as handle: + try: + fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) + except BlockingIOError as exc: + raise ToolError( + "wave_in_progress", + "another mutating call is already running for this project; poll " + ".zenith-runtime/missions//task-state.json and retry when the " + "wave is terminal (do not interrupt or re-issue a running " + "advance_project)", + ) from exc + return fn(*args) + + def _register_orchestrator_tools(mcp: FastMCP, controller: ProjectController) -> None: # Per-project lock around mutating controller calls. The thread hop in each # tool prevents event-loop blocking, but two same-project tool calls could @@ -166,7 +210,14 @@ async def submit_plan( async with await _project_lock(project_id): try: return _to_payload( - await asyncio.to_thread(controller.submit_plan, project_id, task_list) + await asyncio.to_thread( + _run_project_locked, + controller, + project_id, + controller.submit_plan, + project_id, + task_list, + ) ) except ToolError as exc: return _to_payload(exc) @@ -193,7 +244,14 @@ async def advance_project( async with await _project_lock(project_id): try: return _to_payload( - await asyncio.to_thread(controller.advance_project, project_id, max_steps) + await asyncio.to_thread( + _run_project_locked, + controller, + project_id, + controller.advance_project, + project_id, + max_steps, + ) ) except ToolError as exc: return _to_payload(exc) @@ -215,7 +273,13 @@ async def end_mission( async with await _project_lock(project_id): try: return _to_payload( - await asyncio.to_thread(controller.end_mission, project_id) + await asyncio.to_thread( + _run_project_locked, + controller, + project_id, + controller.end_mission, + project_id, + ) ) except ToolError as exc: return _to_payload(exc) @@ -240,7 +304,14 @@ async def decide_attention( async with await _project_lock(project_id): try: return _to_payload( - await asyncio.to_thread(controller.decide_attention, project_id, decisions) + await asyncio.to_thread( + _run_project_locked, + controller, + project_id, + controller.decide_attention, + project_id, + decisions, + ) ) except ToolError as exc: return _to_payload(exc) @@ -276,7 +347,14 @@ async def abort_project( async with await _project_lock(project_id): try: return _to_payload( - await asyncio.to_thread(controller.abort_project, project_id, reason) + await asyncio.to_thread( + _run_project_locked, + controller, + project_id, + controller.abort_project, + project_id, + reason, + ) ) except ToolError as exc: return _to_payload(exc) diff --git a/zenith/tests/test_server_wave_lock.py b/zenith/tests/test_server_wave_lock.py new file mode 100644 index 0000000..7af2999 --- /dev/null +++ b/zenith/tests/test_server_wave_lock.py @@ -0,0 +1,114 @@ +"""Regression tests for the per-project wave flock (`_run_project_locked`). + +Proves the fix for the concurrent-wave disk race: an advisory lock whose +lifetime is tied to the worker THREAD (not the async request) so that an +aborted `advance_project` cannot let a retried call run concurrently and stub +still-in-flight validators. See server.py `_run_project_locked`. +""" +from __future__ import annotations + +import asyncio +import threading +import time +import types +from pathlib import Path + +import pytest + +from zenith_harness.config import HarnessConfig +from zenith_harness.controller import ToolError +from zenith_harness.server import _run_project_locked + +PID = "proj-1" + + +@pytest.fixture +def controller(harness_home: Path): + bundled = Path(__file__).resolve().parents[1] / "src" / "zenith_harness" / "bundled" + cfg = HarnessConfig( + bundled_dir=bundled, + harness_home=harness_home, + projects_dir=harness_home / "projects", + orchestrator_provider_name="claude", + worker_provider_name="claude", + worker_acp_command=None, + validator_provider_name=None, + validator_acp_command=None, + terminal_reviewer_provider_name=None, + terminal_reviewer_acp_command=None, + ) + # The helper only touches controller.config.zenith_runtime_dir(pid). + return types.SimpleNamespace(config=cfg) + + +def test_runs_fn_and_returns_result(controller): + assert _run_project_locked(controller, PID, lambda a, b: a + b, 2, 3) == 5 + assert (controller.config.zenith_runtime_dir(PID) / ".wave.lock").exists() + + +def test_second_call_is_rejected_while_first_holds_lock(controller): + started, release = threading.Event(), threading.Event() + + def slow(): + started.set() + release.wait(5) + return "first" + + t = threading.Thread(target=lambda: _run_project_locked(controller, PID, slow)) + t.start() + assert started.wait(5) + try: + with pytest.raises(ToolError) as ei: + _run_project_locked(controller, PID, lambda: "second") + assert ei.value.code == "wave_in_progress" + finally: + release.set() + t.join(5) + # lock freed after the holder returns + assert _run_project_locked(controller, PID, lambda: "third") == "third" + + +def test_lock_survives_async_cancellation(controller): + """The exact bug: aborting the MCP request must NOT free the wave lock while + the orphaned worker thread keeps running.""" + + async def scenario(): + started, release = threading.Event(), threading.Event() + + def slow(): + started.set() + release.wait(5) + return "done" + + task = asyncio.create_task( + asyncio.to_thread(_run_project_locked, controller, PID, slow) + ) + await asyncio.to_thread(started.wait, 5) # thread now in critical section + + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + # async request is gone, but the thread still runs and still holds the lock + with pytest.raises(ToolError) as ei: + _run_project_locked(controller, PID, lambda: "x") + assert ei.value.code == "wave_in_progress" + + # let the orphaned wave finish; the lock frees only then + release.set() + for _ in range(50): + await asyncio.to_thread(time.sleep, 0.05) + try: + assert _run_project_locked(controller, PID, lambda: "ok") == "ok" + return + except ToolError: + continue + pytest.fail("lock never released after the wave completed") + + asyncio.run(scenario()) + + +if __name__ == "__main__": + import sys + + sys.exit(pytest.main([__file__, "-v"]))