improvement(interception): keep tunnels serving with prime-tunnel

pyproject.toml

@@ -32,7 +32,7 @@ dependencies = [
     "numpy>=2.1.0",
     "math-verify>=0.8.0",
     "openai>=2.54.0,<3.0.0",
-    "prime-tunnel>=0.1.14",
+    "prime-tunnel>=0.1.17",
     "prime-sandboxes>=0.5.4",
     "pydantic>=2.12.3",
     "rich>=11.0.0",

tests/unit/interception/test_tunnel.py

@@ -1,107 +1,171 @@
-"""The prime tunnel watchdog restarts frpc when the tunnel service has the tunnel
-disconnected or frpc has exited, leaves it alone otherwise, and stops watching a tunnel
-whose registration is gone."""
+"""A prime tunnel is kept serving by the SDK through the service's health route. A tunnel
+that can never work again retires its server: its rollouts fail and the pools skip it."""
 
 import asyncio
+import logging
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
 
+import httpx
 import pytest
-from prime_tunnel import TunnelGoneError, TunnelStatus
-
+from prime_tunnel import TunnelGoneError
+
+from verifiers.v1.configs.interception import (
+    InterceptionServerConfig,
+    StaticInterceptionPoolConfig,
+)
+from verifiers.v1.errors import TunnelError
+from verifiers.v1.interception.pool import (
+    ElasticInterceptionPool,
+    StaticInterceptionPool,
+)
+from verifiers.v1.interception.server import InterceptionServer
 from verifiers.v1.interception.tunnel import prime
-
-GONE = (None, "terminated", "expired")
+from verifiers.v1.interception.tunnel.base import HEALTH_PATH
 
 
 class FakeClient:
-    def __init__(self, statuses, running=True, restart_error=None):
-        self.statuses = list(statuses)
-        self.is_running = running
-        self.is_gone = False
-        self.tunnel_id = "t-0-test"
-        self.restart_error = restart_error
-        self.restarts = 0
-        self.checks = 0
-
-    async def restart(self):
-        self.restarts += 1
-        if self.restart_error is not None:
-            raise self.restart_error
-        self.is_running = True
-
-    async def status(self):
-        self.checks += 1
-        registration = self.statuses.pop(0) if self.statuses else "connected"
-        if isinstance(registration, Exception):
-            raise registration
-        self.is_gone = self.is_gone or registration in GONE
-        return TunnelStatus(
-            tunnel_id=self.tunnel_id,
-            running=self.is_running,
-            registration=registration,
-            gone=self.is_gone,
+    gone = False
+
+    def __init__(self, local_port):
+        self.health_paths = []
+        self.stops = 0
+        type(self).instance = self
+
+    async def start(self):
+        return "https://t.example/"
+
+    async def keep_serving(self, health_path):
+        self.health_paths.append(health_path)
+        if self.gone:
+            raise TunnelGoneError("inactive")
+        await asyncio.Future()
+
+    def sync_stop(self):
+        self.stops += 1
+
+
+@asynccontextmanager
+async def unlimited():
+    yield
+
+
+@pytest.fixture
+def tunnel(monkeypatch):
+    monkeypatch.setattr(prime, "ensure_prime_auth", lambda: None)
+    monkeypatch.setattr(prime, "tunnel_limiter", unlimited)
+    monkeypatch.setattr("prime_tunnel.Tunnel", FakeClient)
+    return prime.PrimeTunnel()
+
+
+async def test_held_tunnel_is_kept_serving_and_stopped_on_exit(tunnel):
+    async with tunnel.expose(1234) as url:
+        assert url == "https://t.example"
+        await asyncio.sleep(0)
+        assert FakeClient.instance.health_paths == [HEALTH_PATH]
+        assert FakeClient.instance.stops == 0
+    assert FakeClient.instance.stops == 1
+
+
+async def test_gone_tunnel_reports_once_and_is_still_stopped_on_exit(
+    tunnel, monkeypatch
+):
+    monkeypatch.setattr(FakeClient, "gone", True)
+    reports = []
+    tunnel.on_gone = lambda: reports.append(True)
+    async with tunnel.expose(1234):
+        await asyncio.sleep(0)
+        await asyncio.sleep(0)
+        assert reports == [True]
+    assert FakeClient.instance.stops == 1
+
+
+async def test_sdk_restarts_reach_the_run_logs(caplog):
+    with caplog.at_level(logging.WARNING, logger=prime.logger.name):
+        logging.getLogger("prime_tunnel.tunnel").warning(
+            "Tunnel %s: frpc restarted", "https://t"
         )
+    assert "tunnel https://t: frpc restarted" in caplog.messages
 
 
-async def watch(client, checks, monkeypatch):
-    """Run the watchdog for about `checks` checks; how many restarts it made."""
-    monkeypatch.setattr(prime, "CHECK_SECONDS", 0.01)
-    task = asyncio.create_task(prime._watch(client, "https://t"))
-    await asyncio.sleep(0.0125 * checks)
-    task.cancel()
-    await asyncio.wait({task})
-    return client.restarts
+async def test_server_answers_the_health_route():
+    async with InterceptionServer() as server, httpx.AsyncClient() as client:
+        response = await client.get(f"{server.base_url}{HEALTH_PATH}")
+    assert response.status_code == 200
 
 
-async def stops(client, monkeypatch):
-    """Run the watchdog until it gives the tunnel up."""
-    monkeypatch.setattr(prime, "CHECK_SECONDS", 0.01)
-    await asyncio.wait_for(prime._watch(client, "https://t"), 1)
+def test_gone_tunnel_fails_the_servers_rollouts_and_refuses_new_ones():
+    server = InterceptionServer()
+    released = []
+    session = SimpleNamespace(error=None, release=lambda: released.append(True))
+    server.sessions["secret"] = session
 
+    server._handle_tunnel_gone()
 
-async def test_connected_is_left_alone(monkeypatch):
-    assert await watch(FakeClient(["connected"] * 10), 8, monkeypatch) == 0
+    assert isinstance(session.error, TunnelError)
+    assert released == [True]
+    with pytest.raises(TunnelError):
+        server.register(SimpleNamespace())
 
 
-async def test_one_disconnected_check_is_not_enough(monkeypatch):
-    client = FakeClient(["disconnected", "connected"] * 5)
-    assert await watch(client, 8, monkeypatch) == 0
+async def test_elastic_pool_replaces_a_retired_server_and_frees_it():
+    async with ElasticInterceptionPool() as pool:
+        retired = await pool._warm_task
+        health = f"{retired.base_url}{HEALTH_PATH}"
+        retired._handle_tunnel_gone()
 
+        replacement = await pool._server()
 
-async def test_disconnected_twice_restarts_once(monkeypatch):
-    client = FakeClient(["disconnected", "disconnected"])
-    assert await watch(client, 8, monkeypatch) == 1
+        assert replacement is not retired
+        assert pool.servers == [replacement]
+        async with httpx.AsyncClient() as client:
+            with pytest.raises(httpx.ConnectError):
+                await client.get(health)
 
 
-async def test_exited_frpc_restarts_without_asking(monkeypatch):
-    client = FakeClient([RuntimeError("api down")] * 10, running=False)
-    assert await watch(client, 4, monkeypatch) == 1
+async def test_elastic_pool_keeps_a_retired_server_until_its_rollouts_leave():
+    async with ElasticInterceptionPool() as pool:
+        retired = await pool._warm_task
+        retired.sessions["secret"] = SimpleNamespace(error=None, release=lambda: None)
+        retired._handle_tunnel_gone()
 
+        replacement = await pool._server()
+        assert pool.servers == [retired, replacement]
 
-async def test_unreachable_service_is_no_verdict(monkeypatch):
-    client = FakeClient([RuntimeError("api down")] * 10)
-    assert await watch(client, 8, monkeypatch) == 0
+        retired.sessions.clear()
+        assert await pool._server() is replacement
+        assert pool.servers == [replacement]
 
 
-@pytest.mark.parametrize("registration", GONE)
-async def test_gone_registration_stops_watching(registration, monkeypatch):
-    client = FakeClient([registration])
-    await stops(client, monkeypatch)
-    assert client.restarts == 0
+async def test_cancelling_a_rollout_does_not_interrupt_freeing_a_retired_server():
+    stopping = asyncio.Event()
+    finish = asyncio.Event()
+    stopped = []
 
+    async def stop():
+        stopping.set()
+        await finish.wait()
+        stopped.append(True)
 
-async def test_gone_reported_by_frpc_stops_watching(monkeypatch):
-    client = FakeClient([])
-    client.is_gone = True
-    await stops(client, monkeypatch)
-    assert (client.checks, client.restarts) == (0, 0)
+    pool = ElasticInterceptionPool()
+    pool.servers = [SimpleNamespace(healthy=False, load=0, stop=stop)]
+    freeing = asyncio.create_task(pool._free_retired())
+    await asyncio.wait_for(stopping.wait(), 1)
 
+    freeing.cancel()
+    await asyncio.sleep(0)
+    finish.set()
 
-async def test_restart_refused_as_gone_stops_watching(monkeypatch):
-    client = FakeClient([], running=False, restart_error=TunnelGoneError("inactive"))
-    await stops(client, monkeypatch)
-    assert client.restarts == 1
+    with pytest.raises(asyncio.CancelledError):
+        await freeing
+    assert stopped == [True]
 
 
-async def test_failed_restart_is_tried_again(monkeypatch):
-    client = FakeClient([], running=False, restart_error=RuntimeError("frpc"))
-    assert await watch(client, 6, monkeypatch) >= 2
+async def test_static_pool_with_only_retired_servers_fails_the_rollout():
+    pool = StaticInterceptionPool(
+        StaticInterceptionPoolConfig(servers=[InterceptionServerConfig()])
+    )
+    pool.servers[0].healthy = False
+    with pytest.raises(TunnelError):
+        async with pool.acquire(SimpleNamespace()):
+            pass

verifiers/v1/interception/pool.py

@@ -24,9 +24,11 @@
     PrimeTunnelConfig,
     StaticInterceptionPoolConfig,
 )
+from verifiers.v1.errors import TunnelError
 from verifiers.v1.interception.base import Interception, Slot
 from verifiers.v1.interception.server import InterceptionServer
 from verifiers.v1.session import RolloutSession
+from verifiers.v1.utils.aio import run_shielded
 
 logger = logging.getLogger(__name__)
 
@@ -57,7 +59,10 @@ async def start(self) -> None:
     async def acquire(self, session: RolloutSession) -> AsyncIterator[Slot]:
         # server.acquire registers before its first yield, so concurrent acquires see the
         # updated load before choosing their own least-loaded server.
-        server = min(self.servers, key=lambda s: s.load)
+        healthy = [server for server in self.servers if server.healthy]
+        if not healthy:
+            raise TunnelError("interception tunnel lost")
+        server = min(healthy, key=lambda s: s.load)
         async with server.acquire(session) as slot:
             yield slot
 
@@ -91,12 +96,26 @@ async def stop(self) -> None:
                 await self._warm_task
         await super().stop()
 
+    async def _free_retired(self) -> None:
+        """Stop retired servers whose rollouts have left. A rollout still on one is using
+        its client, so that server waits for a later call. Shielded: an acquire cancelled
+        mid-stop would leave the listener open, with nothing left to close it."""
+        for server in [s for s in self.servers if not s.healthy and s.load == 0]:
+            self.servers.remove(server)
+            try:
+                await run_shielded(server.stop())
+            except Exception:
+                logger.warning(
+                    "interception pool: retired server cleanup failed", exc_info=True
+                )
+
     async def _server(self) -> InterceptionServer:
         """A server with spare capacity — reuse one under `multiplex`, else bring up a new
         one (its own tunnel, on `stack`, torn down with the pool). Acquires hold `_lock`;
         the warm task runs before they reach this path."""
+        await self._free_retired()
         for server in self.servers:
-            if server.load < self.config.multiplex:
+            if server.healthy and server.load < self.config.multiplex:
                 return server
         # Pin prime explicitly — the only tunnel kind that can be minted on demand.
         server = InterceptionServer(

verifiers/v1/interception/server.py

@@ -59,9 +59,11 @@
     ProviderError,
     RolloutError,
     TaskError,
+    TunnelError,
 )
 from verifiers.v1.interception.base import Interception, Slot
 from verifiers.v1.interception.tunnel import Tunnel, make_tunnel
+from verifiers.v1.interception.tunnel.base import HEALTH_PATH
 from verifiers.v1.semantic import ACPInfo, extract_acp_info
 from verifiers.v1.session import IdempotentRequest, ReplayResponse, RolloutSession
 from verifiers.v1.trace import Error, ModelCall, PolicyEvent, TimeSpan
@@ -361,6 +363,9 @@ def __init__(
         self.tunnel: Tunnel | None = (
             make_tunnel(self.config.tunnel) if requires_tunnel else None
         )
+        self.healthy = True
+        if self.tunnel is not None:
+            self.tunnel.on_gone = self._handle_tunnel_gone
         self.host = "127.0.0.1"
         self.port = 0
         self.base_url = ""  # set by `start`
@@ -371,6 +376,14 @@ def load(self) -> int:
         """Rollouts currently registered — what the pools balance on."""
         return len(self.sessions)
 
+    def _handle_tunnel_gone(self) -> None:
+        """Fail this server's rollouts and take no more: their URL will never work again."""
+        self.healthy = False
+        logger.error("interception unhealthy: url=%s", self.base_url)
+        for session in list(self.sessions.values()):
+            session.error = TunnelError("interception tunnel lost")
+            session.release()
+
     def _client(self, config: BaseClientConfig) -> Client:
         """The server-owned client for `config` — one per distinct endpoint config, shared
         by every session registered under it, so the rollouts this server multiplexes reuse
@@ -386,6 +399,8 @@ def _client(self, config: BaseClientConfig) -> Client:
     def register(self, session: RolloutSession) -> tuple[str, str]:
         """Register separate capabilities for model inference and private task state, and
         assign the session its server-owned model client."""
+        if not self.healthy:
+            raise TunnelError("interception tunnel lost")
         session.client = self._client(session.ctx.client)
         model_secret = secrets.token_urlsafe(16)
         state_secret = secrets.token_urlsafe(16)
@@ -544,8 +559,12 @@ async def handler(request: web.Request) -> web.Response:
 
         return handler
 
+    async def handle_health(self, request: web.Request) -> web.Response:
+        return web.Response(text="ok")
+
     async def start(self) -> None:
         app = web.Application(client_max_size=MAX_REQUEST_BODY)
+        app.router.add_get(HEALTH_PATH, self.handle_health)
         for dialect in DIALECTS:
             for route in dialect.routes:
                 app.router.add_post(route, self._handler_for(dialect))