fix: make sure search blocks on load-time indexing gracefully (#92)

Author: georgeh0

src/cocoindex_code/cli.py

@@ -180,15 +180,14 @@ def _search_with_wait_spinner(
     from rich.spinner import Spinner as _Spinner
 
     err_console = _Console(stderr=True)
-    waiting = False
 
-    # Use Live context so the spinner is cleaned up regardless of outcome
-    with _Live(console=err_console, transient=True) as live:
+    with _Live(_Spinner("dots", "Searching..."), console=err_console, transient=True) as live:
 
         def _on_waiting() -> None:
-            nonlocal waiting
-            waiting = True
-            live.update(_Spinner("dots", "Waiting for indexing to complete..."))
+            live.update(
+                _Spinner("dots", "Waiting for indexing to complete..."),
+                refresh=True,
+            )
 
         resp = client.search(
             project_root=project_root,

src/cocoindex_code/daemon.py

@@ -9,7 +9,7 @@
 import sys
 import threading
 import time
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable
 from multiprocessing.connection import Connection, Listener
 from pathlib import Path
 from typing import Any
@@ -102,13 +102,11 @@ class ProjectRegistry:
 
     _projects: dict[str, Project]
     _index_locks: dict[str, asyncio.Lock]
-    _indexing: dict[str, bool]
     _embedder: Embedder
 
     def __init__(self, embedder: Embedder) -> None:
         self._projects = {}
         self._index_locks = {}
-        self._indexing = {}
         self._load_time_done: dict[str, asyncio.Event] = {}
         self._embedder = embedder
 
@@ -127,80 +125,116 @@ async def get_project(self, project_root: str, *, suppress_auto_index: bool = Fa
             project = await Project.create(root, project_settings, self._embedder)
             self._projects[project_root] = project
             self._index_locks[project_root] = asyncio.Lock()
-            self._indexing[project_root] = False
-
-            event = asyncio.Event()
-            self._load_time_done[project_root] = event
-            if suppress_auto_index:
-                event.set()
-            else:
-                asyncio.create_task(self._load_time_index(project_root))
+            self._load_time_done[project_root] = asyncio.Event()
+            if not suppress_auto_index:
+                asyncio.create_task(self._run_index(project_root))
         return self._projects[project_root]
 
-    def is_load_time_indexing(self, project_root: str) -> bool:
-        """Check if load-time indexing is in progress."""
+    def should_wait_for_indexing(self, project_root: str) -> bool:
+        """Check if search should wait before querying.
+
+        Returns True if the index lock is held (indexing actively running)
+        or the initial indexing hasn't completed yet (covers the window
+        between task creation and lock acquisition).
+        """
+        lock = self._index_locks.get(project_root)
+        if lock is not None and lock.locked():
+            return True
         event = self._load_time_done.get(project_root)
         return event is not None and not event.is_set()
 
-    async def wait_for_load_time_indexing(self, project_root: str) -> None:
-        """Wait for load-time indexing to complete. Returns immediately if not in progress."""
+    async def wait_for_indexing_done(self, project_root: str) -> None:
+        """Wait until no indexing is in progress and initial indexing is complete."""
+        # Wait for the initial indexing to complete (if pending)
         event = self._load_time_done.get(project_root)
         if event is not None:
             await event.wait()
+        # Wait for any ongoing indexing to finish (lock released)
+        lock = self._index_locks.get(project_root)
+        if lock is not None and lock.locked():
+            await lock.acquire()
+            lock.release()
 
-    async def _load_time_index(self, project_root: str) -> None:
-        """Background load-time indexing, consuming the update_index stream."""
+    async def _run_index(
+        self,
+        project_root: str,
+        on_progress: Callable[[IndexingProgress], None] | None = None,
+    ) -> None:
+        """Run indexing for a project, acquiring and releasing the per-project lock.
+
+        This is the single place where indexing actually happens.  It is used
+        both as a fire-and-forget background task (load-time indexing) and as a
+        spawned task inside ``update_index`` (client-driven indexing).
+
+        On completion (success or failure) it marks load-time as done
+        (idempotent) and releases the lock.
+        """
+        project = self._projects[project_root]
+        lock = self._index_locks[project_root]
+
+        await lock.acquire()
         try:
-            async for _ in self.update_index(project_root):
-                pass
+            await project.update_index(
+                on_progress=on_progress,
+            )
         except Exception:
-            logger.exception("Load-time indexing failed for %s", project_root)
+            logger.exception("Indexing failed for %s", project_root)
         finally:
             event = self._load_time_done.get(project_root)
             if event is not None:
                 event.set()
+            lock.release()
 
     async def update_index(
         self, project_root: str, *, suppress_auto_index: bool = True
     ) -> AsyncIterator[IndexStreamResponse]:
-        """Update index, yielding progress updates and a final IndexResponse."""
-        project = await self.get_project(project_root, suppress_auto_index=suppress_auto_index)
+        """Update index, yielding progress updates and a final IndexResponse.
+
+        Streams ``IndexProgressUpdate`` messages while indexing is in progress,
+        ending with a terminal ``IndexResponse``.  If the lock is already held,
+        yields ``IndexWaitingNotice`` first.
+
+        The actual indexing runs in a separate task (``_run_index``) so that
+        client disconnects (``GeneratorExit``) do not abort the indexing.
+        """
+        await self.get_project(project_root, suppress_auto_index=suppress_auto_index)
         lock = self._index_locks[project_root]
 
-        # If lock is already held, notify the client and block until released
+        # If lock is already held, notify the client before blocking
         if lock.locked():
             yield IndexWaitingNotice()
 
-        async with lock:
-            self._indexing[project_root] = True
-            try:
-                progress_queue: asyncio.Queue[IndexingProgress] = asyncio.Queue()
-
-                def on_progress(progress: IndexingProgress) -> None:
-                    progress_queue.put_nowait(progress)
-
-                update_task = asyncio.create_task(project.update_index(on_progress=on_progress))
-
-                # Drain the queue until the update completes
-                while not update_task.done():
-                    try:
-                        progress = await asyncio.wait_for(progress_queue.get(), timeout=0.1)
-                        yield IndexProgressUpdate(progress=progress)
-                    except TimeoutError:
-                        continue
-
-                # Drain any remaining items
-                while not progress_queue.empty():
-                    yield IndexProgressUpdate(progress=progress_queue.get_nowait())
-
-                # Propagate any exception from the update task
-                update_task.result()
+        progress_queue: asyncio.Queue[IndexingProgress] = asyncio.Queue()
+        index_task = asyncio.create_task(
+            self._run_index(
+                project_root,
+                on_progress=lambda p: progress_queue.put_nowait(p),
+            )
+        )
 
-                yield IndexResponse(success=True)
-            except Exception as e:
-                yield IndexResponse(success=False, message=str(e))
-            finally:
-                self._indexing[project_root] = False
+        try:
+            # Drain the queue until the task completes
+            while not index_task.done():
+                try:
+                    progress = await asyncio.wait_for(progress_queue.get(), timeout=0.1)
+                    yield IndexProgressUpdate(progress=progress)
+                except TimeoutError:
+                    continue
+
+            # Drain any remaining items
+            while not progress_queue.empty():
+                yield IndexProgressUpdate(progress=progress_queue.get_nowait())
+
+            # Propagate any exception from the index task
+            index_task.result()
+
+            yield IndexResponse(success=True)
+        except GeneratorExit:
+            # Client disconnected — _run_index continues in background and
+            # handles cleanup (release lock, clear _indexing) when done.
+            return
+        except Exception as e:
+            yield IndexResponse(success=False, message=str(e))
 
     async def search(
         self,
@@ -255,7 +289,8 @@ def get_status(self, project_root: str) -> ProjectStatusResponse:
                 " GROUP BY language ORDER BY cnt DESC"
             ).fetchall()
 
-        is_indexing = self._indexing.get(project_root, False)
+        lock = self._index_locks.get(project_root)
+        is_indexing = lock is not None and lock.locked()
         progress = project.indexing_stats if is_indexing else None
         return ProjectStatusResponse(
             indexing=is_indexing,
@@ -272,7 +307,6 @@ def remove_project(self, project_root: str) -> bool:
         was_loaded = project_root in self._projects
         project = self._projects.pop(project_root, None)
         self._index_locks.pop(project_root, None)
-        self._indexing.pop(project_root, None)
         self._load_time_done.pop(project_root, None)
         if project is not None:
             project.close()
@@ -288,7 +322,6 @@ def close_all(self) -> None:
             project.close()
         self._projects.clear()
         self._index_locks.clear()
-        self._indexing.clear()
         self._load_time_done.clear()
         gc.collect()
 
@@ -297,7 +330,7 @@ def list_projects(self) -> list[DaemonProjectInfo]:
         return [
             DaemonProjectInfo(
                 project_root=root,
-                indexing=self._indexing.get(root, False),
+                indexing=self._index_locks[root].locked(),
             )
             for root in self._projects
         ]
@@ -384,9 +417,9 @@ def _recv() -> bytes:
 async def _search_with_wait(
     registry: ProjectRegistry, req: SearchRequest
 ) -> AsyncIterator[SearchStreamResponse]:
-    """Stream search response, waiting for load-time indexing first."""
+    """Stream search response, waiting for ongoing indexing first."""
     yield IndexWaitingNotice()
-    await registry.wait_for_load_time_indexing(req.project_root)
+    await registry.wait_for_indexing_done(req.project_root)
     try:
         results = await registry.search(
             project_root=req.project_root,
@@ -427,7 +460,7 @@ async def _dispatch(
             await registry.get_project(req.project_root)
 
             # If load-time indexing is in progress, return a streaming response
-            if registry.is_load_time_indexing(req.project_root):
+            if registry.should_wait_for_indexing(req.project_root):
                 return _search_with_wait(registry, req)
 
             results = await registry.search(

tests/test_daemon.py

@@ -238,6 +238,53 @@ def test_daemon_remove_project_not_loaded(daemon_sock: str) -> None:
     conn.close()
 
 
+def test_daemon_search_waits_during_explicit_index(daemon_sock: str) -> None:
+    """When IndexRequest is in progress, a concurrent SearchRequest should receive
+    IndexWaitingNotice (Path B: index first, then search)."""
+    # Use enough files to ensure indexing takes long enough for the search to
+    # arrive while it's still in progress.
+    project = Path(tempfile.mkdtemp(prefix="ccc_idx_then_search_"))
+    save_project_settings(project, default_project_settings())
+    for i in range(20):
+        (project / f"module_{i}.py").write_text(
+            f'"""Module {i}."""\n\ndef func_{i}(x: int) -> int:\n'
+            f'    """Compute something for module {i}."""\n'
+            f"    return x * {i} + {i}\n"
+        )
+
+    # Connection 1: start indexing (don't wait for completion)
+    conn1, _ = _connect_and_handshake(daemon_sock)
+    conn1.send_bytes(encode_request(IndexRequest(project_root=str(project))))
+
+    # Send the search request immediately — the daemon processes requests
+    # concurrently across connections, and _run_index needs to acquire the
+    # lock before indexing starts, so a prompt SearchRequest will arrive
+    # while the event is still unset.
+    conn2, _ = _connect_and_handshake(daemon_sock)
+    conn2.send_bytes(encode_request(SearchRequest(project_root=str(project), query="compute")))
+
+    got_waiting = False
+    final_resp: SearchResponse | None = None
+    while True:
+        resp = decode_response(conn2.recv_bytes())
+        if isinstance(resp, IndexWaitingNotice):
+            got_waiting = True
+            continue
+        if isinstance(resp, SearchResponse):
+            final_resp = resp
+            break
+        raise AssertionError(f"Unexpected response on search conn: {type(resp).__name__}")
+
+    assert got_waiting, "Expected IndexWaitingNotice before SearchResponse"
+    assert final_resp is not None
+    assert final_resp.success is True
+
+    # Drain the index stream on connection 1
+    _recv_index_response(conn1)
+    conn1.close()
+    conn2.close()
+
+
 def test_daemon_search_waits_for_load_time_indexing(daemon_sock: str) -> None:
     """Search on a fresh project should wait for load-time indexing, sending IndexWaitingNotice."""
     # Create a new project that the daemon hasn't seen — its first load will