feat: search pending on load-time indexing (#91)

Author: georgeh0

src/cocoindex_code/cli.py

@@ -165,6 +165,44 @@ def _on_progress(progress: IndexingProgress) -> None:
         raise _typer.Exit(code=1)
 
 
+def _search_with_wait_spinner(
+    client: DaemonClient,
+    project_root: str,
+    query: str,
+    languages: list[str] | None = None,
+    paths: list[str] | None = None,
+    limit: int = 10,
+    offset: int = 0,
+) -> SearchResponse:
+    """Run search, showing a spinner if waiting for load-time indexing."""
+    from rich.console import Console as _Console
+    from rich.live import Live as _Live
+    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:
+
+        def _on_waiting() -> None:
+            nonlocal waiting
+            waiting = True
+            live.update(_Spinner("dots", "Waiting for indexing to complete..."))
+
+        resp = client.search(
+            project_root=project_root,
+            query=query,
+            languages=languages,
+            paths=paths,
+            limit=limit,
+            offset=offset,
+            on_waiting=_on_waiting,
+        )
+
+    return resp
+
+
 _GITIGNORE_COMMENT = "# CocoIndex Code (ccc)"
 _GITIGNORE_ENTRY = "/.cocoindex_code/"
 
@@ -300,14 +338,14 @@ def search(
         if default is not None:
             paths = [default]
 
-    resp = client.search(
+    resp = _search_with_wait_spinner(
+        client,
         project_root=project_root,
         query=query_str,
         languages=lang or None,
         paths=paths,
         limit=limit,
         offset=offset,
-        refresh=False,
     )
     print_search_results(resp)
 

src/cocoindex_code/client.py

@@ -101,20 +101,41 @@ def search(
         paths: list[str] | None = None,
         limit: int = 5,
         offset: int = 0,
-        refresh: bool = False,
+        on_waiting: Callable[[], None] | None = None,
     ) -> SearchResponse:
-        """Search the codebase."""
-        return self._send(  # type: ignore[return-value]
-            SearchRequest(
-                project_root=project_root,
-                query=query,
-                languages=languages,
-                paths=paths,
-                limit=limit,
-                offset=offset,
-                refresh=refresh,
+        """Search the codebase.
+
+        If the daemon sends ``IndexWaitingNotice`` (load-time indexing in
+        progress), calls *on_waiting* (if provided) then continues reading
+        until the final ``SearchResponse``.
+        """
+        self._conn.send_bytes(
+            encode_request(
+                SearchRequest(
+                    project_root=project_root,
+                    query=query,
+                    languages=languages,
+                    paths=paths,
+                    limit=limit,
+                    offset=offset,
+                )
             )
         )
+        while True:
+            try:
+                data = self._conn.recv_bytes()
+            except EOFError:
+                raise RuntimeError("Connection to daemon lost during search")
+            resp = decode_response(data)
+            if isinstance(resp, ErrorResponse):
+                raise RuntimeError(f"Daemon error: {resp.message}")
+            if isinstance(resp, IndexWaitingNotice):
+                if on_waiting is not None:
+                    on_waiting()
+                continue
+            if isinstance(resp, SearchResponse):
+                return resp
+            raise RuntimeError(f"Unexpected response: {type(resp).__name__}")
 
     def project_status(self, project_root: str) -> ProjectStatusResponse:
         return self._send(  # type: ignore[return-value]

src/cocoindex_code/daemon.py

@@ -38,6 +38,7 @@
     SearchRequest,
     SearchResponse,
     SearchResult,
+    SearchStreamResponse,
     StopRequest,
     StopResponse,
     decode_request,
@@ -108,15 +109,17 @@ def __init__(self, embedder: Embedder) -> None:
         self._projects = {}
         self._index_locks = {}
         self._indexing = {}
+        self._load_time_done: dict[str, asyncio.Event] = {}
         self._embedder = embedder
 
     async def get_project(self, project_root: str, *, suppress_auto_index: bool = False) -> Project:
         """Get or create a Project for the given root. Lazy initialization.
 
         When a project is newly loaded and *suppress_auto_index* is False,
-        a background indexing task is fired so the project is indexed
-        immediately.  Callers that will index right away (e.g. IndexRequest,
-        SearchRequest with refresh) should pass ``suppress_auto_index=True``.
+        a background indexing task (load-time indexing) is fired so the project
+        is indexed immediately.  Callers that will index right away (e.g.
+        IndexRequest, SearchRequest with refresh) should pass
+        ``suppress_auto_index=True``.
         """
         if project_root not in self._projects:
             root = Path(project_root)
@@ -126,17 +129,36 @@ async def get_project(self, project_root: str, *, suppress_auto_index: bool = Fa
             self._index_locks[project_root] = asyncio.Lock()
             self._indexing[project_root] = False
 
-            if not suppress_auto_index:
-                asyncio.create_task(self._auto_index(project_root))
+            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))
         return self._projects[project_root]
 
-    async def _auto_index(self, project_root: str) -> None:
-        """Background auto-index, consuming the update_index stream."""
+    def is_load_time_indexing(self, project_root: str) -> bool:
+        """Check if load-time indexing is in progress."""
+        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."""
+        event = self._load_time_done.get(project_root)
+        if event is not None:
+            await event.wait()
+
+    async def _load_time_index(self, project_root: str) -> None:
+        """Background load-time indexing, consuming the update_index stream."""
         try:
             async for _ in self.update_index(project_root):
                 pass
         except Exception:
-            logger.exception("Auto-index failed for %s", project_root)
+            logger.exception("Load-time indexing failed for %s", project_root)
+        finally:
+            event = self._load_time_done.get(project_root)
+            if event is not None:
+                event.set()
 
     async def update_index(
         self, project_root: str, *, suppress_auto_index: bool = True
@@ -251,6 +273,7 @@ def remove_project(self, project_root: str) -> bool:
         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()
             del project
@@ -266,6 +289,7 @@ def close_all(self) -> None:
         self._projects.clear()
         self._index_locks.clear()
         self._indexing.clear()
+        self._load_time_done.clear()
         gc.collect()
 
     def list_projects(self) -> list[DaemonProjectInfo]:
@@ -357,26 +381,55 @@ def _recv() -> bytes:
             pass
 
 
+async def _search_with_wait(
+    registry: ProjectRegistry, req: SearchRequest
+) -> AsyncIterator[SearchStreamResponse]:
+    """Stream search response, waiting for load-time indexing first."""
+    yield IndexWaitingNotice()
+    await registry.wait_for_load_time_indexing(req.project_root)
+    try:
+        results = await registry.search(
+            project_root=req.project_root,
+            query=req.query,
+            languages=req.languages,
+            paths=req.paths,
+            limit=req.limit,
+            offset=req.offset,
+        )
+        yield SearchResponse(
+            success=True,
+            results=results,
+            total_returned=len(results),
+            offset=req.offset,
+        )
+    except Exception as e:
+        yield ErrorResponse(message=str(e))
+
+
 async def _dispatch(
     req: Request,
     registry: ProjectRegistry,
     start_time: float,
     shutdown_event: asyncio.Event,
-) -> Response | AsyncIterator[IndexStreamResponse]:
+) -> Response | AsyncIterator[IndexStreamResponse] | AsyncIterator[SearchStreamResponse]:
     """Dispatch a request to the appropriate handler.
 
     Returns a single Response for most requests, or an AsyncIterator for
-    streaming requests (IndexRequest).
+    streaming requests (IndexRequest, or SearchRequest when waiting for
+    load-time indexing).
     """
     try:
         if isinstance(req, IndexRequest):
             return registry.update_index(req.project_root)
 
         if isinstance(req, SearchRequest):
-            if req.refresh:
-                # Consume the index stream silently for refresh
-                async for _ in registry.update_index(req.project_root):
-                    pass
+            # Ensure the project is loaded (may trigger load-time indexing)
+            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):
+                return _search_with_wait(registry, req)
+
             results = await registry.search(
                 project_root=req.project_root,
                 query=req.query,

src/cocoindex_code/protocol.py

@@ -24,7 +24,6 @@ class SearchRequest(_msgspec.Struct, tag="search"):
     paths: list[str] | None = None
     limit: int = 5
     offset: int = 0
-    refresh: bool = False
 
 
 class ProjectStatusRequest(_msgspec.Struct, tag="project_status"):
@@ -154,6 +153,7 @@ class ErrorResponse(_msgspec.Struct, tag="error"):
 )
 
 IndexStreamResponse = IndexProgressUpdate | IndexWaitingNotice | IndexResponse | ErrorResponse
+SearchStreamResponse = IndexWaitingNotice | SearchResponse | ErrorResponse
 
 # ---------------------------------------------------------------------------
 # Encode / decode helpers (msgpack binary)

src/cocoindex_code/server.py

@@ -127,6 +127,8 @@ async def search(
         """Query the codebase index via the daemon."""
         loop = asyncio.get_event_loop()
         try:
+            if refresh_index:
+                await loop.run_in_executor(None, lambda: client.index(project_root))
             resp = await loop.run_in_executor(
                 None,
                 lambda: client.search(
@@ -136,7 +138,6 @@ async def search(
                     paths=paths,
                     limit=limit,
                     offset=offset,
-                    refresh=refresh_index,
                 ),
             )
             return SearchResultModel(

tests/test_daemon.py

@@ -29,6 +29,7 @@
     RemoveProjectRequest,
     Response,
     SearchRequest,
+    SearchResponse,
     StopRequest,
     decode_response,
     encode_request,
@@ -235,3 +236,44 @@ def test_daemon_remove_project_not_loaded(daemon_sock: str) -> None:
     resp = decode_response(conn.recv_bytes())
     assert resp.ok is True  # type: ignore[union-attr]
     conn.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
+    # trigger load-time indexing in the background.
+    project = Path(tempfile.mkdtemp(prefix="ccc_wait_"))
+    save_project_settings(project, default_project_settings())
+    (project / "main.py").write_text(SAMPLE_MAIN_PY)
+
+    conn, _ = _connect_and_handshake(daemon_sock)
+
+    # Send SearchRequest without prior explicit indexing.
+    # The daemon should trigger load-time indexing, detect it's in progress,
+    # and send IndexWaitingNotice before the final SearchResponse.
+    conn.send_bytes(encode_request(SearchRequest(project_root=str(project), query="fibonacci")))
+
+    got_waiting = False
+    final_resp: SearchResponse | None = None
+    while True:
+        resp = decode_response(conn.recv_bytes())
+        if isinstance(resp, IndexWaitingNotice):
+            got_waiting = True
+            continue
+        if isinstance(resp, SearchResponse):
+            final_resp = resp
+            break
+        raise AssertionError(f"Unexpected response: {type(resp).__name__}")
+
+    assert got_waiting, "Expected IndexWaitingNotice before SearchResponse"
+    assert final_resp is not None
+    assert final_resp.success is True
+    assert len(final_resp.results) > 0
+    assert "main.py" in final_resp.results[0].file_path
+
+    # Second search — load-time indexing is done, no waiting expected
+    conn.send_bytes(encode_request(SearchRequest(project_root=str(project), query="fibonacci")))
+    resp2 = decode_response(conn.recv_bytes())
+    assert isinstance(resp2, SearchResponse)
+    assert resp2.success is True
+    conn.close()

tests/test_protocol.py

@@ -47,7 +47,6 @@ def test_encode_decode_search_request_with_defaults() -> None:
     assert decoded.languages is None
     assert decoded.limit == 5
     assert decoded.offset == 0
-    assert decoded.refresh is False
 
 
 def test_encode_decode_search_request_with_all_fields() -> None:
@@ -58,7 +57,6 @@ def test_encode_decode_search_request_with_all_fields() -> None:
         paths=["src/*"],
         limit=20,
         offset=5,
-        refresh=True,
     )
     data = encode_request(req)
     decoded = decode_request(data)
@@ -69,7 +67,6 @@ def test_encode_decode_search_request_with_all_fields() -> None:
     assert decoded.paths == ["src/*"]
     assert decoded.limit == 20
     assert decoded.offset == 5
-    assert decoded.refresh is True
 
 
 def test_encode_decode_search_response_with_results() -> None: