fix: wait for daemon health before reclaiming occupied port

hindsight-embed/hindsight_embed/daemon_embed_manager.py

@@ -6,6 +6,7 @@
 """
 
 import logging
+import math
 import os
 import platform
 import re
@@ -30,12 +31,42 @@
 # Suppress noisy httpx logs
 logging.getLogger("httpx").setLevel(logging.WARNING)
 
+
+def _parse_float_env(name: str, default: float) -> float:
+    """Parse a float environment variable, falling back on invalid values."""
+    try:
+        return float(os.getenv(name, str(default)))
+    except ValueError:
+        return default
+
+
+def _safe_non_negative_float(value: float, fallback: float) -> float:
+    """Return a finite non-negative float, or fallback for invalid values."""
+    return value if math.isfinite(value) and value >= 0 else fallback
+
+
+def _safe_positive_float(value: float, fallback: float) -> float:
+    """Return a finite positive float, or fallback for invalid values."""
+    return value if math.isfinite(value) and value > 0 else fallback
+
+
 # Constants
 # Allow CI/Windows to extend the startup budget — pg0-embedded's Windows wheel
 # unpacks and runs initdb on first boot, which takes noticeably longer on cold
 # runners than POSIX.
 DAEMON_STARTUP_TIMEOUT = int(os.getenv("HINDSIGHT_EMBED_DAEMON_STARTUP_TIMEOUT", "180"))
 DEFAULT_DAEMON_IDLE_TIMEOUT = 0  # 0 = disabled (no auto-exit)
+# When another process is concurrently starting the daemon, the TCP port can be
+# bound before /health returns 200. Give that warming daemon a short grace window
+# before treating the listener as stale/foreign and attempting to reclaim it.
+PORT_HEALTH_GRACE_TIMEOUT = _safe_non_negative_float(
+    _parse_float_env("HINDSIGHT_EMBED_PORT_HEALTH_GRACE_TIMEOUT", 30.0),
+    30.0,
+)
+PORT_HEALTH_CHECK_INTERVAL = _safe_positive_float(
+    _parse_float_env("HINDSIGHT_EMBED_PORT_HEALTH_CHECK_INTERVAL", 0.5),
+    0.5,
+)
 
 
 def _detach_popen_kwargs(log_handle: IO[bytes]) -> dict:
@@ -219,6 +250,40 @@ def _kill_process(pid: int) -> bool:
             return True  # Already gone
         return False
 
+    @staticmethod
+    def _port_health_ok(port: int) -> bool:
+        """Return True when the listener on port responds like initialized Hindsight."""
+        try:
+            with httpx.Client(timeout=2) as client:
+                response = client.get(f"http://127.0.0.1:{port}/health")
+                if response.status_code != 200:
+                    return False
+                try:
+                    health = response.json()
+                except Exception:
+                    return False
+                return health.get("status") == "healthy" and health.get("database") == "connected"
+        except Exception:
+            return False
+
+    def _wait_for_port_health(self, port: int, timeout: float | None = None) -> bool:
+        """Wait briefly for a just-bound daemon port to become healthy."""
+        timeout = _safe_non_negative_float(
+            PORT_HEALTH_GRACE_TIMEOUT if timeout is None else timeout,
+            0.0,
+        )
+        interval = _safe_positive_float(PORT_HEALTH_CHECK_INTERVAL, 0.5)
+        deadline = time.monotonic() + timeout
+        while True:
+            if not self._is_port_in_use(port):
+                return False
+            if self._port_health_ok(port):
+                return True
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                return False
+            time.sleep(min(interval, remaining))
+
     def _clear_port(self, port: int) -> bool:
         """
         Ensure the port is free before starting a daemon.
@@ -229,29 +294,27 @@ def _clear_port(self, port: int) -> bool:
             The caller's "start" is effectively a no-op: the daemon is already up.
             Killing it would race concurrent starts (one process kills the other's
             freshly-started daemon, both rush to rebind the port).
-          * Port occupied but /health is unreachable / non-200 → treat as a stale
-            hindsight daemon (or foreign process) and attempt to reclaim by killing
-            the PID listening on the port. This preserves the original intent of
-            clearing stale daemons from version upgrades.
+          * Port occupied but /health is unreachable, non-200, or does not return
+            Hindsight's initialized health payload → treat as a stale daemon (or
+            foreign process) and attempt to reclaim by killing the PID listening
+            on the port. This preserves the original intent of clearing stale
+            daemons from version upgrades.
           * Kill failed, or non-hindsight process occupying the port → False.
         """
         if not self._is_port_in_use(port):
             return True
 
-        # Port is occupied — check if it's a healthy hindsight daemon.
-        health_ok = False
-        try:
-            with httpx.Client(timeout=2) as client:
-                response = client.get(f"http://127.0.0.1:{port}/health")
-                health_ok = response.status_code == 200
-        except Exception:
-            health_ok = False
-
-        if health_ok:
+        # Port is occupied — check if it's a healthy Hindsight daemon. In
+        # concurrent startup races the socket can bind before /health returns
+        # 200, so wait briefly before deciding the listener is stale/foreign.
+        if self._wait_for_port_health(port):
             logger.debug(f"Port {port} already serving a healthy hindsight daemon; reusing it")
             return True
 
-        # Unhealthy — attempt to reclaim by killing the listener.
+        # Unhealthy after grace window — attempt to reclaim by killing the listener.
+        if not self._is_port_in_use(port):
+            return True
+
         pid = self._find_pid_on_port(port)
         if pid is None:
             logger.warning(f"Port {port} is in use by another process")

hindsight-embed/tests/test_daemon_client.py

@@ -11,6 +11,7 @@
 from hindsight_embed import daemon_client
 from hindsight_embed.daemon_embed_manager import DaemonEmbedManager
 
+
 @pytest.fixture
 def config():
     """Default config for tests."""
@@ -207,7 +208,10 @@ def test_port_occupied_by_healthy_hindsight_is_reused(self):
             mock_client = MagicMock()
             mock_client.__enter__ = Mock(return_value=mock_client)
             mock_client.__exit__ = Mock(return_value=False)
-            mock_client.get.return_value = Mock(status_code=200)
+            mock_client.get.return_value = Mock(
+                status_code=200,
+                json=Mock(return_value={"status": "healthy", "database": "connected"}),
+            )
             mock_httpx_cls.return_value = mock_client
 
             assert manager._clear_port(9555) is True
@@ -220,6 +224,7 @@ def test_port_occupied_by_non_hindsight_returns_false(self):
         with (
             patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True),
             patch("httpx.Client") as mock_httpx_cls,
+            patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0),
         ):
             mock_client = MagicMock()
             mock_client.__enter__ = Mock(return_value=mock_client)
@@ -235,6 +240,7 @@ def test_port_occupied_health_non_200_returns_false(self):
         with (
             patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True),
             patch("httpx.Client") as mock_httpx_cls,
+            patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0),
         ):
             mock_client = MagicMock()
             mock_client.__enter__ = Mock(return_value=mock_client)
@@ -251,6 +257,7 @@ def test_unhealthy_daemon_pid_not_found_returns_false(self):
             patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True),
             patch("httpx.Client") as mock_httpx_cls,
             patch.object(DaemonEmbedManager, "_find_pid_on_port", return_value=None),
+            patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0),
         ):
             mock_client = MagicMock()
             mock_client.__enter__ = Mock(return_value=mock_client)
@@ -268,6 +275,7 @@ def test_unhealthy_daemon_kill_fails_returns_false(self):
             patch("httpx.Client") as mock_httpx_cls,
             patch.object(DaemonEmbedManager, "_find_pid_on_port", return_value=12345),
             patch.object(DaemonEmbedManager, "_kill_process", return_value=False),
+            patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0),
         ):
             mock_client = MagicMock()
             mock_client.__enter__ = Mock(return_value=mock_client)
@@ -285,6 +293,75 @@ def test_unhealthy_daemon_kill_succeeds_returns_true(self):
             patch("httpx.Client") as mock_httpx_cls,
             patch.object(DaemonEmbedManager, "_find_pid_on_port", return_value=12345),
             patch.object(DaemonEmbedManager, "_kill_process", return_value=True),
+            patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0),
+        ):
+            mock_client = MagicMock()
+            mock_client.__enter__ = Mock(return_value=mock_client)
+            mock_client.__exit__ = Mock(return_value=False)
+            mock_client.get.return_value = Mock(status_code=503)
+            mock_httpx_cls.return_value = mock_client
+
+            assert manager._clear_port(9555) is True
+
+    def test_port_occupied_by_warming_hindsight_is_reused(self):
+        """Port bound before /health is ready — wait briefly and reuse when healthy."""
+        manager = DaemonEmbedManager()
+        with (
+            patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True),
+            patch("httpx.Client") as mock_httpx_cls,
+            patch.object(DaemonEmbedManager, "_find_pid_on_port") as mock_find_pid,
+            patch.object(DaemonEmbedManager, "_kill_process") as mock_kill,
+            patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 1.0),
+            patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_CHECK_INTERVAL", 0.01),
+        ):
+            mock_client = MagicMock()
+            mock_client.__enter__ = Mock(return_value=mock_client)
+            mock_client.__exit__ = Mock(return_value=False)
+            mock_client.get.side_effect = [
+                Mock(status_code=503),
+                Mock(status_code=503),
+                Mock(
+                    status_code=200,
+                    json=Mock(return_value={"status": "healthy", "database": "connected"}),
+                ),
+            ]
+            mock_httpx_cls.return_value = mock_client
+
+            assert manager._clear_port(9555) is True
+            assert mock_client.get.call_count == 3
+            mock_find_pid.assert_not_called()
+            mock_kill.assert_not_called()
+
+    def test_port_occupied_by_foreign_health_200_returns_false(self):
+        """HTTP 200 alone is not enough to identify the listener as Hindsight."""
+        manager = DaemonEmbedManager()
+        with (
+            patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True),
+            patch("httpx.Client") as mock_httpx_cls,
+            patch.object(DaemonEmbedManager, "_find_pid_on_port", return_value=None) as mock_find_pid,
+            patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 0.0),
+        ):
+            mock_client = MagicMock()
+            mock_client.__enter__ = Mock(return_value=mock_client)
+            mock_client.__exit__ = Mock(return_value=False)
+            mock_client.get.return_value = Mock(
+                status_code=200,
+                json=Mock(return_value={"status": "ok"}),
+            )
+            mock_httpx_cls.return_value = mock_client
+
+            assert manager._clear_port(9555) is False
+            mock_find_pid.assert_called_once_with(9555)
+
+    def test_port_cleared_during_grace_returns_true(self):
+        """If a stale listener exits during the grace wait, the port is already clear."""
+        manager = DaemonEmbedManager()
+        with (
+            patch.object(DaemonEmbedManager, "_is_port_in_use", side_effect=[True, False, False]),
+            patch("httpx.Client") as mock_httpx_cls,
+            patch.object(DaemonEmbedManager, "_find_pid_on_port") as mock_find_pid,
+            patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_GRACE_TIMEOUT", 1.0),
+            patch("hindsight_embed.daemon_embed_manager.PORT_HEALTH_CHECK_INTERVAL", 0.01),
         ):
             mock_client = MagicMock()
             mock_client.__enter__ = Mock(return_value=mock_client)
@@ -293,6 +370,19 @@ def test_unhealthy_daemon_kill_succeeds_returns_true(self):
             mock_httpx_cls.return_value = mock_client
 
             assert manager._clear_port(9555) is True
+            mock_find_pid.assert_not_called()
+
+    def test_invalid_port_health_timeout_is_bounded(self):
+        """Invalid grace timeout values must not create an unbounded wait."""
+        manager = DaemonEmbedManager()
+        with (
+            patch.object(DaemonEmbedManager, "_is_port_in_use", return_value=True),
+            patch.object(DaemonEmbedManager, "_port_health_ok", return_value=False) as mock_health,
+            patch("hindsight_embed.daemon_embed_manager.time.sleep") as mock_sleep,
+        ):
+            assert manager._wait_for_port_health(9555, timeout=float("nan")) is False
+            mock_health.assert_called_once_with(9555)
+            mock_sleep.assert_not_called()
 
 
 class TestStartDaemonSerialization: