diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 536cb73de9e..c4ff325db1f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -13,8 +13,6 @@ from email.mime.text import MIMEText from typing import ( TYPE_CHECKING, Any, - Callable, - Coroutine, Dict, List, Literal, @@ -2277,6 +2275,11 @@ class PrismaClient: 0.0, float(os.getenv("PRISMA_AUTH_RECONNECT_LOCK_TIMEOUT_SECONDS", "0.1")), ) + self._engine_pidfd: int = -1 + self._engine_pid: int = 0 + self._watching_engine: bool = False + self._engine_confirmed_dead: bool = False + self._engine_wait_thread: Optional[threading.Thread] = None verbose_proxy_logger.debug("Success - Created Prisma Client") def get_request_status( @@ -3544,31 +3547,368 @@ class PrismaClient: ) raise e + def _get_engine_pid(self) -> int: + try: + engine = self.db._original_prisma._engine # type: ignore[attr-defined] + if engine is not None and engine.process is not None: + return engine.process.pid + except (AttributeError, TypeError): + pass + return 0 + + def _is_engine_alive(self) -> bool: + if self._engine_pid <= 0: + return True + try: + os.kill(self._engine_pid, 0) + return True + except ProcessLookupError: + return False + except (PermissionError, OSError): + return True + + @staticmethod + def _reap_all_zombies() -> set: + """Reap ALL zombie child processes via waitpid(-1, WNOHANG). + + Returns a set of reaped PIDs. As PID 1 in Docker (or any + process that spawns children), we must reap ALL terminated + children to prevent zombie accumulation. + """ + reaped: set = set() + while True: + try: + pid, _ = os.waitpid(-1, os.WNOHANG) + if pid == 0: + break + reaped.add(pid) + except ChildProcessError: + break + return reaped + + def _try_waitpid_watch(self, pid: int) -> bool: + """Watch engine PID via os.waitpid() in a dedicated thread. + + The thread blocks on os.waitpid(pid, 0) which is a kernel-level + wait and with zero CPU overhead, instant detection when the process exits. + When the process dies, the thread notifies the asyncio event loop + via call_soon_threadsafe. + + Returns True if the thread was started, False on failure. + """ + try: + probe_pid, _ = os.waitpid(pid, os.WNOHANG) + except ChildProcessError: + verbose_proxy_logger.debug( + "PID %s is not a child process; skipping waitpid watch.", pid, + ) + return False + + if probe_pid == pid: + verbose_proxy_logger.warning( + "prisma-query-engine PID %s already dead at watch start.", pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + asyncio.create_task( + self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + ) + return True + + try: + loop = asyncio.get_running_loop() + except RuntimeError: + return False + + thread = threading.Thread( + target=self._waitpid_thread_func, + args=(pid, loop), + daemon=True, + name=f"prisma-engine-waitpid-{pid}", + ) + thread.start() + self._engine_wait_thread = thread + return True + + def _waitpid_thread_func(self, pid: int, loop: asyncio.AbstractEventLoop) -> None: + """Thread function: block until engine PID exits, then notify event loop. + + Note: uvloop/libuv may reap the child first via waitpid(-1, WNOHANG) + in its SIGCHLD handler. In that case our waitpid raises ChildProcessError. + we still notify the event loop because the engine is dead either way. + """ + try: + os.waitpid(pid, 0) + except ChildProcessError: + pass + except OSError: + pass + try: + loop.call_soon_threadsafe(self._on_engine_death_from_thread, pid) + except RuntimeError: + pass + + def _on_engine_death_from_thread(self, dead_pid: int) -> None: + """Called on the event loop thread when the waitpid thread detects engine death.""" + if self._engine_confirmed_dead: + return + if dead_pid != self._engine_pid: + return + verbose_proxy_logger.error( + "prisma-query-engine PID %s exited (waitpid thread); triggering reconnect.", + dead_pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + asyncio.create_task( + self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + ) + + def _try_pidfd_watch(self, pid: int) -> bool: + """ + Watch engine PID via pidfd_open + asyncio event loop reader. + + Returns True if pidfd watch was set up, False if unavailable or failed. + Broad OSError catch handles both ENOSYS and SECCOMP-blocked syscalls. + """ + if not hasattr(os, "pidfd_open"): + return False + fd = -1 + try: + fd = os.pidfd_open(pid, 0) # type: ignore[attr-defined] + asyncio.get_running_loop().add_reader(fd, self._on_pidfd_readable) + self._engine_pidfd = fd + return True + except OSError: + if fd >= 0: + os.close(fd) + return False + + def _on_pidfd_readable(self) -> None: + """pidfd became readable: engine process exited or became zombie. + + Sets _engine_confirmed_dead BEFORE cleanup so _run_reconnect_cycle + takes the heavy path (recreate Prisma client + re-arm watcher). + """ + if self._engine_confirmed_dead: + # Already handled -- just clean up pidfd resources. + if self._engine_pidfd >= 0: + try: + asyncio.get_running_loop().remove_reader(self._engine_pidfd) + except Exception: + pass + try: + os.close(self._engine_pidfd) + except OSError: + pass + self._engine_pidfd = -1 + return + dead_pid = self._engine_pid + verbose_proxy_logger.error( + "prisma-query-engine PID %s exited (pidfd event); triggering reconnect.", + dead_pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + asyncio.create_task( + self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + ) + + async def _poll_engine_proc(self) -> None: + """poll via os.kill(pid, 0) every 1s. + Only used when BOTH waitpid thread and pidfd are unavailable + (e.g., PID is not our child process and pidfd_open fails) + """ + while self._watching_engine and self._engine_pid > 0: + try: + os.kill(self._engine_pid, 0) + except ProcessLookupError: + verbose_proxy_logger.error( + "prisma-query-engine PID %s gone; triggering reconnect.", + self._engine_pid, + ) + self._engine_confirmed_dead = True + self._reap_all_zombies() + self._cleanup_engine_watcher() + await self.attempt_db_reconnect( + reason="engine_process_death", + force=True, + ) + return + except (PermissionError, OSError): + verbose_proxy_logger.debug( + "Cannot signal PID %s; stopping engine poll.", + self._engine_pid, + ) + self._cleanup_engine_watcher() + return + await asyncio.sleep(1) + + def _cleanup_engine_watcher(self) -> None: + """Clean up pidfd reader, waitpid thread ref, or stop polling and reset state.""" + self._watching_engine = False + if self._engine_pidfd >= 0: + try: + asyncio.get_running_loop().remove_reader(self._engine_pidfd) + except Exception: + pass + try: + os.close(self._engine_pidfd) + except OSError: + pass + self._engine_pidfd = -1 + self._engine_wait_thread = None + self._engine_pid = 0 + + async def _start_engine_watcher(self) -> None: + """ + Start watching the Prisma query engine process for death. + + Detection priority: + 1. os.waitpid() in a dedicated thread, works with all event loops. + 2. pidfd_open kernel fd registered with asyncio. + 3. os.kill(pid, 0) polling (1s), last-resort fallback when neither + waitpid thread nor pidfd are available. + + """ + if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: + return + pid = self._get_engine_pid() + if pid == 0: + verbose_proxy_logger.debug("Could not find prisma-query-engine PID; engine death detection unavailable.") + return + self._engine_pid = pid + self._engine_confirmed_dead = False + verbose_proxy_logger.info("Found prisma-query-engine at PID %s.", pid) + waitpid_ok = self._try_waitpid_watch(pid) + pidfd_ok = False if waitpid_ok else self._try_pidfd_watch(pid) + if waitpid_ok: + verbose_proxy_logger.info( + "Watching engine PID %s via waitpid thread.", pid, + ) + elif pidfd_ok: + verbose_proxy_logger.info( + "Watching engine PID %s via pidfd.", pid, + ) + else: + verbose_proxy_logger.info( + "Watching engine PID %s via os.kill polling.", pid, + ) + self._watching_engine = True + asyncio.create_task(self._poll_engine_proc()) + + def _stop_engine_watcher(self) -> None: + """Stop watching the engine process and clean up all resources.""" + self._cleanup_engine_watcher() + self._engine_confirmed_dead = False + verbose_proxy_logger.debug("Stopped engine process watcher.") + async def _run_reconnect_cycle( self, timeout_seconds: Optional[float] = None ) -> None: """ - Run a reconnect cycle with direct db operations and a single overall timeout - budget to avoid long retries on hot paths (e.g. auth). + Run a reconnect cycle with a single overall timeout budget. + + Uses the _engine_confirmed_dead flag (set by waitpid thread / pidfd / poll + handlers) to choose between heavy reconnect (engine dead -- recreate + Prisma client, re-arm watcher) and lightweight reconnect (network + blip -- disconnect, connect, SELECT 1). """ - async def _do_direct_reconnect() -> None: - try: - await self.db.disconnect() - except Exception as disconnect_err: - verbose_proxy_logger.debug( - "Prisma DB disconnect before reconnect failed (ignored): %s", - disconnect_err, - ) - - await self.db.connect() - await self.db.query_raw("SELECT 1") - effective_timeout = ( - timeout_seconds - if timeout_seconds is not None - else self._db_watchdog_reconnect_timeout_seconds + timeout_seconds if timeout_seconds is not None else self._db_watchdog_reconnect_timeout_seconds ) - await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) + + engine_is_dead = self._engine_confirmed_dead or ( + self._engine_pid > 0 and not self._is_engine_alive() + ) + + if engine_is_dead: + dead_pid = self._engine_pid + verbose_proxy_logger.warning( + "prisma-query-engine PID %s is dead; reconnecting.", + dead_pid, + ) + self._reap_all_zombies() + self._cleanup_engine_watcher() + self._engine_confirmed_dead = False + + async def _do_heavy_reconnect() -> None: + db_url = os.getenv("DATABASE_URL", "") + if not db_url: + verbose_proxy_logger.error("DATABASE_URL not set; cannot recreate Prisma client.") + raise RuntimeError("DATABASE_URL not set") + await self.db.recreate_prisma_client(db_url) + await self._start_engine_watcher() + + await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) + else: + verbose_proxy_logger.debug("Performing Prisma DB reconnect (engine alive or unknown).") + + async def _do_direct_reconnect() -> None: + try: + await self.db.disconnect() + except Exception as disconnect_err: + verbose_proxy_logger.debug( + "Prisma DB disconnect before reconnect failed (ignored): %s", + disconnect_err, + ) + + await self.db.connect() + await self.db.query_raw("SELECT 1") + + await asyncio.wait_for(_do_direct_reconnect(), timeout=effective_timeout) + + async def _attempt_reconnect_inside_lock( + self, + force: bool, + reason: str, + timeout_seconds: Optional[float], + ) -> bool: + now = time.time() + if ( + force is False + and now - self._db_last_reconnect_attempt_ts + < self._db_reconnect_cooldown_seconds + ): + verbose_proxy_logger.debug( + "Skipping DB reconnect attempt inside lock due to cooldown. reason=%s", + reason, + ) + return False + + verbose_proxy_logger.warning( + "Attempting Prisma DB reconnect. reason=%s", reason + ) + + reconnect_succeeded = False + try: + await self._run_reconnect_cycle(timeout_seconds=timeout_seconds) + reconnect_succeeded = True + verbose_proxy_logger.info( + "Prisma DB reconnect succeeded. reason=%s", reason + ) + except Exception as reconnect_err: + verbose_proxy_logger.error( + "Prisma DB reconnect failed. reason=%s error=%s", + reason, + reconnect_err, + ) + finally: + self._db_last_reconnect_attempt_ts = time.time() + + return reconnect_succeeded async def attempt_db_reconnect( self, @@ -3595,59 +3935,10 @@ class PrismaClient: ) return False - async def _attempt_reconnect_inside_lock() -> bool: - now = time.time() - if ( - force is False - and now - self._db_last_reconnect_attempt_ts - < self._db_reconnect_cooldown_seconds - ): - verbose_proxy_logger.debug( - "Skipping DB reconnect attempt inside lock due to cooldown. reason=%s", - reason, - ) - return False - - verbose_proxy_logger.warning( - "Attempting Prisma DB reconnect. reason=%s", reason - ) - - reconnect_succeeded = False - try: - await self._run_reconnect_cycle(timeout_seconds=timeout_seconds) - reconnect_succeeded = True - verbose_proxy_logger.info( - "Prisma DB reconnect succeeded. reason=%s", reason - ) - except Exception as reconnect_err: - verbose_proxy_logger.error( - "Prisma DB reconnect failed. reason=%s error=%s", - reason, - reconnect_err, - ) - finally: - # Start cooldown after reconnect attempt has completed. - self._db_last_reconnect_attempt_ts = time.time() - - return reconnect_succeeded - if lock_timeout_seconds is None: async with self._db_reconnect_lock: - return await _attempt_reconnect_inside_lock() + return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) - return await self._attempt_reconnect_with_lock_timeout( - _attempt_reconnect_inside_lock, - reason=reason, - lock_timeout_seconds=lock_timeout_seconds, - ) - - async def _attempt_reconnect_with_lock_timeout( - self, - reconnect_fn: Callable[[], Coroutine[Any, Any, bool]], - reason: str, - lock_timeout_seconds: float, - ) -> bool: - """Acquire the reconnect lock with a timeout, then run reconnect_fn.""" lock_acquired_by_timeout_task = False async def _acquire_reconnect_lock() -> bool: @@ -3695,14 +3986,14 @@ class PrismaClient: return False try: - return await reconnect_fn() + return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) finally: self._db_reconnect_lock.release() async def start_db_health_watchdog_task(self) -> None: - """ - Start a background task that probes DB health and attempts reconnect on failure. - """ + """Start background tasks that monitor DB health: + - A periodic SELECT 1 probe that triggers reconnect on network/connection failure. + - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling.""" if self._db_health_watchdog_enabled is not True: verbose_proxy_logger.debug( "Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED" @@ -3720,11 +4011,11 @@ class PrismaClient: self._db_health_watchdog_probe_timeout_seconds, self._db_watchdog_reconnect_timeout_seconds, ) + await self._start_engine_watcher() async def stop_db_health_watchdog_task(self) -> None: - """ - Stop DB health watchdog task gracefully. - """ + """Stop DB health watchdog task and engine watcher gracefully.""" + self._stop_engine_watcher() if self._db_health_watchdog_task is None: return self._db_health_watchdog_task.cancel() diff --git a/tests/litellm/proxy/test_prisma_engine_watchdog.py b/tests/litellm/proxy/test_prisma_engine_watchdog.py new file mode 100644 index 00000000000..011b8002db2 --- /dev/null +++ b/tests/litellm/proxy/test_prisma_engine_watchdog.py @@ -0,0 +1,446 @@ +""" +Tests for PrismaClient engine watchdog: death detection and automatic reconnect. + +Covers: +- Engine PID discovery and liveness check +- Engine process gone (os.kill raises ProcessLookupError) → reconnect triggered +- PermissionError from os.kill → treated as alive (process exists but not ours) +- pidfd handler → schedules attempt_db_reconnect even when lock is held +- waitpid thread → instant cross-platform detection, triggers reconnect +- _run_reconnect_cycle branches: heavy path (engine dead) vs lightweight path (engine alive) +- _engine_confirmed_dead flag ensures heavy reconnect even after _engine_pid reset +- Successful heavy reconnect → watcher re-armed for new process +- Missing DATABASE_URL → graceful RuntimeError in reconnect cycle +- Shutdown → polling loop exits cleanly +""" + +import asyncio +import os +import threading +import time +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from litellm.proxy.utils import PrismaClient, ProxyLogging + + +@pytest.fixture(autouse=True) +def mock_prisma_binary(): + """Mock prisma.Prisma to avoid requiring generated Prisma binaries for unit tests.""" + import sys + + mock_module = MagicMock() + with patch.dict(sys.modules, {"prisma": mock_module}): + yield + + +@pytest.fixture +def mock_proxy_logging(): + proxy_logging = AsyncMock(spec=ProxyLogging) + proxy_logging.failure_handler = AsyncMock() + return proxy_logging + + +@pytest.fixture +def engine_client(mock_proxy_logging) -> PrismaClient: + """ + Minimal PrismaClient fixture for engine watchdog tests. + Uses the real constructor pattern from PR #21706 (database_url). + """ + client = PrismaClient(database_url="mock://test", proxy_logging_obj=mock_proxy_logging) + client.db = MagicMock() + client.db.recreate_prisma_client = AsyncMock() + client.db.disconnect = AsyncMock(return_value=None) + client.db.connect = AsyncMock(return_value=None) + client.db.query_raw = AsyncMock(return_value=[{"result": 1}]) + return client + + +# --------------------------------------------------------------------------- +# _is_engine_alive +# --------------------------------------------------------------------------- + + +def test_is_engine_alive_returns_true_when_pid_unknown(engine_client): + """_is_engine_alive returns True when no engine PID is tracked.""" + engine_client._engine_pid = 0 + assert engine_client._is_engine_alive() is True + + +def test_is_engine_alive_returns_false_when_process_gone(engine_client): + """_is_engine_alive returns False when os.kill raises ProcessLookupError.""" + engine_client._engine_pid = 9999 + with patch("os.kill", side_effect=ProcessLookupError): + assert engine_client._is_engine_alive() is False + + +def test_is_engine_alive_returns_true_on_permission_error(engine_client): + """_is_engine_alive returns True when os.kill raises PermissionError (process exists but not ours).""" + engine_client._engine_pid = 1234 + with patch("os.kill", side_effect=PermissionError): + assert engine_client._is_engine_alive() is True + + +def test_is_engine_alive_returns_true_for_running_process(engine_client): + """_is_engine_alive returns True when os.kill succeeds (process running).""" + engine_client._engine_pid = 1234 + with patch("os.kill"): + assert engine_client._is_engine_alive() is True + + +# --------------------------------------------------------------------------- +# _poll_engine_proc — calls attempt_db_reconnect on death +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_poll_missing_process_triggers_reconnect(engine_client) -> None: + """Polling loop triggers attempt_db_reconnect when os.kill raises ProcessLookupError.""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + with patch("os.kill", side_effect=ProcessLookupError): + await engine_client._poll_engine_proc() + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +@pytest.mark.asyncio +async def test_poll_permission_error_stops_polling(engine_client) -> None: + """Polling loop stops cleanly when os.kill raises PermissionError (process not ours).""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + with patch("os.kill", side_effect=PermissionError): + await engine_client._poll_engine_proc() + + # PermissionError means process exists but isn't ours — no reconnect, just stop polling + engine_client.attempt_db_reconnect.assert_not_awaited() + assert engine_client._watching_engine is False + assert engine_client._engine_pid == 0 + + +@pytest.mark.asyncio +async def test_stop_loop_halts_polling(engine_client) -> None: + """Polling loop exits cleanly when _stop_engine_watcher is called.""" + engine_client._engine_pid = 1234 + engine_client._watching_engine = True + + async def stop_during_sleep(_duration: float) -> None: + engine_client._stop_engine_watcher() + + with ( + patch("os.kill"), + patch("asyncio.sleep", side_effect=stop_during_sleep), + ): + await engine_client._poll_engine_proc() + + assert engine_client._watching_engine is False + assert engine_client._engine_pid == 0 + + +# --------------------------------------------------------------------------- +# _on_pidfd_readable — calls attempt_db_reconnect +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_pidfd_readable_schedules_reconnect(engine_client) -> None: + """pidfd handler schedules attempt_db_reconnect via asyncio.create_task.""" + engine_client._engine_pid = 1234 + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + with patch("asyncio.create_task", side_effect=capture_task): + engine_client._on_pidfd_readable() + + # Run the captured coroutine to completion + assert len(created_coros) == 1 + await created_coros[0] + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +@pytest.mark.asyncio +async def test_pidfd_schedules_reconnect_task_when_lock_held(engine_client) -> None: + """pidfd handler schedules reconnect task even when _db_reconnect_lock is held.""" + engine_client._engine_pid = 1234 + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + async with engine_client._db_reconnect_lock: + with patch("asyncio.create_task", side_effect=capture_task): + engine_client._on_pidfd_readable() + + for coro in created_coros: + coro.close() + + assert len(created_coros) == 1 + + +# --------------------------------------------------------------------------- +# _run_reconnect_cycle — engine liveness branching +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_heavy_path_when_engine_dead( + engine_client, +) -> None: + """_run_reconnect_cycle calls recreate_prisma_client when engine is dead.""" + engine_client._engine_pid = 1234 + engine_client._start_engine_watcher = AsyncMock() + + with ( + patch.object(engine_client, "_is_engine_alive", return_value=False), + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + patch("os.waitpid", side_effect=ChildProcessError), + ): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") + engine_client._start_engine_watcher.assert_awaited_once() + engine_client.db.connect.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_heavy_path_when_confirmed_dead( + engine_client, +) -> None: + """_run_reconnect_cycle takes heavy path when _engine_confirmed_dead is set. + + This is the critical race-condition fix: SIGCHLD/pidfd handlers set + _engine_confirmed_dead BEFORE _cleanup_engine_watcher resets _engine_pid + to 0, so the heavy path executes even after cleanup. + """ + engine_client._engine_pid = 0 # Already reset by cleanup! + engine_client._engine_confirmed_dead = True # But flag survives cleanup + engine_client._start_engine_watcher = AsyncMock() + + with ( + patch.dict(os.environ, {"DATABASE_URL": "postgresql://test"}), + patch("os.waitpid", side_effect=ChildProcessError), + ): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_awaited_once_with("postgresql://test") + engine_client._start_engine_watcher.assert_awaited_once() + engine_client.db.connect.assert_not_awaited() + assert engine_client._engine_confirmed_dead is False # Reset after use + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_lightweight_path_when_engine_alive( + engine_client, +) -> None: + """_run_reconnect_cycle uses disconnect/connect when engine is alive.""" + engine_client._engine_pid = 1234 + + with patch.object(engine_client, "_is_engine_alive", return_value=True): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.connect.assert_awaited_once() + engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_uses_lightweight_path_when_pid_unknown( + engine_client, +) -> None: + """_run_reconnect_cycle uses lightweight path when engine PID is not tracked.""" + engine_client._engine_pid = 0 + + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.connect.assert_awaited_once() + engine_client.db.query_raw.assert_awaited_once_with("SELECT 1") + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_run_reconnect_cycle_heavy_path_raises_without_database_url( + engine_client, +) -> None: + """Heavy reconnect raises RuntimeError when DATABASE_URL is not set.""" + engine_client._engine_pid = 1234 + + with ( + patch.object(engine_client, "_is_engine_alive", return_value=False), + patch.dict(os.environ, {}, clear=True), + patch("os.waitpid", side_effect=ChildProcessError), + ): + with pytest.raises(RuntimeError, match="DATABASE_URL not set"): + await engine_client._run_reconnect_cycle(timeout_seconds=5.0) + + engine_client.db.recreate_prisma_client.assert_not_awaited() + + +# --------------------------------------------------------------------------- +# start/stop lifecycle integration +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_start_watchdog_task_also_starts_engine_watcher( + engine_client, +) -> None: + """start_db_health_watchdog_task() also starts engine watcher.""" + engine_client._start_engine_watcher = AsyncMock() + + loop = asyncio.get_running_loop() + dummy_task = loop.create_task(asyncio.sleep(3600)) + + def fake_create_task(coro): + coro.close() + return dummy_task + + with patch("asyncio.create_task", side_effect=fake_create_task): + await engine_client.start_db_health_watchdog_task() + + engine_client._start_engine_watcher.assert_awaited_once() + dummy_task.cancel() + try: + await dummy_task + except asyncio.CancelledError: + pass + + +@pytest.mark.asyncio +async def test_stop_watchdog_task_also_stops_engine_watcher( + engine_client, +) -> None: + """stop_db_health_watchdog_task() also stops engine watcher.""" + engine_client._stop_engine_watcher = MagicMock() + + loop = asyncio.get_running_loop() + dummy_task = loop.create_task(asyncio.sleep(3600)) + engine_client._db_health_watchdog_task = dummy_task + + await engine_client.stop_db_health_watchdog_task() + + engine_client._stop_engine_watcher.assert_called_once() + assert engine_client._db_health_watchdog_task is None + + +# --------------------------------------------------------------------------- +# waitpid thread (cross-platform) +# --------------------------------------------------------------------------- + + +def test_try_waitpid_watch_returns_false_when_not_child(engine_client): + """_try_waitpid_watch returns False when PID is not our child process.""" + engine_client._engine_pid = 9999 + with patch("os.waitpid", side_effect=ChildProcessError): + assert engine_client._try_waitpid_watch(9999) is False + assert engine_client._engine_wait_thread is None + + +def test_try_waitpid_watch_starts_thread_for_child(engine_client): + """_try_waitpid_watch starts a daemon thread when PID is our child.""" + engine_client._engine_pid = 1234 + mock_thread = MagicMock() + mock_loop = MagicMock() + with ( + patch("os.waitpid", return_value=(0, 0)), + patch("asyncio.get_running_loop", return_value=mock_loop), + patch("threading.Thread", return_value=mock_thread) as mock_thread_cls, + ): + result = engine_client._try_waitpid_watch(1234) + assert result is True + mock_thread.start.assert_called_once() + assert engine_client._engine_wait_thread is mock_thread + + +@pytest.mark.asyncio +async def test_try_waitpid_watch_handles_already_dead_engine(engine_client) -> None: + """_try_waitpid_watch detects engine already dead at watch start.""" + engine_client._engine_pid = 1234 + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + waitpid_calls = iter([(1234, 0)]) + + def mock_waitpid(pid, flags): + if pid == -1: + raise ChildProcessError + return next(waitpid_calls) + + with ( + patch("os.waitpid", side_effect=mock_waitpid), + patch("asyncio.create_task", side_effect=capture_task), + ): + result = engine_client._try_waitpid_watch(1234) + + assert result is True + assert engine_client._engine_confirmed_dead is True + assert len(created_coros) == 1 + created_coros[0].close() + + +@pytest.mark.asyncio +async def test_on_engine_death_from_thread_triggers_reconnect(engine_client) -> None: + """waitpid thread callback schedules attempt_db_reconnect.""" + engine_client._engine_pid = 1234 + engine_client.attempt_db_reconnect = AsyncMock(return_value=True) + + created_coros = [] + + def capture_task(coro): + created_coros.append(coro) + return MagicMock() + + with patch("asyncio.create_task", side_effect=capture_task): + engine_client._on_engine_death_from_thread(1234) + + assert len(created_coros) == 1 + await created_coros[0] + + engine_client.attempt_db_reconnect.assert_awaited_once_with( + reason="engine_process_death", + force=True, + ) + + +def test_on_engine_death_from_thread_no_double_trigger(engine_client): + """waitpid thread callback does not trigger reconnect if already confirmed dead.""" + engine_client._engine_pid = 1234 + engine_client._engine_confirmed_dead = True + + with patch("asyncio.create_task") as mock_create_task: + engine_client._on_engine_death_from_thread(1234) + + mock_create_task.assert_not_called() + + +def test_on_engine_death_from_thread_ignores_stale_pid(engine_client): + """waitpid thread callback ignores death notification for a stale PID.""" + engine_client._engine_pid = 5678 + + with patch("asyncio.create_task") as mock_create_task: + engine_client._on_engine_death_from_thread(1234) + + mock_create_task.assert_not_called()