diff --git a/CLAUDE.md b/CLAUDE.md index e902766..be1572f 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -31,6 +31,8 @@ The tool runs as a **long-lived background daemon** (`daemon.py`) that accepts c The daemon is auto-started by `client.py` when not running. It is restarted automatically when the `global_settings.yml` mtime changes (version bump or settings edit), detected via the `HandshakeResponse.global_settings_mtime_us` field. +The daemon idle-exits after `daemon.idle_timeout_minutes` (default 180, 0 = never) without client activity, unless a live MCP session is sending heartbeats; clients transparently restart it on the next request. Graceful exits leave a `last_exit` marker file so the client can tell an expected exit from a crash (which triggers a loud restart, capped after consecutive crashes). + ### Key Layers **`protocol.py`** — All IPC message types as `msgspec.Struct` tagged unions, serialized as msgpack. Every request type, response type, and streaming wrapper lives here. The `Request` and `Response` type aliases are the union types used by the decoder. diff --git a/README.md b/README.md index c0fbf14..4b7c5de 100644 --- a/README.md +++ b/README.md @@ -515,10 +515,15 @@ embedding: envs: # extra environment variables for the daemon OPENAI_API_KEY: your-key # only needed if not already in your shell environment + +daemon: + idle_timeout_minutes: 180 # optional: exit the daemon after this long without client activity (default 180, 0 = never) ``` > **Note:** The daemon inherits your shell environment. If an API key (e.g. `OPENAI_API_KEY`) is already set as an environment variable, you don't need to duplicate it in `envs`. The `envs` field is only for values that aren't in your environment. +> **Idle timeout:** the background daemon holds the embedding model in RAM, so it exits after `daemon.idle_timeout_minutes` without client activity and is restarted automatically on your next `ccc` command or MCP search. A live MCP session sends periodic heartbeats, so the daemon never idles out while your coding agent is connected. Set `0` to keep the daemon running forever. + > **Custom location:** set `COCOINDEX_CODE_DIR` to place `global_settings.yml` somewhere other than `~/.cocoindex_code/` — useful if you want the file to live alongside your projects (e.g. on a synced folder). #### `indexing_params` / `query_params` diff --git a/src/cocoindex_code/cli.py b/src/cocoindex_code/cli.py index 0b0b38c..0621948 100644 --- a/src/cocoindex_code/cli.py +++ b/src/cocoindex_code/cli.py @@ -1003,11 +1003,18 @@ def mcp() -> None: project_root = str(require_project_root()) async def _run_mcp() -> None: - from .server import create_mcp_server + from .server import create_mcp_server, run_heartbeat_loop mcp_server = create_mcp_server(project_root) - asyncio.create_task(_bg_index(project_root)) - await mcp_server.run_stdio_async() + background_tasks = { + asyncio.create_task(_bg_index(project_root)), + asyncio.create_task(run_heartbeat_loop()), + } + try: + await mcp_server.run_stdio_async() + finally: + for task in background_tasks: + task.cancel() asyncio.run(_run_mcp()) @@ -1037,6 +1044,8 @@ def daemon_status() -> None: resp = _client.daemon_status() _typer.echo(f"Daemon version: {resp.version}") _typer.echo(f"Uptime: {resp.uptime_seconds:.1f}s") + timeout_desc = f"{resp.idle_timeout_minutes}m" if resp.idle_timeout_minutes > 0 else "disabled" + _typer.echo(f"Idle: {resp.idle_seconds:.1f}s (timeout: {timeout_desc})") if resp.projects: _typer.echo("Projects:") for p in resp.projects: diff --git a/src/cocoindex_code/client.py b/src/cocoindex_code/client.py index 25f3c28..272eea8 100644 --- a/src/cocoindex_code/client.py +++ b/src/cocoindex_code/client.py @@ -37,6 +37,8 @@ ErrorResponse, HandshakeRequest, HandshakeResponse, + HeartbeatRequest, + HeartbeatResponse, IndexingProgress, IndexProgressUpdate, IndexRequest, @@ -414,6 +416,30 @@ def stop() -> StopResponse: return _send(StopRequest()) # type: ignore[return-value] +def send_heartbeat() -> bool: + """Send one heartbeat so a running daemon counts an active MCP session as activity. + + Uses the raw (non-auto-starting) connect path and returns False when no + compatible daemon answers: a heartbeat keeps an existing daemon warm but + must NEVER start or restart one — otherwise `ccc daemon stop` during an + MCP session would fight the heartbeat loop. The daemon still auto-starts + lazily on the next real search/index request. + """ + try: + conn, _resp = _raw_connect_and_handshake() + except (ConnectionRefusedError, DaemonVersionError, OSError): + return False + try: + conn.send_bytes(encode_request(HeartbeatRequest())) + data = conn.recv_bytes() + except (EOFError, OSError): + return False + finally: + conn.close() + resp = decode_response(data) + return isinstance(resp, HeartbeatResponse) and resp.ok + + def daemon_env() -> DaemonEnvResponse: """Get environment variable names from the daemon.""" return _send(DaemonEnvRequest()) # type: ignore[return-value] diff --git a/src/cocoindex_code/daemon.py b/src/cocoindex_code/daemon.py index 1a339b9..d2b38ab 100644 --- a/src/cocoindex_code/daemon.py +++ b/src/cocoindex_code/daemon.py @@ -12,7 +12,8 @@ import time import traceback from collections.abc import AsyncIterator, Callable -from multiprocessing.connection import Connection, Listener +from datetime import timedelta +from multiprocessing.connection import Client, Connection, Listener from pathlib import Path from typing import Any @@ -42,6 +43,8 @@ ErrorResponse, HandshakeRequest, HandshakeResponse, + HeartbeatRequest, + HeartbeatResponse, IndexRequest, IndexStreamResponse, IndexWaitingNotice, @@ -60,6 +63,7 @@ ) from .settings import ( ChunkerMapping, + DaemonSettings, UserSettings, format_path_for_display, get_host_path_mappings, @@ -196,6 +200,50 @@ def list_projects(self) -> list[DaemonProjectInfo]: for root, project in self._projects.items() ] + def any_indexing(self) -> bool: + """True when any loaded project currently holds its index lock.""" + return any(project._index_lock.locked() for project in self._projects.values()) + + +# --------------------------------------------------------------------------- +# Idle reaper +# --------------------------------------------------------------------------- + + +class IdleReaper: + """Tracks client activity and decides when the daemon should idle-exit. + + ``last_activity`` is refreshed on every accepted connection and again when + each handler task finishes, so a long streaming request (e.g. an index run) + counts as activity up to its end. + """ + + def __init__(self, timeout: timedelta, *, supervised: bool) -> None: + self.timeout = timeout + self.supervised = supervised + self.last_activity = time.monotonic() + + def record_activity(self) -> None: + self.last_activity = time.monotonic() + + def idle_seconds(self, now: float | None = None) -> float: + return (time.monotonic() if now is None else now) - self.last_activity + + def should_exit(self, *, now: float, active_handlers: int, indexing: bool) -> bool: + """Pure exit predicate — testable without an event loop. + + Exit only when a timeout is configured, no external supervisor owns + the daemon lifecycle, no handler task is live, no project is indexing, + and the idle period exceeds the timeout. + """ + if self.timeout <= timedelta(0): + return False + if self.supervised: + return False + if active_handlers > 0 or indexing: + return False + return now - self.last_activity > self.timeout.total_seconds() + # --------------------------------------------------------------------------- # Connection handler @@ -210,6 +258,7 @@ async def handle_connection( settings_mtime_us: int | None, settings_env_names: list[str], handshake_warnings: list[str], + reaper: IdleReaper, ) -> None: """Handle a single client connection (per-request model). @@ -247,7 +296,7 @@ async def handle_connection( data = await loop.run_in_executor(None, conn.recv_bytes) req = decode_request(data) - result = await _dispatch(req, registry, start_time, on_shutdown, settings_env_names) + result = await _dispatch(req, registry, start_time, on_shutdown, settings_env_names, reaper) if isinstance(result, AsyncIterator): try: async for resp in result: @@ -463,6 +512,7 @@ async def _dispatch( start_time: float, on_shutdown: Callable[[], None], settings_env_names: list[str], + reaper: IdleReaper, ) -> ( Response | AsyncIterator[IndexStreamResponse] @@ -506,10 +556,13 @@ async def _dispatch( return project.get_status() if isinstance(req, DaemonStatusRequest): + now = time.monotonic() return DaemonStatusResponse( version=__version__, - uptime_seconds=time.monotonic() - start_time, + uptime_seconds=now - start_time, projects=registry.list_projects(), + idle_seconds=reaper.idle_seconds(now), + idle_timeout_minutes=round(reaper.timeout / timedelta(minutes=1)), ) if isinstance(req, RemoveProjectRequest): @@ -520,6 +573,11 @@ async def _dispatch( on_shutdown() return StopResponse(ok=True) + if isinstance(req, HeartbeatRequest): + # A heartbeat's only effect is the connection itself, which the + # accept path already recorded as activity. + return HeartbeatResponse(ok=True) + if isinstance(req, DaemonEnvRequest): from .protocol import DbPathMappingEntry from .settings import get_db_path_mappings @@ -551,12 +609,21 @@ async def _dispatch( # --------------------------------------------------------------------------- -def run_daemon() -> None: +def run_daemon( + *, + idle_timeout: timedelta | None = None, + idle_check_interval: timedelta = timedelta(minutes=1), +) -> None: """Main entry point for the daemon process (blocking). Sets up the listener, runs the asyncio event loop (``loop.run_forever``) to serve connections, and performs cleanup when shutdown is requested via - ``StopRequest`` or a signal (SIGTERM / SIGINT). + ``StopRequest``, a signal (SIGTERM / SIGINT), or the idle timeout. + + ``idle_timeout`` and ``idle_check_interval`` exist so tests can run a + seconds-scale idle timeout in-process; production leaves them at their + defaults and the timeout comes from ``daemon.idle_timeout_minutes`` in + ``global_settings.yml`` (the dataclass default when the file is missing). """ daemon_runtime_dir().mkdir(parents=True, exist_ok=True) @@ -570,8 +637,10 @@ def run_daemon() -> None: indexing_params: dict[str, Any] = {} query_params: dict[str, Any] = {} handshake_warnings: list[str] = [] + daemon_settings = DaemonSettings() if user_settings_path().is_file(): user_settings = load_user_settings() + daemon_settings = user_settings.daemon settings_env_keys = list(user_settings.envs.keys()) for key, value in user_settings.envs.items(): os.environ[key] = value @@ -630,10 +699,19 @@ def run_daemon() -> None: loop = asyncio.new_event_loop() tasks: set[asyncio.Task[Any]] = set() + reaper = IdleReaper( + timeout=( + idle_timeout + if idle_timeout is not None + else timedelta(minutes=daemon_settings.idle_timeout_minutes) + ), + supervised=os.environ.get("COCOINDEX_CODE_DAEMON_SUPERVISED") == "1", + ) + shutdown_reason: str | None = None def _request_shutdown(reason: str) -> None: - """Trigger daemon shutdown — called by StopRequest or signal handler. + """Trigger daemon shutdown — called by StopRequest, signal, or idle reaper. Records *reason* (first one wins) for the graceful-exit marker written during cleanup. @@ -644,6 +722,7 @@ def _request_shutdown(reason: str) -> None: loop.stop() def _spawn_handler(conn: Connection) -> None: + reaper.record_activity() task = loop.create_task( handle_connection( conn, @@ -653,10 +732,37 @@ def _spawn_handler(conn: Connection) -> None: settings_mtime_us, settings_env_keys, handshake_warnings, + reaper, ) ) tasks.add(task) - task.add_done_callback(tasks.discard) + + def _on_done(t: asyncio.Task[Any]) -> None: + tasks.discard(t) + # A finished handler is activity too — a long streaming request + # (e.g. an index run) counts up to its end, not just its start. + reaper.record_activity() + + task.add_done_callback(_on_done) + + async def _idle_reaper_loop() -> None: + """Periodically check the idle predicate; trigger shutdown when it fires.""" + while True: + await asyncio.sleep(idle_check_interval.total_seconds()) + if reaper.should_exit( + now=time.monotonic(), + active_handlers=len(tasks), + indexing=registry.any_indexing(), + ): + logger.info( + "Idle for %.1fm (timeout %.1fm), shutting down", + reaper.idle_seconds() / 60, + reaper.timeout / timedelta(minutes=1), + ) + _request_shutdown("idle_timeout") + return + + reaper_task = loop.create_task(_idle_reaper_loop()) # Handle signals for graceful shutdown try: @@ -666,15 +772,28 @@ def _spawn_handler(conn: Connection) -> None: pass # Not in main thread, or not supported on this platform (e.g. Windows) # Accept loop runs in a background thread; new connections are dispatched - # to the event loop via call_soon_threadsafe. The loop exits when - # listener.close() (called during shutdown) causes accept() to raise. + # to the event loop via call_soon_threadsafe. On POSIX the loop exits when + # listener.close() (called during shutdown) causes accept() to raise. On + # Windows, PipeListener.close() does not cancel an in-flight + # ConnectNamedPipe wait, so shutdown additionally sets `shutting_down` and + # wakes accept() with a dummy client connection; the accepted connection + # (dummy or real) is dropped and the thread exits. + shutting_down = threading.Event() + def _accept_loop() -> None: while True: try: conn = listener.accept() - loop.call_soon_threadsafe(_spawn_handler, conn) except OSError: break + if shutting_down.is_set(): + conn.close() + break + try: + loop.call_soon_threadsafe(_spawn_handler, conn) + except RuntimeError: # loop already closed — shutdown race + conn.close() + break accept_thread = threading.Thread(target=_accept_loop, daemon=True) accept_thread.start() @@ -683,20 +802,34 @@ def _accept_loop() -> None: try: loop.run_forever() finally: - # 0. Record the graceful exit (StopRequest, signal). Written first so - # a client that races the rest of the cleanup already sees it. A - # crashed daemon never gets here — the marker's absence is how the - # client detects a crash. + # 0. Record the graceful exit (idle timeout, StopRequest, signal). + # Written first so a client that races the rest of the cleanup + # already sees it. A crashed daemon never gets here — the marker's + # absence is how the client detects a crash. write_last_exit_marker(pid=os.getpid(), reason=shutdown_reason or "unknown") - # 1. Stop accepting new connections. + # 1. Stop accepting new connections. On Windows a blocked accept() + # holds an open pipe instance that listener.close() can't release + # (it only closes the queued next-instance handle), which would keep + # the named pipe alive after exit — wake it with a dummy connection + # so the accept thread closes it and exits. POSIX doesn't need the + # wake-up (listener.close() makes accept() raise), but it's harmless + # there and keeps this path exercised on every platform rather than + # only on Windows CI. + shutting_down.set() + try: + Client(sock_path, family=connection_family()).close() + except OSError: + pass listener.close() + accept_thread.join(timeout=5) - # 2. Cancel handler tasks (they may be blocked in run_in_executor). - for task in tasks: + # 2. Cancel handler tasks (they may be blocked in run_in_executor) + # and the idle-reaper task. + pending = [*tasks, reaper_task] + for task in pending: task.cancel() - if tasks: - loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True)) + loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) # 3. Release project resources. registry.close_all() diff --git a/src/cocoindex_code/protocol.py b/src/cocoindex_code/protocol.py index 6bcde27..164dedf 100644 --- a/src/cocoindex_code/protocol.py +++ b/src/cocoindex_code/protocol.py @@ -50,6 +50,16 @@ class DaemonEnvRequest(_msgspec.Struct, tag="daemon_env"): pass +class HeartbeatRequest(_msgspec.Struct, tag="heartbeat"): + """Counts as client activity for the daemon's idle timeout. + + Sent periodically by long-lived MCP servers so the daemon does not + idle-exit under an active session. + """ + + pass + + Request = ( HandshakeRequest | IndexRequest @@ -60,6 +70,7 @@ class DaemonEnvRequest(_msgspec.Struct, tag="daemon_env"): | StopRequest | DoctorRequest | DaemonEnvRequest + | HeartbeatRequest ) # --------------------------------------------------------------------------- @@ -144,6 +155,10 @@ class DaemonStatusResponse(_msgspec.Struct, tag="daemon_status"): version: str uptime_seconds: float projects: list[DaemonProjectInfo] + # Idle-timeout observability: seconds since the last client activity and + # the configured timeout (0 = never exit). + idle_seconds: float + idle_timeout_minutes: int class RemoveProjectResponse(_msgspec.Struct, tag="remove_project"): @@ -154,6 +169,10 @@ class StopResponse(_msgspec.Struct, tag="stop"): ok: bool +class HeartbeatResponse(_msgspec.Struct, tag="heartbeat"): + ok: bool + + class DoctorCheckResult(_msgspec.Struct): name: str ok: bool @@ -198,6 +217,7 @@ class ErrorResponse(_msgspec.Struct, tag="error"): | DaemonStatusResponse | RemoveProjectResponse | StopResponse + | HeartbeatResponse | DoctorResponse | DaemonEnvResponse | ErrorResponse diff --git a/src/cocoindex_code/server.py b/src/cocoindex_code/server.py index 2708c86..8b1ffaf 100644 --- a/src/cocoindex_code/server.py +++ b/src/cocoindex_code/server.py @@ -17,6 +17,8 @@ from mcp.server.fastmcp import FastMCP from pydantic import BaseModel, Field +from .settings import DaemonSettings, load_user_settings + _MCP_INSTRUCTIONS = ( "Code search and codebase understanding tools." "\n" @@ -159,6 +161,45 @@ async def search( return mcp +# === Daemon heartbeat loop === + + +def heartbeat_interval_s(idle_timeout_minutes: int) -> int: + """Heartbeat period for a given idle timeout: a third of the timeout, + clamped to [30 s, 300 s] so several heartbeats always fit in one idle + window without hammering the daemon on very long timeouts. + """ + return max(30, min(300, idle_timeout_minutes * 60 // 3)) + + +async def run_heartbeat_loop() -> None: + """Periodically heartbeat the daemon so it never idle-exits under this + live MCP session; when this process dies (even SIGKILL) the heartbeats + stop and the daemon exits on its normal idle timeout. + + Reads the timeout from the same ``global_settings.yml`` the daemon reads + (dataclass default when the file is missing or invalid). Returns + immediately when the timeout is 0 (daemon never idle-exits). The + heartbeat itself never starts or restarts a daemon — see + ``client.send_heartbeat``. + """ + from .client import send_heartbeat + + try: + timeout_minutes = load_user_settings().daemon.idle_timeout_minutes + except (FileNotFoundError, ValueError): + timeout_minutes = DaemonSettings().idle_timeout_minutes + if timeout_minutes <= 0: + return + + interval = heartbeat_interval_s(timeout_minutes) + loop = asyncio.get_event_loop() + while True: + await asyncio.sleep(interval) + # Client I/O is blocking — run off the event loop thread. + await loop.run_in_executor(None, send_heartbeat) + + # Keep the old `mcp` global for backward compatibility in __init__.py mcp: FastMCP | None = None @@ -326,7 +367,14 @@ def _on_progress(progress: IndexingProgress) -> None: async def _serve() -> None: from .cli import _bg_index - asyncio.create_task(_bg_index(str(project_root))) - await mcp_server.run_stdio_async() + background_tasks = { + asyncio.create_task(_bg_index(str(project_root))), + asyncio.create_task(run_heartbeat_loop()), + } + try: + await mcp_server.run_stdio_async() + finally: + for task in background_tasks: + task.cancel() asyncio.run(_serve()) diff --git a/src/cocoindex_code/settings.py b/src/cocoindex_code/settings.py index 57d117f..e228839 100644 --- a/src/cocoindex_code/settings.py +++ b/src/cocoindex_code/settings.py @@ -105,10 +105,18 @@ class EmbeddingSettings: query_params: dict[str, Any] | None = None +@dataclass +class DaemonSettings: + # Minutes without client activity before the daemon exits (0 = never exit). + # Clients auto-restart the daemon on the next request, so exiting is cheap. + idle_timeout_minutes: int = 180 + + @dataclass class UserSettings: embedding: EmbeddingSettings envs: dict[str, str] = field(default_factory=dict) + daemon: DaemonSettings = field(default_factory=DaemonSettings) @dataclass @@ -433,6 +441,8 @@ def _user_settings_to_dict(settings: UserSettings) -> dict[str, Any]: d: dict[str, Any] = {"embedding": _embedding_settings_to_dict(settings.embedding)} if settings.envs: d["envs"] = dict(settings.envs) + if settings.daemon != DaemonSettings(): + d["daemon"] = {"idle_timeout_minutes": settings.daemon.idle_timeout_minutes} return d @@ -457,7 +467,13 @@ def _user_settings_from_dict(d: dict[str, Any]) -> UserSettings: emb_kwargs["query_params"] = dict(emb_dict["query_params"] or {}) embedding = EmbeddingSettings(**emb_kwargs) envs = d.get("envs", {}) - return UserSettings(embedding=embedding, envs=envs) + # `daemon:` section is optional — missing keys use the dataclass defaults. + daemon_dict = d.get("daemon") or {} + daemon_kwargs: dict[str, Any] = {} + if "idle_timeout_minutes" in daemon_dict: + daemon_kwargs["idle_timeout_minutes"] = int(daemon_dict["idle_timeout_minutes"]) + daemon = DaemonSettings(**daemon_kwargs) + return UserSettings(embedding=embedding, envs=envs, daemon=daemon) def _project_settings_to_dict(settings: ProjectSettings) -> dict[str, Any]: diff --git a/tests/test_daemon_idle.py b/tests/test_daemon_idle.py new file mode 100644 index 0000000..e7c4f25 --- /dev/null +++ b/tests/test_daemon_idle.py @@ -0,0 +1,294 @@ +"""Idle-timeout tests: IdleReaper predicate units + in-process daemon E2E. + +The E2E tests start the daemon in-process (a background thread) with a +seconds-scale idle timeout injected via ``run_daemon(idle_timeout=..., +idle_check_interval=...)`` and verify it exits — or refuses to exit — at +the right times, and that the socket and PID file are cleaned up. +""" + +from __future__ import annotations + +import os +import tempfile +import threading +import time +from datetime import timedelta +from multiprocessing.connection import Client, Connection +from pathlib import Path + +import pytest +from conftest import make_test_user_settings + +import cocoindex_code.daemon as dm +from cocoindex_code._daemon_paths import ( + connection_family, + daemon_pid_path, + read_last_exit_marker, +) +from cocoindex_code._version import __version__ +from cocoindex_code.daemon import IdleReaper +from cocoindex_code.protocol import ( + HandshakeRequest, + IndexProgressUpdate, + IndexRequest, + IndexResponse, + IndexWaitingNotice, + StopRequest, + decode_response, + encode_request, +) +from cocoindex_code.settings import ( + default_project_settings, + save_project_settings, + save_user_settings, +) + +# --------------------------------------------------------------------------- +# IdleReaper predicate (pure, no event loop) +# --------------------------------------------------------------------------- + + +def test_reaper_exits_after_timeout_elapsed() -> None: + reaper = IdleReaper(timeout=timedelta(seconds=60), supervised=False) + now = reaper.last_activity + 61.0 + assert reaper.should_exit(now=now, active_handlers=0, indexing=False) is True + + +def test_reaper_stays_before_timeout() -> None: + reaper = IdleReaper(timeout=timedelta(seconds=60), supervised=False) + now = reaper.last_activity + 59.0 + assert reaper.should_exit(now=now, active_handlers=0, indexing=False) is False + + +def test_reaper_stays_with_live_handler() -> None: + reaper = IdleReaper(timeout=timedelta(seconds=60), supervised=False) + now = reaper.last_activity + 61.0 + assert reaper.should_exit(now=now, active_handlers=1, indexing=False) is False + + +def test_reaper_stays_while_indexing() -> None: + reaper = IdleReaper(timeout=timedelta(seconds=60), supervised=False) + now = reaper.last_activity + 61.0 + assert reaper.should_exit(now=now, active_handlers=0, indexing=True) is False + + +def test_reaper_timeout_zero_never_exits() -> None: + reaper = IdleReaper(timeout=timedelta(0), supervised=False) + now = reaper.last_activity + 1_000_000.0 + assert reaper.should_exit(now=now, active_handlers=0, indexing=False) is False + + +def test_reaper_supervised_never_exits() -> None: + reaper = IdleReaper(timeout=timedelta(seconds=60), supervised=True) + now = reaper.last_activity + 1_000_000.0 + assert reaper.should_exit(now=now, active_handlers=0, indexing=False) is False + + +def test_reaper_activity_resets_idle_clock() -> None: + reaper = IdleReaper(timeout=timedelta(seconds=60), supervised=False) + reaper.last_activity -= 120.0 # long idle... + reaper.record_activity() # ...then a connection arrives + now = reaper.last_activity + 1.0 + assert reaper.should_exit(now=now, active_handlers=0, indexing=False) is False + assert reaper.idle_seconds(now) == pytest.approx(1.0) + + +# --------------------------------------------------------------------------- +# MCP heartbeat interval +# --------------------------------------------------------------------------- + + +def test_heartbeat_interval_is_a_third_of_the_timeout_clamped() -> None: + from cocoindex_code.server import heartbeat_interval_s + + assert heartbeat_interval_s(1) == 30 # 20 s → clamped up to 30 s + assert heartbeat_interval_s(10) == 200 # 600 s / 3 + assert heartbeat_interval_s(180) == 300 # 3600 s → clamped down to 300 s + + +# --------------------------------------------------------------------------- +# In-process daemon E2E +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def idle_env(monkeypatch: pytest.MonkeyPatch) -> Path: + """Point COCOINDEX_CODE_DIR at a fresh short temp dir (AF_UNIX path limit).""" + base_dir = Path(tempfile.mkdtemp(prefix="ccc_idle_")) + monkeypatch.setenv("COCOINDEX_CODE_DIR", str(base_dir)) + return base_dir + + +def _start_daemon_thread( + *, idle_timeout_s: float, idle_check_interval_s: float +) -> tuple[threading.Thread, str]: + """Run the daemon in-process and wait for its socket to appear.""" + thread = threading.Thread( + target=lambda: dm.run_daemon( + idle_timeout=timedelta(seconds=idle_timeout_s), + idle_check_interval=timedelta(seconds=idle_check_interval_s), + ), + daemon=True, + ) + thread.start() + + sock_path = dm.daemon_socket_path() + deadline = time.monotonic() + 30 + while time.monotonic() < deadline: + if os.path.exists(sock_path): + return thread, sock_path + if not thread.is_alive(): + raise RuntimeError("Daemon thread exited before its socket appeared") + time.sleep(0.05) + raise TimeoutError("In-process daemon did not start") + + +def _stop_daemon_thread(thread: threading.Thread, sock_path: str) -> None: + """Best-effort graceful stop for tests where the daemon is still running.""" + if not thread.is_alive(): + return + try: + conn = Client(sock_path, family=connection_family()) + conn.send_bytes(encode_request(HandshakeRequest(version=__version__))) + conn.recv_bytes() + conn.send_bytes(encode_request(StopRequest())) + conn.recv_bytes() + conn.close() + except Exception: + pass + thread.join(timeout=10) + + +def _assert_cleaned_up(sock_path: str) -> None: + assert not os.path.exists(sock_path), "socket file not cleaned up" + assert not daemon_pid_path().exists(), "PID file not cleaned up" + + +def _recv_index_stream(conn: Connection) -> IndexResponse: + while True: + resp = decode_response(conn.recv_bytes()) + if isinstance(resp, IndexProgressUpdate | IndexWaitingNotice): + continue + if isinstance(resp, IndexResponse): + return resp + raise AssertionError(f"Unexpected response during indexing: {type(resp).__name__}") + + +def test_daemon_exits_when_idle(idle_env: Path) -> None: + """With no client activity the daemon exits after the timeout, and the + socket + PID file are cleaned up. Runs in no-settings mode (no embedder) + to prove an unconfigured daemon still idle-exits. + """ + thread, sock_path = _start_daemon_thread(idle_timeout_s=1.0, idle_check_interval_s=0.2) + try: + thread.join(timeout=15) + assert not thread.is_alive(), "daemon did not idle-exit" + _assert_cleaned_up(sock_path) + + # The graceful exit left a marker recording the idle timeout as the + # reason (the in-process daemon shares pytest's pid). + marker = read_last_exit_marker() + assert marker is not None + assert marker.reason == "idle_timeout" + assert marker.pid == os.getpid() + finally: + _stop_daemon_thread(thread, sock_path) + + +def test_daemon_does_not_exit_with_live_connection(idle_env: Path) -> None: + """A live handler task (open connection awaiting its request) blocks the + idle exit; once the connection closes the daemon exits on the next timeout. + """ + thread, sock_path = _start_daemon_thread(idle_timeout_s=0.5, idle_check_interval_s=0.1) + try: + conn = Client(sock_path, family=connection_family()) + conn.send_bytes(encode_request(HandshakeRequest(version=__version__))) + conn.recv_bytes() + # Hold the connection open well past the timeout without sending a request. + time.sleep(1.5) + assert thread.is_alive(), "daemon idle-exited despite a live connection" + + conn.close() + thread.join(timeout=15) + assert not thread.is_alive(), "daemon did not exit after the connection closed" + _assert_cleaned_up(sock_path) + finally: + _stop_daemon_thread(thread, sock_path) + + +def test_daemon_does_not_exit_while_heartbeats_arrive(idle_env: Path) -> None: + """Periodic heartbeats keep the daemon alive past its idle timeout; once + they stop, it exits on the normal timeout. Mirrors a live MCP session + dying (heartbeats stop) without any explicit shutdown. + """ + from cocoindex_code.client import send_heartbeat + + thread, sock_path = _start_daemon_thread(idle_timeout_s=1.0, idle_check_interval_s=0.2) + try: + # Heartbeat every 0.4 s for ~3 s — several idle timeouts' worth. + deadline = time.monotonic() + 3.0 + while time.monotonic() < deadline: + assert send_heartbeat() is True + time.sleep(0.4) + assert thread.is_alive(), "daemon idle-exited despite heartbeats" + + # Heartbeats stopped — the daemon exits on the next idle timeout. + thread.join(timeout=15) + assert not thread.is_alive(), "daemon did not exit after heartbeats stopped" + _assert_cleaned_up(sock_path) + finally: + _stop_daemon_thread(thread, sock_path) + + +def test_send_heartbeat_does_not_spawn_a_daemon(idle_env: Path) -> None: + """Against a stopped daemon, send_heartbeat returns False and never starts + one — otherwise `ccc daemon stop` during an MCP session would fight the + heartbeat loop. + """ + from cocoindex_code.client import send_heartbeat + + sock_path = dm.daemon_socket_path() + assert not os.path.exists(sock_path) + + assert send_heartbeat() is False + + # Give a hypothetical spawned daemon time to appear, then confirm nothing did. + time.sleep(1.0) + assert not os.path.exists(sock_path) + assert not daemon_pid_path().exists() + + +def test_daemon_does_not_exit_during_index_run(idle_env: Path) -> None: + """An in-flight index run that outlives the idle timeout is not killed: + the stream completes successfully, then the daemon idle-exits afterwards. + """ + save_user_settings(make_test_user_settings()) + project = idle_env / "proj" + project.mkdir() + save_project_settings(project, default_project_settings()) + for i in range(60): + (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" + ) + + thread, sock_path = _start_daemon_thread(idle_timeout_s=1.0, idle_check_interval_s=0.2) + try: + conn = Client(sock_path, family=connection_family()) + conn.send_bytes(encode_request(HandshakeRequest(version=__version__))) + conn.recv_bytes() + conn.send_bytes(encode_request(IndexRequest(project_root=str(project)))) + # If the daemon idle-exited mid-run, the stream would break with + # EOFError instead of delivering a successful IndexResponse. + final = _recv_index_stream(conn) + assert final.success is True + conn.close() + + # After the run finishes and the connection closes, the idle clock + # restarts and the daemon exits on the next timeout. + thread.join(timeout=30) + assert not thread.is_alive(), "daemon did not idle-exit after the index run" + _assert_cleaned_up(sock_path) + finally: + _stop_daemon_thread(thread, sock_path) diff --git a/tests/test_protocol.py b/tests/test_protocol.py index 765cdad..f026764 100644 --- a/tests/test_protocol.py +++ b/tests/test_protocol.py @@ -14,6 +14,8 @@ ErrorResponse, HandshakeRequest, HandshakeResponse, + HeartbeatRequest, + HeartbeatResponse, IndexingProgress, IndexProgressUpdate, IndexRequest, @@ -124,6 +126,8 @@ def test_encode_decode_daemon_status_response() -> None: projects=[ DaemonProjectInfo(project_root="/tmp/proj", indexing=False), ], + idle_seconds=12.5, + idle_timeout_minutes=180, ) data = encode_response(resp) decoded = decode_response(data) @@ -133,6 +137,19 @@ def test_encode_decode_daemon_status_response() -> None: assert len(decoded.projects) == 1 assert decoded.projects[0].project_root == "/tmp/proj" assert decoded.projects[0].indexing is False + assert decoded.idle_seconds == 12.5 + assert decoded.idle_timeout_minutes == 180 + + +def test_encode_decode_heartbeat_round_trip() -> None: + req = HeartbeatRequest() + decoded_req = decode_request(encode_request(req)) + assert isinstance(decoded_req, HeartbeatRequest) + + resp = HeartbeatResponse(ok=True) + decoded_resp = decode_response(encode_response(resp)) + assert isinstance(decoded_resp, HeartbeatResponse) + assert decoded_resp.ok is True def test_tagged_union_dispatch() -> None: @@ -212,6 +229,7 @@ def test_all_request_types_round_trip() -> None: StopRequest(), DoctorRequest(project_root="/tmp"), DaemonEnvRequest(), + HeartbeatRequest(), ] for req in requests: data = encode_request(req) @@ -300,9 +318,16 @@ def test_all_response_types_round_trip() -> None: IndexWaitingNotice(), SearchResponse(success=True), ProjectStatusResponse(indexing=False, total_chunks=0, total_files=0, languages={}), - DaemonStatusResponse(version="1.0.0", uptime_seconds=0.0, projects=[]), + DaemonStatusResponse( + version="1.0.0", + uptime_seconds=0.0, + projects=[], + idle_seconds=0.0, + idle_timeout_minutes=180, + ), RemoveProjectResponse(ok=True), StopResponse(ok=True), + HeartbeatResponse(ok=True), DoctorResponse( result=DoctorCheckResult(name="test", ok=True, details=[], errors=[]), ), diff --git a/tests/test_settings.py b/tests/test_settings.py index b50bd07..f9be64b 100644 --- a/tests/test_settings.py +++ b/tests/test_settings.py @@ -14,6 +14,7 @@ DEFAULT_EXCLUDED_PATTERNS, DEFAULT_INCLUDED_PATTERNS, ChunkerMapping, + DaemonSettings, EmbeddingSettings, LanguageOverride, ProjectSettings, @@ -473,6 +474,62 @@ def test_invalid_env_raises(self, monkeypatch: pytest.MonkeyPatch) -> None: get_host_path_mappings() +# --------------------------------------------------------------------------- +# daemon settings (idle timeout) +# --------------------------------------------------------------------------- + + +@pytest.mark.usefixtures("_patch_user_dir") +def test_daemon_settings_absent_section_uses_default(tmp_path: Path) -> None: + path = tmp_path / ".cocoindex_code" / "global_settings.yml" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("embedding:\n provider: litellm\n model: m\n") + loaded = load_user_settings() + assert loaded.daemon.idle_timeout_minutes == 180 + + +@pytest.mark.usefixtures("_patch_user_dir") +def test_daemon_settings_parses_idle_timeout(tmp_path: Path) -> None: + path = tmp_path / ".cocoindex_code" / "global_settings.yml" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "embedding:\n provider: litellm\n model: m\ndaemon:\n idle_timeout_minutes: 30\n" + ) + loaded = load_user_settings() + assert loaded.daemon.idle_timeout_minutes == 30 + + +@pytest.mark.usefixtures("_patch_user_dir") +def test_daemon_settings_explicit_zero_means_never(tmp_path: Path) -> None: + path = tmp_path / ".cocoindex_code" / "global_settings.yml" + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + "embedding:\n provider: litellm\n model: m\ndaemon:\n idle_timeout_minutes: 0\n" + ) + loaded = load_user_settings() + assert loaded.daemon.idle_timeout_minutes == 0 + + +@pytest.mark.usefixtures("_patch_user_dir") +def test_daemon_settings_round_trip() -> None: + settings = UserSettings( + embedding=EmbeddingSettings(provider="litellm", model="m"), + daemon=DaemonSettings(idle_timeout_minutes=45), + ) + save_user_settings(settings) + loaded = load_user_settings() + assert loaded.daemon.idle_timeout_minutes == 45 + + +@pytest.mark.usefixtures("_patch_user_dir") +def test_daemon_settings_default_omitted_from_yaml() -> None: + from cocoindex_code.settings import user_settings_path + + settings = UserSettings(embedding=EmbeddingSettings(provider="litellm", model="m")) + save_user_settings(settings) + assert "daemon" not in user_settings_path().read_text() + + # --------------------------------------------------------------------------- # find_parent_with_marker — global-only should not match # ---------------------------------------------------------------------------