Production entrypoint, health/drain, and cloud image for the CDP proxy (SKY-12533) (#7546)

Author: AronPerez

skyvern/proxy/adapters/websocket_server.py

@@ -16,6 +16,7 @@
 import json
 import logging
 import os
+import signal
 import time
 from dataclasses import dataclass, field
 from http import HTTPStatus
@@ -123,6 +124,15 @@
 # so /json/version and /json/list match before the bare /json alias.
 _DISCOVERY_SUFFIXES = ("/json/version", "/json/list", "/json")
 
+# Unauthenticated HTTP health target for orchestrator probes (httpGet, not
+# tcpSocket: a wedged event loop still accepts TCP via the listen backlog).
+_HEALTH_PATH = "/healthz"
+# After SIGTERM/SIGINT the listener closes and existing relays drain. Past this
+# budget remaining connections are force-closed (1001) so the process exits with
+# clean WS closes instead of the supervisor's SIGKILL cutting sockets mid-frame.
+_DEFAULT_DRAIN_TIMEOUT_SECONDS = 3600
+_MAX_DRAIN_TIMEOUT_SECONDS = 24 * 3600
+
 # Enqueued by the drain path so the writer flushes all real frames then stops.
 _DRAIN_SENTINEL = object()
 
@@ -814,15 +824,67 @@ def __init__(
         self._shared_upstreams_lock = asyncio.Lock()
 
     async def serve_forever(self) -> None:
-        async with websockets.serve(
+        loop = asyncio.get_running_loop()
+        stop = asyncio.Event()
+        installed: list[signal.Signals] = []
+        # Handlers go in BEFORE the listener opens so no signal can land in a
+        # window where connections are being accepted but drain is not armed.
+        for sig in (signal.SIGTERM, signal.SIGINT):
+            try:
+                loop.add_signal_handler(sig, stop.set)
+                installed.append(sig)
+            except (NotImplementedError, RuntimeError, ValueError):
+                # Platforms/embeddings without loop signal support (Windows,
+                # non-main threads) still serve; drain then relies on cancellation.
+                pass
+        server = await websockets.serve(
             self._handle_client, self._host, self._port, max_size=None, process_request=self._process_request
-        ):
-            LOG.info("CDP proxy listening", host=self._host, port=self._port)
-            await asyncio.Future()
+        )
+        LOG.info("CDP proxy listening", host=self._host, port=self._port)
+        try:
+            await stop.wait()
+            # Drain: close the listener so no new client lands here, keep every
+            # established relay running, and only force-close stragglers once the
+            # drain budget is spent. The orchestrator's grace period is the outer
+            # bound; this inner one exists so a healthy drain ends with clean
+            # 1001 closes rather than a SIGKILL.
+            drain_timeout = _positive_env_int(
+                "CDP_PROXY_DRAIN_TIMEOUT_SECONDS", _DEFAULT_DRAIN_TIMEOUT_SECONDS, _MAX_DRAIN_TIMEOUT_SECONDS
+            )
+            LOG.info("CDP proxy draining", drain_timeout_seconds=drain_timeout)
+            server.close(close_connections=False)
+            try:
+                await asyncio.wait_for(server.wait_closed(), timeout=drain_timeout)
+            except asyncio.TimeoutError:
+                # Server.close() is idempotent — a second call cannot upgrade the
+                # drain to a force-close — so stragglers are closed directly. Each
+                # close() bounds its own handshake, so this always terminates.
+                # `handlers` is an undocumented websockets.asyncio Server attribute
+                # (its _close() iterates the same dict); the lifecycle tests here
+                # exercise it, so a websockets bump that renames it fails loudly.
+                LOG.warning("CDP proxy drain timeout expired, force-closing remaining connections")
+                stragglers = [asyncio.create_task(conn.close(1001)) for conn in list(server.handlers)]
+                if stragglers:
+                    await asyncio.wait(stragglers)
+        finally:
+            for sig in installed:
+                loop.remove_signal_handler(sig)
+            server.close(close_connections=True)
+            await server.wait_closed()
+            LOG.info("CDP proxy stopped")
 
     async def _process_request(self, connection: websockets.ServerConnection, request: Request) -> Response | None:
         """Serve the pre-upgrade CDP discovery GETs; return None to let a real WS
         upgrade proceed. Authenticated the same way as the WS connection."""
+        try:
+            is_health_probe = urlsplit(request.path).path == _HEALTH_PATH
+        except ValueError:
+            is_health_probe = False
+        if is_health_probe:
+            # Unauthenticated on purpose: probes carry no credentials, and the
+            # response reveals nothing but liveness. Draining needs no branch —
+            # the closed listener refuses the probe's connect outright.
+            return connection.respond(HTTPStatus.OK, "ok")
         try:
             split = _split_discovery_target(request.path)
             if split is None:

tests/unit/proxy/test_discovery.py

@@ -235,3 +235,29 @@ async def test_discovery_never_resolves_or_leaks_the_upstream() -> None:
     assert sessions.resolve_calls == 0
     text = bytes(response.body).decode()
     assert SECRET not in text and "vendor.internal" not in text
+
+
+# ---- health endpoint --------------------------------------------------------
+
+
+@pytest.mark.asyncio
+async def test_healthz_answers_without_credentials() -> None:
+    server = _server(StaticKeyAuth({VALID_KEY: Principal(principal_id="p")}))
+
+    response = await _discover(server, "/healthz")
+
+    assert response is not None
+    assert response.status_code == 200
+    assert response.body == b"ok"
+
+
+@pytest.mark.asyncio
+async def test_healthz_never_reaches_discovery_or_ws_upgrade() -> None:
+    # A probe target is answered directly: not None (which would hand the GET to
+    # the WS handshake) and not gated on the discovery auth path.
+    server = _server(StaticKeyAuth({VALID_KEY: Principal(principal_id="p")}))
+
+    response = await _discover(server, "/healthz?probe=1")
+
+    assert response is not None
+    assert response.status_code == 200

tests/unit/proxy/test_websocket_server.py

@@ -1,9 +1,16 @@
 from __future__ import annotations
 
+import asyncio
+import contextlib
+import os
+import signal
+import socket
 from types import SimpleNamespace
 from typing import Mapping
 
 import pytest
+from websockets import exceptions as websockets_exceptions
+from websockets.asyncio import client as websockets_client
 from websockets.datastructures import Headers
 
 from skyvern.proxy.adapters.memory import (
@@ -365,3 +372,102 @@ async def test_session_is_resolved_once_per_connection_not_per_message() -> None
     await make_server(upstream, counting)._handle_client(ws)  # type: ignore[arg-type]
 
     assert counting.resolve_calls == 1
+
+
+# ---- serve_forever lifecycle: health endpoint + signal drain ----------------
+
+
+def _free_port() -> int:
+    with socket.socket() as sock:
+        sock.bind(("127.0.0.1", 0))
+        return int(sock.getsockname()[1])
+
+
+async def _wait_for_listener(port: int, *, up: bool, attempts: int = 500) -> None:
+    for _ in range(attempts):
+        try:
+            _, writer = await asyncio.open_connection("127.0.0.1", port)
+        except OSError:
+            if not up:
+                return
+        else:
+            writer.close()
+            if up:
+                return
+        await asyncio.sleep(0.01)
+    raise AssertionError(f"listener on port {port} never became {'reachable' if up else 'refused'}")
+
+
+async def _http_get(port: int, path: str) -> tuple[int, bytes]:
+    reader, writer = await asyncio.open_connection("127.0.0.1", port)
+    writer.write(f"GET {path} HTTP/1.1\r\nHost: localhost\r\nConnection: close\r\n\r\n".encode())
+    await writer.drain()
+    raw = await reader.read(-1)
+    writer.close()
+    head, _, body = raw.partition(b"\r\n\r\n")
+    return int(head.split(b" ", 2)[1]), body
+
+
+def _lifecycle_server(port: int) -> CdpProxyServer:
+    sessions = InMemorySessionRegistry()
+    sessions.put(make_resolved_session())
+    return CdpProxyServer(
+        upstream=RecordingUpstreamBrowser(),
+        sessions=sessions,
+        auth=AllowAllAuth(),
+        metrics=NoOpMetrics(),
+        event_policy=ForwardAllEventPolicy(),
+        host="127.0.0.1",
+        port=port,
+    )
+
+
+@pytest.mark.asyncio
+async def test_serve_forever_healthz_and_sigterm_drain() -> None:
+    port = _free_port()
+    task = asyncio.create_task(_lifecycle_server(port).serve_forever())
+    try:
+        await _wait_for_listener(port, up=True)
+        status, body = await _http_get(port, "/healthz")
+        assert (status, body) == (200, b"ok")
+
+        async with websockets_client.connect(f"ws://127.0.0.1:{port}/s1") as client:
+            await client.ping()
+            os.kill(os.getpid(), signal.SIGTERM)
+            # Drain: the listener refuses new work while the live relay survives.
+            await _wait_for_listener(port, up=False)
+            await client.ping()
+        await asyncio.wait_for(task, timeout=10)
+    finally:
+        if not task.done():
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task
+
+
+@pytest.mark.asyncio
+async def test_serve_forever_force_closes_stragglers_after_drain_timeout(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setenv("CDP_PROXY_DRAIN_TIMEOUT_SECONDS", "1")
+    port = _free_port()
+    task = asyncio.create_task(_lifecycle_server(port).serve_forever())
+    client = None
+    try:
+        await _wait_for_listener(port, up=True)
+        client = await websockets_client.connect(f"ws://127.0.0.1:{port}/s1")
+        os.kill(os.getpid(), signal.SIGTERM)
+
+        # The straggler never closes; the drain budget expires and the server
+        # exits anyway, closing the client with a clean 1001 instead of a RST.
+        await asyncio.wait_for(task, timeout=10)
+        with pytest.raises(websockets_exceptions.ConnectionClosed):
+            await asyncio.wait_for(client.recv(), timeout=5)
+        assert client.protocol.close_code == 1001
+    finally:
+        if client is not None:
+            await client.close()
+        if not task.done():
+            task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await task