diff --git a/gateway/run.py b/gateway/run.py index d9edbae96022a..e5483319e831e 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -7838,7 +7838,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew return self.adapters.get(Platform.RELAY) async def _scale_to_zero_watcher(self, interval: float = 30.0) -> None: - """Watch for idle and drive the relay dormant so the platform can suspend. + """Watch for idle, drive the relay dormant, then self-suspend the machine. Started ONLY when _scale_to_zero_should_arm() (opted in via the Labs HERMES_SCALE_TO_ZERO stamp + relay-only/absent messaging + a wakeUrl). @@ -7847,12 +7847,23 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew machine, §3.4(6); does NOT set _running=False), - relay adapter.go_dormant() — going_idle->ack + supervisor-preserving socket close (NOT disconnect(), NOT the run.py stop path), - - deliberately NO mark_resume_pending (D13 — suspend preserves RAM). - The process stays alive; the platform (Fly autostop:"suspend") suspends - the now-traffic-idle machine and autostart wakes it on the wakeUrl poke, - at which point the preserved reconnect supervisor re-dials and the + - deliberately NO mark_resume_pending (D13 — suspend preserves RAM), + - THEN suspend this machine through the local flaps socket + (gateway.scale_to_zero.suspend_self). The gateway owns the suspend + because Fly Proxy autostop judges idle on INBOUND connections only: + it cannot see an in-flight agent turn (outbound-only LLM traffic) + and, since the mid-2026 proxy change, an open outbound relay socket + no longer holds the machine awake — autostop:"suspend" would freeze + the machine mid-job or before the relay flip (the buffered-event + black hole). NAS therefore provisions scale-to-zero machines with + autostop:"off"; the suspend only ever happens HERE, strictly after + the idle predicate held and the dormant quiesce completed. + Autostart stays platform-side: the connector's wakeUrl poke (Fly-proxied) + wakes the machine, the preserved reconnect supervisor re-dials, and the connector drains the buffered backlog. After driving dormant we set a re-arm cooldown so a wake's drained backlog isn't immediately re-quiesced. + Off-Fly (no flaps socket / machine identity) the suspend step is skipped: + dormancy still happens, the process just stays running — fail-awake. """ await asyncio.sleep(min(interval, 30.0)) # let startup settle while self._running: @@ -7872,28 +7883,68 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew continue logger.info( "scale-to-zero: gateway idle for >= %.0fs — going dormant " - "(relay buffered, socket closed, awaiting platform suspend)", + "(relay buffered, socket closed) then self-suspending", self._scale_to_zero_idle_timeout_seconds(), ) try: self._update_runtime_status("draining") except Exception: # noqa: BLE001 - status is best-effort logger.debug("scale-to-zero: status mark failed", exc_info=True) + dormant_ok = True try: result = go_dormant() if asyncio.iscoroutine(result): await result except Exception: # noqa: BLE001 - dormancy is best-effort + dormant_ok = False logger.debug("scale-to-zero: go_dormant failed", exc_info=True) # 0.F: after a wake the drained inbound updates _last_inbound_at, # but give it a window so we don't immediately re-go-dormant on the # same idle reading before traffic lands. self._scale_to_zero_cooldown_until = time.time() + max(interval, 60.0) + # Self-suspend ONLY after a clean quiesce: the relay flip must be + # set (buffered delivery + wake poke armed) before the freeze, or + # inbound events black-hole while we sleep. Re-check idle one last + # time — inbound may have landed during the quiesce await. + if not dormant_ok: + continue + if not self._scale_to_zero_is_idle(): + logger.info( + "scale-to-zero: inbound arrived during quiesce — skipping suspend" + ) + continue + await self._scale_to_zero_self_suspend() except asyncio.CancelledError: raise except Exception: # noqa: BLE001 - the watcher must never crash the gateway logger.debug("scale-to-zero watcher iteration error", exc_info=True) + async def _scale_to_zero_self_suspend(self) -> None: + """Suspend this Fly machine via the local flaps socket (fail-awake). + + Runs the blocking unix-socket call in a worker thread so the event loop + stays live right up to the kernel freeze. On success the process is + frozen shortly after — nothing meaningful runs until the wake resume. + Off-Fly (self_suspend_available() False) this is a silent no-op. + """ + from gateway.scale_to_zero import self_suspend_available, suspend_self + + try: + if not self_suspend_available(): + logger.debug( + "scale-to-zero: flaps socket / machine identity absent — " + "dormant without platform suspend" + ) + return + accepted = await asyncio.to_thread(suspend_self) + if not accepted: + logger.warning( + "scale-to-zero: self-suspend not accepted — machine stays " + "awake (fail-awake); will retry on the next idle window" + ) + except Exception: # noqa: BLE001 - suspend is best-effort, never crash + logger.debug("scale-to-zero: self-suspend failed", exc_info=True) + def _status_action_label(self) -> str: return "restart" if self._restart_requested else "shutdown" @@ -11895,7 +11946,9 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew # relay-only/absent, and a wakeUrl is registered (decisions.md D1/D11/ # §3.4(1)). A non-opted instance never starts it, so behaviour is exactly # as today. When armed, the watcher drives the relay dormant on sustained - # idle so the platform (Fly autostop:"suspend") can suspend the machine. + # idle and then suspends the machine itself via the local flaps socket + # (Fly Proxy autostop is inbound-only and job-blind, so the gateway owns + # the suspend decision; NAS provisions these machines autostop:"off"). try: if self._scale_to_zero_should_arm(): logger.info( diff --git a/gateway/scale_to_zero.py b/gateway/scale_to_zero.py index 107c45d764e52..60b68e9d0191d 100644 --- a/gateway/scale_to_zero.py +++ b/gateway/scale_to_zero.py @@ -3,9 +3,19 @@ This is the gateway-side BEHAVIOUR layer that consumes the relay scale-to-zero PRIMITIVES (gateway-gateway Phase 5: the buffered-flip, the durable per-instance buffer, the wakeUrl poke, the reconnect supervisor). It owns the *decision* to go -idle and drives the relay transport's ``go_dormant()`` (D12) — it does NOT itself -suspend the machine. On Fly, the now-traffic-idle machine is suspended by -``autostop:"suspend"`` and woken by autostart-on-wakeUrl (decisions.md Q3=C′). +idle, drives the relay transport's ``go_dormant()`` (D12), and then SUSPENDS the +machine itself through the local Fly Machines API socket. Wake stays platform-side: +autostart-on-wakeUrl (decisions.md Q3=C′). + +Why the gateway self-suspends instead of relying on ``autostop:"suspend"``: Fly +Proxy judges idle exclusively on INBOUND proxied connections — it cannot see an +in-flight agent turn (outbound-only LLM traffic) and there is no way for the app +to signal "not ready to suspend". Mid-2026 the proxy also stopped treating open +OUTBOUND sockets as activity, so the relay WebSocket no longer masks the race: +Fly would suspend a machine mid-job, and could suspend BEFORE ``go_dormant()`` +flipped the relay destination (the buffered-event black hole). Owning the suspend +call closes both: it only ever fires after the idle predicate (no running agents, +no live background work, inbound-quiet) holds AND the dormant quiesce completed. Design constraints (decisions.md): - Per-instance enable is gated SOLELY by the NAS "Labs" toggle, carried to the @@ -28,13 +38,30 @@ inputs so they unit-test without a live gateway. from __future__ import annotations +import json +import logging import os +import socket from typing import Any, Iterable, Optional +logger = logging.getLogger(__name__) + # Env flag stamped by NAS when the scaleToZero Labs toggle is on (D11/Q8=A), # mirroring how the `relay` feature stamps GATEWAY_RELAY_URL. Truthy values only. SCALE_TO_ZERO_ENV = "HERMES_SCALE_TO_ZERO" +# Fly-injected machine identity (present on every Fly machine). Used by the +# self-suspend call; both must be present for self_suspend_available(). +FLY_APP_NAME_ENV = "FLY_APP_NAME" +FLY_MACHINE_ID_ENV = "FLY_MACHINE_ID" + +# The local flaps (Fly Machines API) unix socket, available inside every Fly +# machine. A POST to /v1/apps/{app}/machines/{id}/suspend snapshots RAM and +# suspends THIS machine — the Fly-endorsed replacement for proxy autostop when +# the app must own the idle decision (https://fly.io/docs/reference/suspend-resume/). +FLY_API_SOCKET = "/.fly/api" + + # config.yaml default (D2). Behavioural setting -> config, not env. DEFAULT_IDLE_TIMEOUT_MINUTES = 5 @@ -122,3 +149,84 @@ def is_idle( if has_live_background_work: return False return seconds_since_last_inbound >= idle_timeout_seconds + + +def self_suspend_available(environ: Optional[dict] = None) -> bool: + """Whether this process can suspend its own machine via the flaps socket. + + True iff the Fly-injected machine identity is present AND the local Machines + API socket exists. Off-Fly (local dev, Azure ACA, tests) this is False and + the watcher simply skips the suspend step — dormancy still happens, the + platform just never freezes the process. + """ + env = environ if environ is not None else os.environ + return bool( + str(env.get(FLY_APP_NAME_ENV, "")).strip() + and str(env.get(FLY_MACHINE_ID_ENV, "")).strip() + and os.path.exists(FLY_API_SOCKET) + ) + + +def suspend_self( + environ: Optional[dict] = None, + *, + socket_path: str = FLY_API_SOCKET, + timeout: float = 10.0, +) -> bool: + """POST /v1/apps/{app}/machines/{id}/suspend on the local flaps socket. + + Fly's in-machine Machines API needs no token — the socket itself is the + credential. Equivalent to: + curl --unix-socket /.fly/api -X POST \\ + http://flaps/v1/apps/$FLY_APP_NAME/machines/$FLY_MACHINE_ID/suspend + + Returns True when flaps accepted the request (2xx). The caller should treat + this as fire-and-forget: on success the kernel freezes this process shortly + after, so there may be nothing meaningful to run afterwards. Never raises — + a failed suspend just leaves the machine running (fail-awake, never + fail-frozen), which costs money but loses no work. + + stdlib-only on purpose: a plain unix-socket HTTP/1.1 request, no httpx/ + requests dependency in the hot path and no async plumbing to freeze + mid-await. + """ + env = environ if environ is not None else os.environ + app = str(env.get(FLY_APP_NAME_ENV, "")).strip() + machine_id = str(env.get(FLY_MACHINE_ID_ENV, "")).strip() + if not app or not machine_id: + logger.warning("scale-to-zero: suspend_self called without Fly machine identity") + return False + request = ( + f"POST /v1/apps/{app}/machines/{machine_id}/suspend HTTP/1.1\r\n" + "Host: flaps\r\n" + "Content-Length: 0\r\n" + "Connection: close\r\n" + "\r\n" + ) + try: + with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as sock: + sock.settimeout(timeout) + sock.connect(socket_path) + sock.sendall(request.encode("ascii")) + response = b"" + while len(response) < 65536: + chunk = sock.recv(4096) + if not chunk: + break + response += chunk + except OSError as exc: + logger.warning("scale-to-zero: flaps suspend request failed: %s", exc) + return False + status_line = response.split(b"\r\n", 1)[0].decode("ascii", "replace") + parts = status_line.split() + ok = len(parts) >= 2 and parts[1].isdigit() and 200 <= int(parts[1]) < 300 + if ok: + logger.info("scale-to-zero: machine suspend accepted by flaps (%s)", status_line) + else: + body = response.split(b"\r\n\r\n", 1)[-1][:500].decode("utf-8", "replace") + logger.warning( + "scale-to-zero: flaps suspend rejected: %s %s", + status_line, + json.dumps(body)[:500], + ) + return ok diff --git a/tests/gateway/test_scale_to_zero.py b/tests/gateway/test_scale_to_zero.py index 9935b23df3b4d..570d4f5f17b35 100644 --- a/tests/gateway/test_scale_to_zero.py +++ b/tests/gateway/test_scale_to_zero.py @@ -89,3 +89,92 @@ def test_idle_exactly_at_threshold(): assert is_idle(**_idle_kwargs(seconds_since_last_inbound=300.0)) is True + + +# ── suspend_self / self_suspend_available (the gateway-owned suspend call) ─── +# +# Fly Proxy autostop is inbound-only and job-blind (and since mid-2026 no longer +# counts outbound sockets as activity), so the gateway suspends its own machine +# via the local flaps unix socket strictly after the idle predicate + dormant +# quiesce. These exercise the wire call against a real unix-socket fake flaps. + + +import os +import socket as _socket +import threading + + +from gateway.scale_to_zero import ( # noqa: E402 - grouped with their section + FLY_APP_NAME_ENV, + FLY_MACHINE_ID_ENV, + self_suspend_available, + suspend_self, +) + +_FLY_ENV = {FLY_APP_NAME_ENV: "hermes-agent-stg-test", FLY_MACHINE_ID_ENV: "d891234f"} + + +def _fake_flaps(tmp_path, status_line, capture): + """One-shot unix-socket HTTP server standing in for flaps.""" + sock_path = str(tmp_path / "fly-api.sock") + server = _socket.socket(_socket.AF_UNIX, _socket.SOCK_STREAM) + server.bind(sock_path) + server.listen(1) + + def serve(): + conn, _ = server.accept() + with conn: + conn.settimeout(5) + data = b"" + while b"\r\n\r\n" not in data: + chunk = conn.recv(4096) + if not chunk: + break + data += chunk + capture.append(data) + conn.sendall( + f"HTTP/1.1 {status_line}\r\nContent-Length: 2\r\nConnection: close\r\n\r\n{{}}".encode() + ) + server.close() + + t = threading.Thread(target=serve, daemon=True) + t.start() + return sock_path, t + + +def test_suspend_self_posts_suspend_for_this_machine(tmp_path): + captured: list[bytes] = [] + sock_path, t = _fake_flaps(tmp_path, "200 OK", captured) + assert suspend_self(_FLY_ENV, socket_path=sock_path) is True + t.join(timeout=5) + request = captured[0].decode() + # The request must target THIS machine's suspend endpoint, per the Fly + # Machines API (POST /v1/apps/{app}/machines/{id}/suspend on /.fly/api). + assert request.startswith( + "POST /v1/apps/hermes-agent-stg-test/machines/d891234f/suspend HTTP/1.1\r\n" + ) + assert "Host: flaps\r\n" in request + + +def test_suspend_self_non_2xx_is_false_not_raise(tmp_path): + captured: list[bytes] = [] + sock_path, t = _fake_flaps(tmp_path, "412 Precondition Failed", captured) + assert suspend_self(_FLY_ENV, socket_path=sock_path) is False + t.join(timeout=5) + + +def test_suspend_self_missing_socket_is_false_not_raise(tmp_path): + # Fail-awake: a dead/absent flaps socket must never raise out of the watcher. + assert suspend_self(_FLY_ENV, socket_path=str(tmp_path / "nope.sock")) is False + + +def test_suspend_self_requires_machine_identity(tmp_path): + assert suspend_self({}, socket_path=str(tmp_path / "unused.sock")) is False + + +def test_self_suspend_available_needs_identity_and_socket(): + # No socket at /.fly/api in a test environment -> unavailable even with env. + if not os.path.exists("/.fly/api"): + assert self_suspend_available(_FLY_ENV) is False + # Missing identity -> unavailable regardless of socket. + assert self_suspend_available({}) is False diff --git a/tests/gateway/test_scale_to_zero_watcher.py b/tests/gateway/test_scale_to_zero_watcher.py index fa42e0a242966..6cb4ad5f15328 100644 --- a/tests/gateway/test_scale_to_zero_watcher.py +++ b/tests/gateway/test_scale_to_zero_watcher.py @@ -162,3 +162,89 @@ def test_no_arm_when_a_direct_platform_is_actually_enabled(monkeypatch): assert r._scale_to_zero_should_arm() is False + +# ── the self-suspend step: fires only after a clean quiesce, in order ───────── +# +# The gateway owns the suspend (Fly Proxy autostop is inbound-only/job-blind and +# no longer held open by outbound sockets), so the watcher must (a) suspend only +# AFTER go_dormant succeeded — the relay flip precedes the freeze, closing the +# buffered-event black hole — and (b) never suspend when the quiesce failed or +# inbound landed mid-quiesce. + + +@pytest.mark.asyncio +async def test_watcher_self_suspends_after_dormant(monkeypatch): + r, adapter = _runner_with(monkeypatch, idle=True) + calls = [] + + async def fake_suspend(): + calls.append(("suspend", adapter.go_dormant_calls)) + r._running = False # stop the loop after the first full sequence + + monkeypatch.setattr(r, "_scale_to_zero_self_suspend", fake_suspend, raising=False) + task = asyncio.create_task(r._scale_to_zero_watcher(interval=0.01)) + await asyncio.wait_for(task, timeout=2) + # Suspend fired exactly once, and only AFTER go_dormant ran (flip-before-freeze). + assert calls == [("suspend", 1)] + + +@pytest.mark.asyncio +async def test_watcher_skips_suspend_when_dormant_fails(monkeypatch): + r, adapter = _runner_with(monkeypatch, idle=True) + + async def broken_dormant(): + raise RuntimeError("quiesce failed") + + adapter.go_dormant = broken_dormant + suspend_calls = [] + + async def fake_suspend(): + suspend_calls.append(1) + + monkeypatch.setattr(r, "_scale_to_zero_self_suspend", fake_suspend, raising=False) + task = asyncio.create_task(r._scale_to_zero_watcher(interval=0.01)) + await asyncio.sleep(0.1) + r._running = False + await asyncio.wait_for(task, timeout=2) + # A failed quiesce means an UNFLIPPED relay — suspending would black-hole + # inbound events. Must stay awake. + assert suspend_calls == [] + + +@pytest.mark.asyncio +async def test_watcher_skips_suspend_when_inbound_lands_mid_quiesce(monkeypatch): + r, adapter = _runner_with(monkeypatch, idle=True) + # First idle check (loop gate) True, second (post-quiesce re-check) False. + reads = iter([True, False, False, False, False, False]) + monkeypatch.setattr( + r, "_scale_to_zero_is_idle", lambda: next(reads, False), raising=False + ) + suspend_calls = [] + + async def fake_suspend(): + suspend_calls.append(1) + + monkeypatch.setattr(r, "_scale_to_zero_self_suspend", fake_suspend, raising=False) + task = asyncio.create_task(r._scale_to_zero_watcher(interval=0.01)) + await asyncio.sleep(0.15) + r._running = False + await asyncio.wait_for(task, timeout=2) + assert adapter.go_dormant_calls == 1 + assert suspend_calls == [] + + +@pytest.mark.asyncio +async def test_self_suspend_noop_off_fly(monkeypatch): + """Off-Fly (no flaps socket/identity) the helper is a silent no-op — + dormancy without platform suspend, never an error.""" + r = GatewayRunner.__new__(GatewayRunner) + monkeypatch.setattr( + "gateway.scale_to_zero.self_suspend_available", lambda *a, **k: False + ) + called = [] + monkeypatch.setattr( + "gateway.scale_to_zero.suspend_self", + lambda *a, **k: called.append(1) or True, + ) + await r._scale_to_zero_self_suspend() + assert called == []