From 19be5be1c02c881ab3b15a2f591d41a56cb7b974 Mon Sep 17 00:00:00 2001
From: Xubin Ren <52506698+Re-bin@users.noreply.github.com>
Date: Sun, 16 Aug 2026 23:12:51 +0800
Subject: [PATCH] refactor(gateway): scope foreground registration

---
 nanobot/cli/gateway_runtime.py |  5 +----
 nanobot/gateway/runtime.py     | 18 +++++++++++-----
 tests/gateway/test_runtime.py  | 38 ++++++++++++++++++++--------------
 3 files changed, 36 insertions(+), 25 deletions(-)

diff --git a/nanobot/cli/gateway_runtime.py b/nanobot/cli/gateway_runtime.py
index 1fe8e0da4..0de9fbe02 100644
--- a/nanobot/cli/gateway_runtime.py
+++ b/nanobot/cli/gateway_runtime.py
@@ -954,8 +954,5 @@ def _run_gateway(
             finally:
                 restore_shutdown_handlers()
 
-    gateway_runtime.claim_current_process(gateway_start_options)
-    try:
+    with gateway_runtime.foreground_instance(gateway_start_options):
         asyncio.run(run())
-    finally:
-        gateway_runtime.release_current_process()
diff --git a/nanobot/gateway/runtime.py b/nanobot/gateway/runtime.py
index 1cfe9ce01..fd6a11b24 100644
--- a/nanobot/gateway/runtime.py
+++ b/nanobot/gateway/runtime.py
@@ -10,10 +10,11 @@ import tempfile
 import time
 import uuid
 from collections.abc import Callable
+from contextlib import contextmanager
 from dataclasses import dataclass
 from datetime import UTC, datetime
 from pathlib import Path
-from typing import Any, cast
+from typing import Any, Generator, cast
 
 from filelock import FileLock
 
@@ -103,8 +104,16 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
     def _build_child_command(self, options: ProcessStartOptions) -> list[str]:
         return build_gateway_command(self.python_executable, options)
 
-    def claim_current_process(self, options: ProcessStartOptions) -> None:
-        """Publish the running foreground gateway for local client discovery."""
+    @contextmanager
+    def foreground_instance(self, options: ProcessStartOptions) -> Generator[None]:
+        """Publish this foreground gateway while it is available to local clients."""
+        self._claim_current_process(options)
+        try:
+            yield
+        finally:
+            self._release_current_process()
+
+    def _claim_current_process(self, options: ProcessStartOptions) -> None:
         with self._lifecycle_lock():
             state = self._read_state() or {}
             pid = os.getpid()
@@ -123,8 +132,7 @@ class GatewayRuntime(ManagedProcessRuntime[ProcessStartOptions]):
             )
             self._write_state(state)
 
-    def release_current_process(self) -> None:
-        """Remove this foreground gateway's state without touching a replacement."""
+    def _release_current_process(self) -> None:
         with self._lifecycle_lock():
             state = self._read_state()
             if state and self._record_matches_process(state, os.getpid()):
diff --git a/tests/gateway/test_runtime.py b/tests/gateway/test_runtime.py
index 45d0748dc..968b1480f 100644
--- a/tests/gateway/test_runtime.py
+++ b/tests/gateway/test_runtime.py
@@ -118,15 +118,12 @@ def test_foreground_gateway_claim_is_discoverable_and_released(tmp_path, monkeyp
         config_path="/tmp/config.json",
     )
 
-    runtime.claim_current_process(options)
-
-    status = runtime.status()
-    assert status.running is True
-    assert status.pid == os.getpid()
-    assert status.port == 18790
-    assert status.command == tuple(runtime._build_child_command(options))
-
-    runtime.release_current_process()
+    with runtime.foreground_instance(options):
+        status = runtime.status()
+        assert status.running is True
+        assert status.pid == os.getpid()
+        assert status.port == 18790
+        assert status.command == tuple(runtime._build_child_command(options))
 
     assert runtime.status().running is False
     assert not runtime.paths.state_path.exists()
@@ -135,17 +132,26 @@ def test_foreground_gateway_claim_is_discoverable_and_released(tmp_path, monkeyp
 def test_foreground_gateway_release_preserves_a_replacement_state(tmp_path, monkeypatch):
     runtime = GatewayRuntime(paths=_paths(tmp_path), platform_name="Darwin")
     monkeypatch.setattr(runtime, "_process_identity", lambda pid: pid)
-    runtime.claim_current_process(GatewayStartOptions(port=18790))
-    replacement = json.loads(runtime.paths.state_path.read_text(encoding="utf-8"))
-    replacement["pid"] = os.getpid() + 1
-    replacement["identity"] = replacement["pid"]
-    runtime.paths.state_path.write_text(json.dumps(replacement), encoding="utf-8")
-
-    runtime.release_current_process()
+    with runtime.foreground_instance(GatewayStartOptions(port=18790)):
+        replacement = json.loads(runtime.paths.state_path.read_text(encoding="utf-8"))
+        replacement["pid"] = os.getpid() + 1
+        replacement["identity"] = replacement["pid"]
+        runtime.paths.state_path.write_text(json.dumps(replacement), encoding="utf-8")
 
     assert runtime.paths.state_path.exists()
 
 
+def test_foreground_gateway_clears_its_state_after_an_error(tmp_path, monkeypatch):
+    runtime = GatewayRuntime(paths=_paths(tmp_path), platform_name="Darwin")
+    monkeypatch.setattr(runtime, "_process_identity", lambda pid: pid)
+
+    with pytest.raises(RuntimeError, match="startup failed"):
+        with runtime.foreground_instance(GatewayStartOptions(port=18790)):
+            raise RuntimeError("startup failed")
+
+    assert not runtime.paths.state_path.exists()
+
+
 def test_stop_reaps_an_owned_child_without_consuming_the_shutdown_timeout(
     tmp_path,
     monkeypatch,