Fix lifecycle-lock acquire() race that allowed duplicate daemons

src/iai_mcp/lifecycle_lock.py

@@ -23,6 +23,13 @@ def _default_lock_path() -> Path:
 
 SCHEMA_VERSION: int = 1
 
+# acquire() serialises its read-check-write behind an OS advisory lock on a
+# sidecar guard file. How long to wait for that guard before giving up (a
+# healthy acquire() holds it for microseconds; a multi-second wait means a
+# peer is wedged mid-acquire).
+_GUARD_TIMEOUT_SEC: float = 10.0
+_GUARD_POLL_SEC: float = 0.02
+
 
 class LifecycleLockConflict(RuntimeError):
 
@@ -47,6 +54,63 @@ def _current_hostname() -> str:
     return socket.gethostname()
 
 
+def _guard_lock(fd: int) -> None:
+    """Take an exclusive OS advisory lock on the guard fd, cross-process.
+
+    Blocks (bounded by _GUARD_TIMEOUT_SEC) until the lock is granted. Uses
+    BSD flock on POSIX and LockFile via msvcrt on Windows -- both are held by
+    the open file handle and released automatically if the process dies, so a
+    crashed acquirer can never wedge the guard.
+    """
+    import time
+
+    deadline = time.monotonic() + _GUARD_TIMEOUT_SEC
+    if platform.system() == "Windows":
+        import msvcrt
+
+        while True:
+            try:
+                os.lseek(fd, 0, os.SEEK_SET)
+                msvcrt.locking(fd, msvcrt.LK_NBLCK, 1)
+                return
+            except OSError:
+                if time.monotonic() >= deadline:
+                    raise TimeoutError(
+                        "lifecycle guard lock busy: another acquire() appears "
+                        "wedged"
+                    )
+                time.sleep(_GUARD_POLL_SEC)
+    else:
+        import fcntl
+
+        while True:
+            try:
+                fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                return
+            except OSError:
+                if time.monotonic() >= deadline:
+                    raise TimeoutError(
+                        "lifecycle guard lock busy: another acquire() appears "
+                        "wedged"
+                    )
+                time.sleep(_GUARD_POLL_SEC)
+
+
+def _guard_unlock(fd: int) -> None:
+    try:
+        if platform.system() == "Windows":
+            import msvcrt
+
+            os.lseek(fd, 0, os.SEEK_SET)
+            msvcrt.locking(fd, msvcrt.LK_UNLCK, 1)
+        else:
+            import fcntl
+
+            fcntl.flock(fd, fcntl.LOCK_UN)
+    except OSError:
+        log.debug("lifecycle guard unlock failed", exc_info=True)
+
+
 def _is_pid_alive(pid: int) -> bool:
     if pid <= 0:
         return False
@@ -154,6 +218,33 @@ def is_held_by_self(self) -> bool:
         )
 
     def acquire(self) -> None:
+        # Two daemons launched in the same instant (e.g. the Windows scheduled
+        # task and a wrapper/doctor respawn) both used to pass this method: the
+        # read-check-then-write sequence is not atomic, so neither saw a
+        # conflict and both ran, one becoming a store-contending orphan.
+        #
+        # A bare O_EXCL create is not enough on its own: reclaiming a *stale*
+        # dead-pid lock still needs a check-then-remove, and there is no atomic
+        # "remove this file only if it is dead" primitive -- under contention a
+        # racer can evict a lock a peer just created and briefly reopen the
+        # window. So serialise the whole read-check-write behind an OS advisory
+        # lock (flock / LockFileEx) on a sidecar guard file. Only one acquirer
+        # across all processes runs the critical section at a time, which makes
+        # the simple sequence in _acquire_locked() correct; the guard is
+        # auto-released if the holder dies, so it cannot wedge.
+        self._lock_path.parent.mkdir(parents=True, exist_ok=True)
+        guard_path = self._lock_path.parent / (self._lock_path.name + ".guard")
+        gfd = os.open(str(guard_path), os.O_CREAT | os.O_RDWR, 0o600)
+        try:
+            _guard_lock(gfd)
+            try:
+                self._acquire_locked()
+            finally:
+                _guard_unlock(gfd)
+        finally:
+            os.close(gfd)
+
+    def _acquire_locked(self) -> None:
         existing = self.read()
         if existing is not None:
             if existing["hostname"] == _current_hostname() and _is_pid_alive(
@@ -173,7 +264,6 @@ def acquire(self) -> None:
             "schema_version": SCHEMA_VERSION,
         }
 
