diff --git a/gateway/control_socket.py b/gateway/control_socket.py index 383b3c8769..50c516f393 100644 --- a/gateway/control_socket.py +++ b/gateway/control_socket.py @@ -543,3 +543,18 @@ def _query_windows_pipe( def identify_gateway(home: Path, *, timeout: float = _DEFAULT_CLIENT_TIMEOUT) -> Optional[dict[str, Any]]: """Convenience wrapper: ``identify`` the gateway serving ``home``.""" return query_gateway_control(home, "identify", timeout=timeout) + + +def pause_gateway_for_update( + home: Path, *, timeout: float = _DEFAULT_CLIENT_TIMEOUT +) -> Optional[dict[str, Any]]: + """Ask the gateway serving ``home`` to drain and exit for an update. + + Step 2 of the socket migration (#92091). Returns the gateway's ACK — + ``{"pausing": bool, "already_stopping": bool, "pid": int, + "drain_timeout": float}`` — or None when no gateway answers (older + gateway without the verb, no socket, dead socket). None means the + caller falls back to the legacy pause path (signals / tree-kill), + exactly as before this verb existed. + """ + return query_gateway_control(home, "pause-for-update", timeout=timeout) diff --git a/gateway/run.py b/gateway/run.py index 915eee4d49..7e16191321 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -31437,7 +31437,47 @@ async def start_gateway(config: Optional[GatewayConfig] = None, replace: bool = try: from gateway.control_socket import GatewayControlServer - _control_server = GatewayControlServer() + # pause-for-update (#92091 step 2): the updater asks this gateway to + # drain in-flight turns and exit cleanly — releasing every venv file + # handle — instead of being tree-killed mid-turn. Same drain path as + # SIGUSR1/service restarts (request_restart(via_service=True)); the + # updater (or the service manager) relaunches after the code swap. + # The handler runs on the socket's executor thread, so the restart + # request is marshalled onto the loop thread; the ACK returns the + # drain budget so the caller knows how long to wait for exit. + _main_loop = asyncio.get_running_loop() + + def _pause_for_update_handler() -> dict: + try: + from hermes_cli.gateway import _get_restart_drain_timeout + + _drain = float(_get_restart_drain_timeout()) + except Exception: + _drain = 30.0 + accepted_box: list[bool] = [] + _done = threading.Event() + + def _request() -> None: + try: + accepted_box.append( + runner.request_restart(detached=False, via_service=True) + ) + finally: + _done.set() + + _main_loop.call_soon_threadsafe(_request) + _done.wait(timeout=5.0) + accepted = bool(accepted_box and accepted_box[0]) + return { + "pausing": accepted, + "already_stopping": not accepted, + "pid": os.getpid(), + "drain_timeout": _drain, + } + + _control_server = GatewayControlServer( + verb_handlers={"pause-for-update": _pause_for_update_handler} + ) if not await _control_server.start(): _control_server = None else: diff --git a/hermes_cli/update_cmd.py b/hermes_cli/update_cmd.py index fa527fe55a..3a62699bc3 100644 --- a/hermes_cli/update_cmd.py +++ b/hermes_cli/update_cmd.py @@ -5116,6 +5116,7 @@ def _pause_windows_gateways_for_update() -> dict | None: profiles: dict[str, int] = {} mapped_pids = [] + socket_acks: list[dict] = [] for pid in running_pids: proc = profile_processes.get(pid) if proc is None: @@ -5123,6 +5124,23 @@ def _pause_windows_gateways_for_update() -> dict | None: profiles[str(proc.profile)] = int(pid) mapped_pids.append(int(pid)) _write_update_planned_stop_marker(Path(proc.path), int(pid)) + # Socket-first pause (#92091 step 2): ask the gateway to drain and + # exit itself instead of relying on the marker poll + force-kill + # ladder. A positive ACK means the gateway is running its own + # graceful restart path (same drain as SIGUSR1/service restarts) and + # will release its venv handles on the way out. No answer (older + # gateway, no socket) → the marker watcher / force-kill fallback + # below behaves exactly as before this verb existed. + try: + from gateway.control_socket import pause_gateway_for_update + + ack = pause_gateway_for_update(Path(proc.path)) + if ack and (ack.get("pausing") or ack.get("already_stopping")): + socket_acks.append(ack) + except Exception as exc: + logger.debug( + "Socket pause unavailable for gateway %s: %s", pid, exc + ) # Resolve each mapped worker's venv-side launcher BEFORE draining: the # drain stops tracking a PID exactly when it dies, so a gracefully @@ -5143,6 +5161,23 @@ def _pause_windows_gateways_for_update() -> dict | None: drain_timeout = max(float(_get_restart_drain_timeout()), 1.0) except Exception: drain_timeout = 10.0 + if socket_acks: + # A socket-paused gateway drains its ACTIVE TURN before exiting; give + # it the budget it declared (plus teardown grace) rather than only + # the local default, so a mid-turn gateway isn't force-killed at the + # end of a too-short wait — the exact outcome the verb exists to + # prevent. + try: + declared = max( + float(a.get("drain_timeout") or 0.0) for a in socket_acks + ) + drain_timeout = max(drain_timeout, declared + 10.0) + except Exception: + pass + print( + f" → {len(socket_acks)} gateway(s) ACKed socket pause; " + f"waiting up to {int(drain_timeout)}s for graceful exit" + ) survivors = _m()._wait_for_windows_update_gateway_exit( mapped_pids, timeout=drain_timeout, diff --git a/tests/gateway/test_control_socket_pause.py b/tests/gateway/test_control_socket_pause.py new file mode 100644 index 0000000000..07998ea36a --- /dev/null +++ b/tests/gateway/test_control_socket_pause.py @@ -0,0 +1,109 @@ +"""pause-for-update control-socket verb (#92091 step 2, campaign #91277). + +The updater asks a running gateway to drain and exit cleanly (releasing its +venv file handles) instead of tree-killing it mid-turn. Fallback contract: +older gateways without the verb answer nothing, and callers keep the legacy +marker/force-kill path. +""" + +from __future__ import annotations + +import asyncio +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import pytest + +from gateway.control_socket import ( + GatewayControlServer, + pause_gateway_for_update, + query_gateway_control, +) + + +def _make_server(tmp_path, handler): + server = GatewayControlServer( + home=tmp_path, verb_handlers={"pause-for-update": handler} + ) + return server + + +def test_pause_verb_dispatches_and_returns_ack(tmp_path): + calls = [] + + def handler(): + calls.append(1) + return {"pausing": True, "already_stopping": False, "pid": 111, + "drain_timeout": 30.0} + + server = _make_server(tmp_path, handler) + raw = json.dumps({"verb": "pause-for-update", "id": 7}).encode() + response = json.loads(server.handle_request_line(raw).decode()) + assert response["ok"] is True + assert response["result"]["pausing"] is True + assert response["result"]["drain_timeout"] == 30.0 + assert response["id"] == 7 + assert calls == [1] + + +def test_unknown_verb_still_lists_pause(tmp_path): + server = _make_server(tmp_path, lambda: {}) + raw = json.dumps({"verb": "nope"}).encode() + response = json.loads(server.handle_request_line(raw).decode()) + assert response["ok"] is False + assert "pause-for-update" in response["supported_verbs"] + + +@pytest.mark.skipif(sys.platform == "win32", reason="unix socket transport") +def test_pause_client_roundtrip_over_real_socket(tmp_path): + """Full client→socket→handler→ACK path over a REAL unix socket.""" + + async def scenario(): + acks = [] + + def handler(): + acks.append(1) + return {"pausing": True, "already_stopping": False, + "pid": 4242, "drain_timeout": 12.5} + + server = _make_server(tmp_path, handler) + assert await server.start() + try: + loop = asyncio.get_running_loop() + result = await loop.run_in_executor( + None, lambda: pause_gateway_for_update(tmp_path, timeout=5.0) + ) + finally: + await server.stop() + return result, acks + + result, acks = asyncio.run(scenario()) + assert acks == [1] + assert result is not None + assert result["pausing"] is True and result["drain_timeout"] == 12.5 + + +@pytest.mark.skipif(sys.platform == "win32", reason="unix socket transport") +def test_pause_client_none_when_gateway_lacks_verb(tmp_path): + """Back-compat: a step-1 gateway (identify/status only) answers ok:false + for the unknown verb → the client returns None → caller keeps the legacy + kill path.""" + + async def scenario(): + server = GatewayControlServer(home=tmp_path) # no pause handler + assert await server.start() + try: + loop = asyncio.get_running_loop() + return await loop.run_in_executor( + None, lambda: pause_gateway_for_update(tmp_path, timeout=5.0) + ) + finally: + await server.stop() + + assert asyncio.run(scenario()) is None + + +def test_pause_client_none_when_no_socket(tmp_path): + assert pause_gateway_for_update(tmp_path, timeout=0.5) is None