fix(client): serialize daemon startup across processes

src/cocoindex_code/client.py

@@ -13,7 +13,8 @@
 import subprocess
 import sys
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
 from multiprocessing.connection import Client, Connection
 from pathlib import Path
 from typing import NamedTuple
@@ -125,6 +126,49 @@ def _is_daemon_supervised() -> bool:
     return os.environ.get("COCOINDEX_CODE_DAEMON_SUPERVISED") == "1"
 
 
+@contextmanager
+def _daemon_start_lock() -> Iterator[None]:
+    """Serialize daemon startup across independent client processes.
+
+    Each MCP server is a separate process, so the process-local
+    ``_daemon_ensured`` flag cannot prevent a startup stampede. Keep the lock
+    file itself after releasing the OS lock: unlinking it would let a new
+    opener lock a different inode while another process still holds the old
+    one.
+    """
+    daemon_runtime_dir().mkdir(parents=True, exist_ok=True)
+    lock_path = daemon_runtime_dir() / "daemon-start.lock"
+
+    with open(lock_path, "a+b") as lock_file:
+        if sys.platform == "win32":
+            import msvcrt
+
+            lock_file.seek(0)
+            if lock_file.read(1) == b"":
+                lock_file.write(b"\0")
+                lock_file.flush()
+            lock_file.seek(0)
+            while True:
+                try:
+                    msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1)
+                    break
+                except OSError:
+                    time.sleep(0.1)
+        else:
+            import fcntl
+
+            fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
+
+        try:
+            yield
+        finally:
+            if sys.platform == "win32":
+                lock_file.seek(0)
+                msvcrt.locking(lock_file.fileno(), msvcrt.LK_UNLCK, 1)
+            else:
+                fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN)
+
+
 def _connect_and_handshake() -> Connection:
     """Connect to the daemon and perform the version handshake.
 
@@ -142,51 +186,85 @@ def _connect_and_handshake() -> Connection:
     """
     global _daemon_ensured, _ensured_daemon_pid, _consecutive_crash_restarts  # noqa: PLW0603
 
+    startup_error: Exception
     try:
         result = _raw_connect_and_handshake()
         _daemon_ensured = True
         _ensured_daemon_pid = result.resp.pid
         # The daemon answered without needing a restart — any crash streak is over.
         _consecutive_crash_restarts = 0
         return result.conn
-    except DaemonVersionError as e:
+    except DaemonVersionError as error:
+        startup_error = error
         # `resp.ok` is False only for a real version mismatch. Once we have
         # ensured a matching daemon, a fresh version mismatch means the binary
-        # was swapped under us — fail fast. A settings-only restart request
-        # (resp.ok True, but the loaded settings mtime moved) is expected;
-        # restart the daemon below so it reloads them. (stop_daemon sends a
-        # StopRequest, so this restart leaves a graceful-exit marker and the
-        # next connect stays silent.)
-        if _daemon_ensured and not e.resp.ok:
+        # was swapped under us — fail fast.
+        if _daemon_ensured and not error.resp.ok:
             raise
-        stop_daemon()
-    except DaemonProtocolError:
+    except DaemonProtocolError as error:
+        startup_error = error
         # The reply didn't even decode — a daemon whose wire protocol
         # drifted beyond what the version-mismatch path can express (e.g.
-        # still running a pre-upgrade binary). Same treatment as a version
-        # mismatch: restart it below; once a matching daemon was already
-        # ensured this cannot legitimately happen, so fail fast instead of
-        # looping on restarts.
+        # still running a pre-upgrade binary).
         if _daemon_ensured:
             raise
-        stop_daemon()
-    except (ConnectionRefusedError, OSError):
-        # No daemon answered. Normal on the first call of this process (start
-        # one below). If we had already ensured one, decide crash vs graceful
-        # exit — a crash restarts loudly and is capped, a graceful exit (e.g.
-        # a manual stop) restarts silently.
-        if _daemon_ensured:
-            _handle_vanished_daemon()
+    except (ConnectionRefusedError, OSError) as error:
+        startup_error = error
 
     if _is_daemon_supervised():
         # Supervisor is responsible for (re)starting the daemon — just wait
         # for the socket to reappear.
