fix: make sure daemon restarted on version mismatch (#107)

Author: georgeh0

src/cocoindex_code/cli.py

@@ -48,22 +48,6 @@ def require_project_root() -> Path:
     return root
 
 
-def require_daemon_for_project() -> str:
-    """Resolve project root, then ensure daemon is running (auto-starting if needed).
-
-    Returns ``project_root_str``. Exits on failure.
-    """
-    from .client import ensure_daemon
-
-    project_root = require_project_root()
-    try:
-        ensure_daemon()
-    except Exception as e:
-        _typer.echo(f"Error: Failed to connect to daemon: {e}", err=True)
-        raise _typer.Exit(code=1)
-    return str(project_root)
-
-
 def resolve_default_path(project_root: Path) -> str | None:
     """Compute default ``--path`` filter from CWD relative to project root."""
     cwd = Path.cwd().resolve()
@@ -306,7 +290,7 @@ def index() -> None:
     """Create/update index for the codebase."""
     from . import client as _client
 
-    project_root = require_daemon_for_project()
+    project_root = str(require_project_root())
     print_project_header(project_root)
     _run_index_with_progress(project_root)
     print_index_stats(_client.project_status(project_root))
@@ -322,7 +306,7 @@ def search(
     refresh: bool = _typer.Option(False, "--refresh", help="Refresh index before searching"),
 ) -> None:
     """Semantic search across the codebase."""
-    project_root = require_daemon_for_project()
+    project_root = str(require_project_root())
     query_str = " ".join(query)
 
     if refresh:
@@ -353,7 +337,7 @@ def status() -> None:
     """Show project status."""
     from . import client as _client
 
-    project_root = require_daemon_for_project()
+    project_root = str(require_project_root())
     print_project_header(project_root)
     print_index_stats(_client.project_status(project_root))
 
@@ -436,7 +420,7 @@ def mcp() -> None:
     """Run as MCP server (stdio mode)."""
     import asyncio
 
-    project_root = require_daemon_for_project()
+    project_root = str(require_project_root())
 
     async def _run_mcp() -> None:
         from .server import create_mcp_server
@@ -469,12 +453,6 @@ def daemon_status() -> None:
     """Show daemon status."""
     from . import client as _client
 
-    try:
-        _client.ensure_daemon()
-    except Exception as e:
-        _typer.echo(f"Error: {e}", err=True)
-        raise _typer.Exit(code=1)
-
     resp = _client.daemon_status()
     _typer.echo(f"Daemon version: {resp.version}")
     _typer.echo(f"Uptime: {resp.uptime_seconds:.1f}s")

src/cocoindex_code/client.py

@@ -51,13 +51,51 @@
 # ---------------------------------------------------------------------------
 
 
+_daemon_ensured = False
+
+
 def _connect_and_handshake() -> Connection:
     """Connect to the daemon and perform the version handshake.
 
     Returns the open connection for the caller to send exactly one request.
-    Raises ``ConnectionRefusedError`` if the daemon is not running, or
-    ``RuntimeError`` on protocol/version errors.
+
+    On the first call, automatically starts or
+    restarts the daemon if needed.  Subsequent calls fail fast with
+    ``DaemonVersionError`` on mismatch (indicating the daemon was replaced
+    mid-session, e.g. after a tool upgrade).
     """
+    global _daemon_ensured  # noqa: PLW0603
+
+    if _daemon_ensured:
+        return _raw_connect_and_handshake()
+
+    # First connection — auto-start/restart as needed.
+    try:
+        conn = _raw_connect_and_handshake()
+        _daemon_ensured = True
+        return conn
+    except DaemonVersionError:
+        stop_daemon()
+    except (ConnectionRefusedError, OSError):
+        pass
+
+    start_daemon()
+    _wait_for_daemon()
+
+    # Verify the fresh daemon is reachable
+    for _attempt in range(10):
+        try:
+            conn = _raw_connect_and_handshake()
+            _daemon_ensured = True
+            return conn
+        except (ConnectionRefusedError, OSError):
+            time.sleep(0.5)
+
+    raise RuntimeError("Failed to connect to daemon after starting it")
+
+
+def _raw_connect_and_handshake() -> Connection:
+    """Low-level connect + handshake without auto-start logic."""
     sock = daemon_socket_path()
     if sys.platform != "win32" and not os.path.exists(sock):
         raise ConnectionRefusedError(f"Daemon socket not found: {sock}")
@@ -80,20 +118,25 @@ def _connect_and_handshake() -> Connection:
     if not isinstance(resp, HandshakeResponse):
         conn.close()
         raise RuntimeError(f"Unexpected handshake response: {type(resp).__name__}")
-    if not resp.ok:
+    if not resp.ok or _needs_restart(resp):
         conn.close()
-        raise _VersionMismatchError(resp)
+        raise DaemonVersionError(resp)
     return conn
 
 
-class _VersionMismatchError(Exception):
-    """Raised when the daemon version or settings are stale."""
+class DaemonVersionError(RuntimeError):
+    """Raised when the daemon has a version or settings mismatch.
+
+    The first ``_connect_and_handshake()`` call handles this by restarting
+    the daemon.  If a mismatch occurs on a subsequent call, it means the
+    daemon was replaced mid-session (e.g. after a tool upgrade).
+    """
 
     def __init__(self, resp: HandshakeResponse) -> None:
         self.resp = resp
         super().__init__(
-            f"Daemon version mismatch: {resp.daemon_version} "
-            f"(settings_mtime={resp.global_settings_mtime_us})"
+            f"Daemon version mismatch (daemon={resp.daemon_version}, "
+            f"client={__version__}). Please retry — the daemon may need a restart."
         )
 
 
@@ -319,6 +362,8 @@ def stop_daemon() -> None:
 
     Escalation: StopRequest → SIGTERM → SIGKILL.
     """
+    global _daemon_ensured  # noqa: PLW0603
+    _daemon_ensured = False
     pid_path = daemon_pid_path()
 
     pid: int | None = None
@@ -329,10 +374,15 @@ def stop_daemon() -> None:
     except (FileNotFoundError, ValueError):
         pass
 
-    # 1) Graceful StopRequest via socket
+    # 1) Graceful StopRequest via socket (bypass auto-start)
     try:
-        stop()
-    except (ConnectionRefusedError, OSError, RuntimeError):
+        conn = _raw_connect_and_handshake()
+        try:
+            conn.send_bytes(encode_request(StopRequest()))
+            conn.recv_bytes()
+        finally:
+            conn.close()
+    except (ConnectionRefusedError, OSError, RuntimeError, DaemonVersionError):
         pass
 
     if _wait_for_daemon_exit(timeout=3.0):
@@ -408,37 +458,3 @@ def _needs_restart(resp: HandshakeResponse) -> bool:
     if current_mtime != resp.global_settings_mtime_us:
         return True
     return False
-
-
-def ensure_daemon() -> None:
-    """Ensure daemon is running with correct version. Starts or restarts as needed.
-
-    After this returns, per-request functions (``index``, ``search``, etc.)
-    can be called directly — each opens its own connection.
-    """
-    # Try connecting to existing daemon
-    try:
-        conn = _connect_and_handshake()
-        conn.close()
-        return  # daemon is up and version matches
-    except _VersionMismatchError:
-        stop_daemon()
-    except (ConnectionRefusedError, OSError):
-        pass
-
-    # Start daemon
-    start_daemon()
-    _wait_for_daemon()
-
-    # Verify with retries
-    for _attempt in range(10):
-        try:
-            conn = _connect_and_handshake()
-            conn.close()
-            return
-        except _VersionMismatchError as e:
-            raise RuntimeError(f"Daemon mismatch after fresh start: {e}") from e
-        except (ConnectionRefusedError, OSError):
-            time.sleep(0.5)
-
-    raise RuntimeError("Failed to connect to daemon after starting it")

src/cocoindex_code/server.py

@@ -274,8 +274,6 @@ def main() -> None:
     from . import client as _client
     from .protocol import IndexingProgress
 
-    _client.ensure_daemon()
-
     if args.command == "index":
         import sys
 

tests/test_client.py

@@ -18,4 +18,4 @@ def test_client_connect_refuses_when_no_daemon(
     monkeypatch.setattr("cocoindex_code.client.daemon_socket_path", lambda: sock_path)
 
     with pytest.raises(ConnectionRefusedError):
-        client._connect_and_handshake()
+        client._raw_connect_and_handshake()