-        self._lock_path.parent.mkdir(parents=True, exist_ok=True)
         fd, tmp = tempfile.mkstemp(
             prefix=".locked.",
             suffix=".tmp",

tests/test_lifecycle_lock.py

@@ -219,6 +219,127 @@ def test_read_returns_none_for_invalid_schema(tmp_path: Path) -> None:
     assert lock.read() is None
 
 
+def test_acquire_serialises_and_guard_is_released(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # The guard must serialise the critical section and then be released so the
+    # next acquirer can proceed. Acquire, release, and re-acquire in sequence
+    # from the same process: if the guard leaked, the second acquire would hang
+    # until the timeout and raise TimeoutError instead of succeeding.
+    import iai_mcp.lifecycle_lock as ll
+
+    monkeypatch.setattr(ll, "_current_hostname", lambda: "host.local")
+    # No live prior holder, so each acquire should cleanly take the lock.
+    monkeypatch.setattr(ll, "_is_pid_alive", lambda pid: pid == os.getpid())
+
+    lock_path = tmp_path / ".locked"
+    lock = LifecycleLock(lock_path)
+
+    lock.acquire()
+    assert lock.is_held_by_self()
+    lock.release()
+    # Guard must be free now; this must not block/raise TimeoutError.
+    lock.acquire()
+    assert lock.is_held_by_self()
+
+
+def test_acquire_concedes_to_existing_live_lock(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # A live local lock is already present; a second acquirer must raise inside
+    # the guarded critical section and leave the holder's pid untouched.
+    lock_path = tmp_path / ".locked"
+    lock_path.write_text(
+        json.dumps(
+            {
+                "pid": 31_337,
+                "hostname": "host.local",
+                "started_at": "2026-05-01T00:00:00+00:00",
+                "schema_version": SCHEMA_VERSION,
+            }
+        ),
+        encoding="utf-8",
+    )
+    import iai_mcp.lifecycle_lock as ll
+
+    monkeypatch.setattr(ll, "_current_hostname", lambda: "host.local")
+    monkeypatch.setattr(ll, "_is_pid_alive", lambda pid: True)
+
+    lock = LifecycleLock(lock_path)
+    with pytest.raises(LifecycleLockConflict) as exc_info:
+        lock.acquire()
+
+    assert exc_info.value.existing is not None
+    assert exc_info.value.existing["pid"] == 31_337
+    payload = json.loads(lock_path.read_text(encoding="utf-8"))
+    assert payload["pid"] == 31_337
+
+
+def test_concurrent_acquire_yields_exactly_one_winner(
+    tmp_path: Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    # End-to-end race: N threads acquire() the same lock path simultaneously,
+    # starting from a stale dead-pid lock (the real Windows dual-launch boot
+    # scenario). Exactly one must win; the rest must raise. Pre-fix this failed
+    # because read()-check-then-os.replace() is last-writer-wins, so multiple
+    # threads "won" and real daemons ended up doubled.
+    import threading
+
+    import iai_mcp.lifecycle_lock as ll
+
+    real_pid = os.getpid()
+    monkeypatch.setattr(ll, "_current_hostname", lambda: "host.local")
+    # Only this live process counts as a daemon; the seeded stale pid is dead.
+    monkeypatch.setattr(ll, "_is_pid_alive", lambda pid: pid == real_pid)
+
+    lock_path = tmp_path / ".locked"
+    lock_path.write_text(
+        json.dumps(
+            {
+                "pid": 999_999,
+                "hostname": "host.local",
+                "started_at": "2026-04-30T15:00:00+00:00",
+                "schema_version": SCHEMA_VERSION,
+            }
+        ),
+        encoding="utf-8",
+    )
+
+    n_threads = 12
+    winners: list[int] = []
+    conflicts: list[int] = []
+    guard = threading.Lock()
+    gate = threading.Barrier(n_threads)
+
+    def worker(worker_id: int) -> None:
+        gate.wait()
+        try:
+            LifecycleLock(lock_path).acquire()
+            with guard:
+                winners.append(worker_id)
+        except LifecycleLockConflict:
+            with guard:
+                conflicts.append(worker_id)
+
+    threads = [
+        threading.Thread(target=worker, args=(i,)) for i in range(n_threads)
+    ]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    assert len(winners) == 1, (
+        f"expected exactly one winner, got {len(winners)}: {winners}"
+    )
+    assert len(conflicts) == n_threads - 1
+    final = json.loads(lock_path.read_text(encoding="utf-8"))
+    assert final["pid"] == real_pid
+
+
 def test_acquire_writes_mode_0600(tmp_path: Path) -> None:
     lock_path = tmp_path / ".locked"
     lock = LifecycleLock(lock_path)