+        _prepare_daemon_restart(startup_error)
         _wait_for_daemon()
-    else:
+        return _connect_after_daemon_start()
+
+    # Multiple MCP servers can arrive here together. The first process starts
+    # the daemon while holding the OS lock. Every waiter rechecks after
+    # acquiring it and joins that daemon instead of starting another one.
+    with _daemon_start_lock():
+        try:
+            result = _raw_connect_and_handshake()
+        except (
+            DaemonVersionError,
+            DaemonProtocolError,
+            ConnectionRefusedError,
+            OSError,
+        ) as locked_error:
+            _prepare_daemon_restart(locked_error)
+        else:
+            _daemon_ensured = True
+            _ensured_daemon_pid = result.resp.pid
+            _consecutive_crash_restarts = 0
+            return result.conn
+
         proc = start_daemon()
         _wait_for_daemon(proc=proc)
+        return _connect_after_daemon_start()
+
+
+def _prepare_daemon_restart(error: Exception) -> None:
+    """Apply restart policy after a failed handshake has been confirmed."""
+    if isinstance(error, DaemonVersionError):
+        if _daemon_ensured and not error.resp.ok:
+            raise error
+        # A settings-only mismatch is expected. Stop the daemon so the new
+        # process reloads global_settings.yml.
+        stop_daemon()
+    elif isinstance(error, DaemonProtocolError):
+        if _daemon_ensured:
+            raise error
+        stop_daemon()
+    elif _daemon_ensured:
+        # No daemon answered. Decide crash vs graceful exit only after the
+        # startup lock recheck, so a transient backlog refusal is not reported
+        # as a crash.
+        _handle_vanished_daemon()
+
+
+def _connect_after_daemon_start() -> Connection:
+    """Wait for a newly started daemon to complete a real handshake."""
+    global _daemon_ensured, _ensured_daemon_pid  # noqa: PLW0603
 
-    # Verify the fresh daemon is reachable
     for _attempt in range(10):
         try:
             result = _raw_connect_and_handshake()

src/cocoindex_code/daemon.py

