From eafc0bc6eb1da4dc80460c4a26cbd9038b53d6d7 Mon Sep 17 00:00:00 2001 From: Xubin Ren <52506698+Re-bin@users.noreply.github.com> Date: Mon, 17 Aug 2026 11:53:24 +0800 Subject: [PATCH] fix(gateway): serialize shared runtime lifecycle --- nanobot/cli/gateway.py | 11 +-- nanobot/gateway/runtime.py | 149 ++++++++++++++++++++++------- tests/cli/test_commands.py | 2 + tests/cli/test_gateway_commands.py | 10 +- tests/cli/test_tui_launcher.py | 2 + tests/gateway/test_runtime.py | 141 ++++++++++++++++++++++++++- 6 files changed, 263 insertions(+), 52 deletions(-) diff --git a/nanobot/cli/gateway.py b/nanobot/cli/gateway.py index 5c00302ed..20e82d68f 100644 --- a/nanobot/cli/gateway.py +++ b/nanobot/cli/gateway.py @@ -14,7 +14,6 @@ from rich.console import Console from nanobot.config.schema import Config from nanobot.gateway import ( - GatewayClientLease, GatewayRuntime, GatewayRuntimePaths, GatewayStartOptions, @@ -209,18 +208,12 @@ def create_gateway_app( ) print_status(result.status) raise typer.Exit(1) - promoted = False - if result.ok or result.message == "gateway_already_running": - promoted = GatewayClientLease( - runtime, - kind="gateway-background", - ).mark_persistent() if result.ok: console.print("[green]Gateway started in the background.[/green]") print_status(runtime.status()) return if result.message == "gateway_already_running": - if promoted: + if result.promoted: console.print( "[green]Existing on-demand gateway promoted to persistent " "background mode.[/green]" @@ -290,8 +283,6 @@ def create_gateway_app( """Stop the background gateway.""" runtime = runtime_for_instance(workspace=workspace, config=config) result = runtime.stop(timeout_s=timeout) - if result.ok or result.message == "gateway_not_running": - GatewayClientLease(runtime, kind="gateway-stop").clear() if result.ok: console.print("[green]Gateway stopped.[/green]") else: diff --git a/nanobot/gateway/runtime.py b/nanobot/gateway/runtime.py index 7ed332e09..f376c7d36 100644 --- a/nanobot/gateway/runtime.py +++ b/nanobot/gateway/runtime.py @@ -1,5 +1,7 @@ """Gateway-specific configuration for the shared background process runtime.""" +# pyright: reportPrivateUsage=false + from __future__ import annotations import asyncio @@ -57,6 +59,7 @@ class RuntimeResult(ProcessResult): """Result of a gateway lifecycle operation.""" status: GatewayStatus + promoted: bool = False def build_gateway_command(python_executable: str, options: GatewayStartOptions) -> list[str]: @@ -133,8 +136,15 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]): def start_background(self, options: ProcessStartOptions) -> RuntimeResult: """Start the gateway detached from the current terminal.""" - with self._lifecycle_lock(): - return self._start_background(options) + lease = GatewayClientLease(self, kind="gateway-background") + while True: + lease.wait_for_shutdown() + with self._lifecycle_lock(): + promoted = lease._try_mark_persistent_locked() + if promoted is None: + continue + result = self._start_background(options) + return RuntimeResult(result.ok, result.message, result.status, promoted) def start_on_demand(self, options: ProcessStartOptions) -> RuntimeResult: """Atomically reuse a gateway or start one owned by local client leases.""" @@ -142,7 +152,7 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]): status = self.status() if status.running: return RuntimeResult(False, "gateway_already_running", status) - GatewayClientLease(self, kind="gateway-start").mark_ephemeral() + GatewayClientLease(self, kind="gateway-start")._mark_ephemeral_locked() return self._start_background(options) def _start_background(self, options: ProcessStartOptions) -> RuntimeResult: @@ -158,7 +168,10 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]): def stop(self, *, timeout_s: int = 20) -> RuntimeResult: """Stop the gateway recorded by this runtime.""" with self._lifecycle_lock(): - return self._result(self._stop(timeout_s=timeout_s)) + result = self._stop(timeout_s=timeout_s) + if result.ok or result.message in {"gateway_not_running", "gateway_state_stale"}: + GatewayClientLease(self, kind="gateway-stop")._clear_locked() + return self._result(result) def status(self, *, reason: str | None = None) -> GatewayStatus: """Return process, launch, and client lifetime state in one snapshot.""" @@ -186,9 +199,7 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]): @contextmanager def foreground_instance(self, options: ProcessStartOptions) -> Generator[None]: """Publish this foreground gateway while it is available to local clients.""" - launch_mode = self._claim_current_process(options) - if launch_mode == "foreground": - GatewayClientLease(self, kind="gateway-foreground").mark_persistent() + self._claim_current_process(options) try: yield finally: @@ -218,6 +229,11 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]): } ) self._write_state(state) + if launch_mode == "foreground": + GatewayClientLease( + self, + kind="gateway-foreground", + )._try_mark_persistent_locked() return launch_mode def _release_current_process(self) -> None: @@ -225,6 +241,7 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]): state = self._read_state() if state and self._record_matches_process(state, os.getpid()): self._clear_state() + GatewayClientLease(self, kind="gateway-exit")._finish_shutdown_locked() def restart(self, options: ProcessStartOptions, *, timeout_s: int = 20) -> RuntimeResult: """Restart an existing gateway without creating a new persistent instance.""" @@ -264,21 +281,20 @@ class GatewayClientLease: self.state_path = state_path.with_name( f"{state_path.stem}.clients{state_path.suffix}" ) + self.lifecycle_lock = FileLock(f"{state_path}.lock") self.lock = FileLock(f"{self.state_path}.lock") self._acquired = False def acquire(self) -> None: """Register this client before it starts or attaches to the gateway.""" - with self.lock: - state = self._live_state() - clients = self._clients(state) - clients[self.token] = { - "pid": self.pid, - "kind": self.kind, - "identity": self._process_identity(self.pid), - } - self._write_state(state) - self._acquired = True + while True: + self.wait_for_shutdown() + with self.lifecycle_lock, self.lock: + state = self._live_state() + if state.get("stopping"): + continue + self._register(state) + return def ensure_on_demand_gateway(self, options: GatewayStartOptions) -> RuntimeResult: """Atomically reuse a gateway or start one owned by local client leases.""" @@ -288,6 +304,10 @@ class GatewayClientLease: def mark_ephemeral(self) -> None: """Mark a gateway started by a client for last-client shutdown.""" + with self.lifecycle_lock: + self._mark_ephemeral_locked() + + def _mark_ephemeral_locked(self) -> None: with self.lock: state = self._live_state() state["auto_stop"] = True @@ -295,8 +315,18 @@ class GatewayClientLease: def mark_persistent(self) -> bool: """Keep an explicitly backgrounded gateway alive; return whether it was promoted.""" + while True: + self.wait_for_shutdown() + with self.lifecycle_lock: + promoted = self._try_mark_persistent_locked() + if promoted is not None: + return promoted + + def _try_mark_persistent_locked(self) -> bool | None: with self.lock: state = self._live_state() + if state.get("stopping"): + return None promoted = bool(state.get("auto_stop")) state["auto_stop"] = False self._write_or_clear(state) @@ -304,6 +334,10 @@ class GatewayClientLease: def clear(self) -> None: """Forget leases after an explicit gateway stop.""" + with self.lifecycle_lock: + self._clear_locked() + + def _clear_locked(self) -> None: with self.lock: self.state_path.unlink(missing_ok=True) @@ -317,35 +351,76 @@ class GatewayClientLease: clients=len(self._clients(state)), ) - def orphaned_on_demand(self) -> bool: - """Return whether an on-demand gateway has lost every live client.""" - snapshot = self.snapshot() - return snapshot.auto_stop and snapshot.clients == 0 + def begin_orphan_shutdown(self) -> bool: + """Commit shutdown only while an on-demand gateway still has no clients.""" + with self.lifecycle_lock, self.lock: + state = self._live_state() + if not bool(state.get("auto_stop")) or self._clients(state): + self._write_or_clear(state) + return False + state["stopping"] = True + self._write_state(state) + return True def release(self, *, timeout_s: int = 20) -> bool: """Release this client and stop an ephemeral gateway when it was the last.""" if not self._acquired: return False - should_stop = False + with self.lifecycle_lock: + with self.lock: + state = self._live_state() + clients = self._clients(state) + clients.pop(self.token, None) + self._acquired = False + should_stop = not clients and bool(state.get("auto_stop")) + self._write_or_clear(state) + if not should_stop: + return False + result = self.runtime._stop(timeout_s=timeout_s) + stopped = result.ok or result.message in { + "gateway_not_running", + "gateway_state_stale", + } + if stopped: + self._clear_locked() + else: + self._mark_ephemeral_locked() + return stopped + + def wait_for_shutdown(self, *, timeout_s: float = 20) -> None: + """Wait until a committed orphan shutdown can no longer accept clients.""" + deadline = time.monotonic() + timeout_s + while True: + with self.lifecycle_lock: + with self.lock: + state = self._live_state() + if not state.get("stopping"): + return + if not self.runtime.status().running: + self._finish_shutdown_locked() + return + if time.monotonic() >= deadline: + raise RuntimeError("gateway is still shutting down; try again shortly") + time.sleep(0.05) + + def _finish_shutdown_locked(self) -> None: with self.lock: state = self._live_state() - clients = self._clients(state) - clients.pop(self.token, None) - self._acquired = False - should_stop = not clients and bool(state.get("auto_stop")) - self._write_or_clear(state) - if not should_stop: - return False - result = self.runtime.stop(timeout_s=timeout_s) - with self.lock: - stopped = result.ok or result.message == "gateway_not_running" - if stopped: + state.pop("stopping", None) + if not self._clients(state): self.state_path.unlink(missing_ok=True) else: - state = self._live_state() - state["auto_stop"] = True self._write_state(state) - return stopped + + def _register(self, state: dict[str, object]) -> None: + clients = self._clients(state) + clients[self.token] = { + "pid": self.pid, + "kind": self.kind, + "identity": self._process_identity(self.pid), + } + self._write_state(state) + self._acquired = True def _live_state(self) -> dict[str, object]: state = self._read_state() @@ -432,7 +507,7 @@ async def monitor_gateway_clients( try: await asyncio.wait_for(shutdown_event.wait(), timeout=poll_interval_s) except TimeoutError: - if lease.orphaned_on_demand(): + if lease.begin_orphan_shutdown(): shutdown_event.set() return True return False diff --git a/tests/cli/test_commands.py b/tests/cli/test_commands.py index b72dadbe7..09e577c6e 100644 --- a/tests/cli/test_commands.py +++ b/tests/cli/test_commands.py @@ -1991,6 +1991,8 @@ def _patch_webui_managed_gateway( self.running = False return RuntimeResult(True, "gateway_stopped", self.status()) + _stop = stop + monkeypatch.setattr("nanobot.gateway.GatewayRuntime", _FakeRuntime) monkeypatch.setattr( "nanobot.cli.webui._prepare_webui_bundle_for_gateway", diff --git a/tests/cli/test_gateway_commands.py b/tests/cli/test_gateway_commands.py index 6684e7b3b..c651eca58 100644 --- a/tests/cli/test_gateway_commands.py +++ b/tests/cli/test_gateway_commands.py @@ -57,6 +57,7 @@ class FakeRuntime: def stop(self, *, timeout_s: int) -> RuntimeResult: self.stop_timeout = timeout_s + self.paths.state_path.with_name("gateway.clients.json").unlink(missing_ok=True) return RuntimeResult(True, "gateway_stopped", self.status_value) def status(self) -> GatewayStatus: @@ -223,7 +224,13 @@ def test_gateway_background_adopts_an_existing_on_demand_gateway(tmp_path): lease_state.write_text('{"auto_stop": true, "clients": {}}', encoding="utf-8") def already_running(_options: GatewayStartOptions) -> RuntimeResult: - return RuntimeResult(False, "gateway_already_running", fake_runtime.status_value) + lease_state.unlink(missing_ok=True) + return RuntimeResult( + False, + "gateway_already_running", + fake_runtime.status_value, + promoted=True, + ) fake_runtime.start_background = already_running # type: ignore[method-assign] @@ -311,6 +318,7 @@ def test_gateway_stop_treats_not_running_as_clean(tmp_path): def fake_stop(*, timeout_s: int) -> RuntimeResult: fake_runtime.stop_timeout = timeout_s + lease_state.unlink(missing_ok=True) return RuntimeResult(False, "gateway_not_running", fake_runtime.status_value) fake_runtime.stop = fake_stop # type: ignore[method-assign] diff --git a/tests/cli/test_tui_launcher.py b/tests/cli/test_tui_launcher.py index 3bb3206b4..3e22679ae 100644 --- a/tests/cli/test_tui_launcher.py +++ b/tests/cli/test_tui_launcher.py @@ -537,6 +537,8 @@ def test_gateway_started_for_tui_stops_when_its_last_lease_exits( stopped = True return SimpleNamespace(ok=True, message="gateway_stopped") + _stop = stop + monkeypatch.setattr("nanobot.gateway.GatewayRuntime", FakeRuntime) monkeypatch.setattr( "nanobot.cli.tui_launcher._webui_endpoint_reachable", diff --git a/tests/gateway/test_runtime.py b/tests/gateway/test_runtime.py index 54c011b6f..847206ef8 100644 --- a/tests/gateway/test_runtime.py +++ b/tests/gateway/test_runtime.py @@ -4,6 +4,7 @@ import os import signal import subprocess import sys +import threading from pathlib import Path from types import SimpleNamespace @@ -284,7 +285,7 @@ def test_last_interactive_client_stops_an_on_demand_gateway(tmp_path, monkeypatc stopped.append(timeout_s) return SimpleNamespace(ok=True, message="gateway_stopped") - monkeypatch.setattr(runtime, "stop", stop) + monkeypatch.setattr(runtime, "_stop", stop) tui = GatewayClientLease(runtime, kind="tui", pid=os.getpid(), token="tui") webui = GatewayClientLease(runtime, kind="webui", pid=os.getpid(), token="webui") @@ -299,6 +300,95 @@ def test_last_interactive_client_stops_an_on_demand_gateway(tmp_path, monkeypatc assert not webui.state_path.exists() +def test_last_client_shutdown_preserves_a_replacement_lease(tmp_path, monkeypatch): + runtime = GatewayRuntime(paths=_paths(tmp_path), platform_name="Linux") + monkeypatch.setattr(runtime, "_process_identity", lambda pid: pid) + monkeypatch.setattr(runtime, "_is_pid_running", lambda _pid: True) + stop_started = threading.Event() + finish_stop = threading.Event() + replacement_acquired = threading.Event() + + def stop(*, timeout_s: int): + assert timeout_s == 20 + stop_started.set() + assert finish_stop.wait(timeout=2) + return SimpleNamespace( + ok=True, + message="gateway_stopped", + status=runtime.status(), + ) + + monkeypatch.setattr(runtime, "_stop", stop) + original = GatewayClientLease(runtime, kind="tui", token="original") + replacement = GatewayClientLease(runtime, kind="webui", token="replacement") + original.acquire() + original.mark_ephemeral() + + release_thread = threading.Thread(target=original.release) + release_thread.start() + assert stop_started.wait(timeout=2) + + acquire_thread = threading.Thread( + target=lambda: (replacement.acquire(), replacement_acquired.set()) + ) + acquire_thread.start() + assert not replacement_acquired.wait(timeout=0.05) + + finish_stop.set() + release_thread.join(timeout=2) + acquire_thread.join(timeout=2) + + assert replacement_acquired.is_set() + state = json.loads(replacement.state_path.read_text(encoding="utf-8")) + assert set(state["clients"]) == {"replacement"} + + +def test_explicit_stop_clears_leases_before_accepting_a_replacement( + tmp_path, + monkeypatch, +): + runtime = GatewayRuntime(paths=_paths(tmp_path), platform_name="Linux") + monkeypatch.setattr(runtime, "_process_identity", lambda pid: pid) + monkeypatch.setattr(runtime, "_is_pid_running", lambda _pid: True) + stop_started = threading.Event() + finish_stop = threading.Event() + replacement_acquired = threading.Event() + + stale = GatewayClientLease(runtime, kind="tui", token="stale") + replacement = GatewayClientLease(runtime, kind="webui", token="replacement") + stale.acquire() + stale.mark_ephemeral() + + def stop(*, timeout_s: int): + assert timeout_s == 20 + stop_started.set() + assert finish_stop.wait(timeout=2) + return SimpleNamespace( + ok=True, + message="gateway_stopped", + status=runtime.status(), + ) + + monkeypatch.setattr(runtime, "_stop", stop) + stop_thread = threading.Thread(target=runtime.stop) + stop_thread.start() + assert stop_started.wait(timeout=2) + + acquire_thread = threading.Thread( + target=lambda: (replacement.acquire(), replacement_acquired.set()) + ) + acquire_thread.start() + assert not replacement_acquired.wait(timeout=0.05) + + finish_stop.set() + stop_thread.join(timeout=2) + acquire_thread.join(timeout=2) + + assert replacement_acquired.is_set() + state = json.loads(replacement.state_path.read_text(encoding="utf-8")) + assert set(state["clients"]) == {"replacement"} + + def test_on_demand_lifetime_is_recorded_before_the_gateway_spawns(tmp_path, monkeypatch): observed_auto_stop: list[bool] = [] @@ -331,18 +421,28 @@ def test_explicit_background_gateway_survives_the_last_client(tmp_path, monkeypa runtime = GatewayRuntime(paths=_paths(tmp_path), platform_name="Linux") monkeypatch.setattr(runtime, "_process_identity", lambda pid: pid) monkeypatch.setattr(runtime, "_is_pid_running", lambda _pid: True) + runtime._write_state( + { + "pid": os.getpid(), + "identity": os.getpid(), + "launch_mode": "background", + } + ) stopped: list[int] = [] monkeypatch.setattr( runtime, - "stop", + "_stop", lambda *, timeout_s: stopped.append(timeout_s), ) client = GatewayClientLease(runtime, kind="webui", pid=os.getpid()) client.acquire() client.mark_ephemeral() - assert GatewayClientLease(runtime, kind="gateway-background").mark_persistent() is True + result = runtime.start_background(GatewayStartOptions(port=18790)) + assert result.ok is False + assert result.message == "gateway_already_running" + assert result.promoted is True assert client.release() is False assert stopped == [] assert not client.state_path.exists() @@ -354,7 +454,7 @@ def test_failed_last_client_shutdown_remains_retryable(tmp_path, monkeypatch): monkeypatch.setattr(runtime, "_is_pid_running", lambda _pid: True) monkeypatch.setattr( runtime, - "stop", + "_stop", lambda *, timeout_s: SimpleNamespace(ok=False, message="gateway_stop_timeout"), ) client = GatewayClientLease(runtime, kind="tui", pid=os.getpid()) @@ -400,6 +500,39 @@ async def test_client_monitor_stops_an_orphaned_on_demand_gateway(tmp_path): assert orphaned is True assert shutdown_event.is_set() + assert json.loads(lease.state_path.read_text(encoding="utf-8"))["stopping"] is True + + +async def test_client_monitor_blocks_replacement_until_gateway_exit(tmp_path, monkeypatch): + runtime = GatewayRuntime(paths=_paths(tmp_path), platform_name="Linux") + monkeypatch.setattr(runtime, "_process_identity", lambda pid: pid) + monkeypatch.setattr(runtime, "_is_pid_running", lambda _pid: True) + runtime._write_state({"pid": os.getpid(), "identity": os.getpid()}) + monitor = GatewayClientLease(runtime, kind="gateway-monitor") + monitor.mark_ephemeral() + shutdown_event = asyncio.Event() + + assert await monitor_gateway_clients( + monitor, + shutdown_event, + poll_interval_s=0.001, + ) is True + + replacement = GatewayClientLease(runtime, kind="webui", token="replacement") + replacement_acquired = threading.Event() + acquire_thread = threading.Thread( + target=lambda: (replacement.acquire(), replacement_acquired.set()) + ) + acquire_thread.start() + assert not replacement_acquired.wait(timeout=0.05) + + runtime._release_current_process() + acquire_thread.join(timeout=2) + + assert replacement_acquired.is_set() + state = json.loads(replacement.state_path.read_text(encoding="utf-8")) + assert set(state["clients"]) == {"replacement"} + assert "stopping" not in state def test_start_background_uses_windows_process_group_flags(tmp_path, monkeypatch):