@@ -693,7 +693,10 @@ def run_daemon(
         except Exception:
             pass
 
-    listener = Listener(sock_path, family=connection_family())
+    # multiprocessing.Listener defaults to backlog=1. Several MCP clients can
+    # connect at once when an agent session starts, so a larger queue prevents
+    # transient refusals from being mistaken for a vanished daemon.
+    listener = Listener(sock_path, family=connection_family(), backlog=64)
     logger.info("Listening on %s", sock_path)
 
     loop = asyncio.new_event_loop()

tests/test_client.py

@@ -2,10 +2,13 @@
 
 from __future__ import annotations
 
+import multiprocessing
+import os
+import subprocess
 import tempfile
 from multiprocessing.connection import Connection
 from pathlib import Path
-from typing import cast
+from typing import Any, cast
 
 import pytest
 
@@ -14,6 +17,86 @@
 from cocoindex_code.protocol import HandshakeResponse
 
 
+def _concurrent_connect_worker(
+    runtime_dir: str,
+    barrier: Any,
+    starts_path: str,
+    ready_path: str,
+    result_queue: Any,
+) -> None:
+    """Run one isolated client process in the daemon-start stampede harness."""
+    os.environ["COCOINDEX_CODE_RUNTIME_DIR"] = runtime_dir
+
+    from cocoindex_code import client as worker_client
+
+    raw_calls = 0
+    sentinel_conn = cast(Connection, object())
+    ok_resp = HandshakeResponse(ok=True, daemon_version="test", pid=42)
+
+    def fake_raw() -> worker_client._HandshakeResult:
+        nonlocal raw_calls
+        raw_calls += 1
+        if raw_calls == 1:
+            barrier.wait(timeout=10)
+            raise ConnectionRefusedError("simultaneous initial miss")
+        if not Path(ready_path).exists():
+            raise ConnectionRefusedError("daemon is not ready")
+        return worker_client._HandshakeResult(conn=sentinel_conn, resp=ok_resp)
+
+    def fake_start() -> subprocess.Popen[bytes]:
+        with Path(starts_path).open("a") as starts_file:
+            starts_file.write(f"{os.getpid()}\n")
+        Path(ready_path).touch()
+        return cast(subprocess.Popen[bytes], object())
+
+    worker_client._raw_connect_and_handshake = fake_raw
+    worker_client.start_daemon = fake_start
+    worker_client._wait_for_daemon = lambda **_kwargs: None
+    worker_client._is_daemon_supervised = lambda: False
+    worker_client._daemon_ensured = False
+
+    try:
+        worker_client._connect_and_handshake()
+    except Exception as exc:
+        result_queue.put(f"{type(exc).__name__}: {exc}")
+    else:
+        result_queue.put(None)
+
+
+def test_concurrent_clients_start_only_one_daemon(tmp_path: Path) -> None:
+    """Independent MCP processes must serialize and deduplicate daemon startup."""
+    process_count = 6
+    ctx = multiprocessing.get_context("spawn")
+    barrier = ctx.Barrier(process_count)
+    result_queue = ctx.Queue()
+    runtime_dir = tmp_path / "runtime"
+    runtime_dir.mkdir()
+    starts_path = tmp_path / "starts"
+    ready_path = tmp_path / "ready"
+
+    processes = [
+        ctx.Process(
+            target=_concurrent_connect_worker,
+            args=(
+                str(runtime_dir),
+                barrier,
+                str(starts_path),
+                str(ready_path),
+                result_queue,
+            ),
+        )
+        for _ in range(process_count)
+    ]
+    for process in processes:
+        process.start()
+    for process in processes:
+        process.join(timeout=15)
+
+    assert all(process.exitcode == 0 for process in processes)
+    assert [result_queue.get(timeout=1) for _ in processes] == [None] * process_count
+    assert len(starts_path.read_text().splitlines()) == 1
+
+
 def test_client_connect_refuses_when_no_daemon(
     tmp_path: Path, monkeypatch: pytest.MonkeyPatch
 ) -> None:
@@ -100,7 +183,7 @@ def test_connect_restarts_ensured_daemon_on_stale_settings(
 
     def fake_raw() -> client._HandshakeResult:
         calls["raw"] += 1
-        if calls["raw"] == 1:
+        if calls["raw"] <= 2:
             raise client.DaemonVersionError(
                 HandshakeResponse(ok=True, daemon_version="v1", pid=42, global_settings_mtime_us=1)
             )
@@ -117,7 +200,7 @@ def fake_raw() -> client._HandshakeResult:
     assert conn is sentinel_conn
     assert calls["stop"] == 1  # old daemon stopped
     assert calls["start"] == 1  # fresh daemon started to reload settings
-    assert calls["raw"] == 2  # reconnected after restart
+    assert calls["raw"] == 3  # lock recheck, then reconnect after restart
 
 
 def test_connect_fails_fast_on_version_mismatch_after_ensured(
@@ -163,7 +246,7 @@ def test_connect_restarts_daemon_on_undecodable_handshake(
 
     def fake_raw() -> client._HandshakeResult:
         calls["raw"] += 1
-        if calls["raw"] == 1:
+        if calls["raw"] <= 2:
             raise client.DaemonProtocolError("Undecodable handshake reply from daemon")
         return client._HandshakeResult(conn=cast(Connection, sentinel_conn), resp=ok_resp)
 
@@ -178,7 +261,7 @@ def fake_raw() -> client._HandshakeResult:
     assert conn is sentinel_conn
     assert calls["stop"] == 1  # incompatible daemon stopped
     assert calls["start"] == 1  # fresh daemon started
-    assert calls["raw"] == 2  # reconnected after restart
+    assert calls["raw"] == 3  # lock recheck, then reconnect after restart
 
 
 def test_connect_fails_fast_on_undecodable_handshake_after_ensured(
@@ -257,7 +340,7 @@ def _setup_vanished_daemon(
 
     def fake_raw() -> client._HandshakeResult:
         calls["raw"] += 1
-        if calls["raw"] == 1:
+        if calls["raw"] <= 2:
             raise ConnectionRefusedError("daemon socket not found")
         return client._HandshakeResult(conn=cast(Connection, sentinel_conn), resp=ok_resp)