An attached client re-acquires its runtime once its spawner exits

src/hai_agents/client.py

@@ -35,6 +35,7 @@
 
 if typing.TYPE_CHECKING:
     from hai_agents_local.runtime import Inference, LocalRuntime
+    from hai_agents_local.runtime.acquire import RuntimeHold
 
 API_KEY_VAR = "HAI_API_KEY"
 
@@ -79,12 +80,24 @@ def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
     return wrapper
 
 
-class Client(BaseClient):
-    __init__ = _default_api_key(BaseClient.__init__)
-    local_runtime: typing.Optional[LocalRuntime] = None
-    _owns_runtime = False
+class _LocalRuntimeAccess:
+    """The runtime behind a local client, read through its hold so a re-acquired runtime shows up everywhere."""
+
+    _hold: typing.Optional[RuntimeHold] = None
     _auto_bridges = True
 
+    @property
+    def local_runtime(self) -> typing.Optional[LocalRuntime]:
+        return None if self._hold is None else self._hold.runtime
+
+    @property
+    def _owns_runtime(self) -> bool:
+        return self._hold is not None and self._hold.owned
+
+
+class Client(_LocalRuntimeAccess, BaseClient):
+    __init__ = _default_api_key(BaseClient.__init__)
+
     @classmethod
     def local(
         cls,
@@ -96,16 +109,16 @@ def local(
         timeout: typing.Optional[float] = None,
     ) -> Client:
         """A client on a local agent runtime: ``runtime`` if given, else one this client starts and owns."""
-        from hai_agents_local.runtime import acquire_runtime
+        from hai_agents_local.runtime.acquire import RuntimeHold
 
-        runtime, owned = acquire_runtime(runtime, inference=inference, local_options=local_options)
+        hold = RuntimeHold.acquire(runtime, inference=inference, local_options=local_options)
         try:
-            client = cls(base_url=runtime.base_url, api_key=runtime.api_key, httpx_client=runtime.http_client(timeout))
+            client = cls(base_url=hold.runtime.base_url, api_key=hold.api_key, httpx_client=hold.http_client(timeout))
         except BaseException:
-            if owned:
-                runtime.shutdown()
+            if hold.owned:
+                hold.runtime.shutdown()
             raise
-        client.local_runtime, client._owns_runtime, client._auto_bridges = runtime, owned, auto_bridges
+        client._hold, client._auto_bridges = hold, auto_bridges
         return client
 
     def close(self) -> None:
@@ -180,16 +193,13 @@ def sessions(self) -> SessionsClient:
             from hai_agents_local.sessions import LocalSessionsClient
 
             self._sessions = LocalSessionsClient(
-                client_wrapper=self._client_wrapper, runtime=self.local_runtime, auto_bridges=self._auto_bridges
+                client_wrapper=self._client_wrapper, hold=self._hold, auto_bridges=self._auto_bridges
             )
         return self._sessions
 
 
-class AsyncClient(AsyncBaseClient):
+class AsyncClient(_LocalRuntimeAccess, AsyncBaseClient):
     __init__ = _default_api_key(AsyncBaseClient.__init__)
-    local_runtime: typing.Optional[LocalRuntime] = None
-    _owns_runtime = False
-    _auto_bridges = True
 
     @classmethod
     async def local(
@@ -202,18 +212,18 @@ async def local(
         timeout: typing.Optional[float] = None,
     ) -> AsyncClient:
         """A client on a local agent runtime: ``runtime`` if given, else one this client starts and owns."""
-        from hai_agents_local.runtime import acquire_runtime_async
+        from hai_agents_local.runtime.acquire import RuntimeHold
 
-        runtime, owned = await acquire_runtime_async(runtime, inference=inference, local_options=local_options)
+        hold = await RuntimeHold.acquire_async(runtime, inference=inference, local_options=local_options)
         try:
             client = cls(
-                base_url=runtime.base_url, api_key=runtime.api_key, httpx_client=runtime.async_http_client(timeout)
+                base_url=hold.runtime.base_url, api_key=hold.api_key, httpx_client=hold.async_http_client(timeout)
             )
         except BaseException:
-            if owned:
-                await asyncio.to_thread(runtime.shutdown)
+            if hold.owned:
+                await asyncio.to_thread(hold.runtime.shutdown)
             raise
-        client.local_runtime, client._owns_runtime, client._auto_bridges = runtime, owned, auto_bridges
+        client._hold, client._auto_bridges = hold, auto_bridges
         return client
 
     async def aclose(self) -> None:
@@ -290,6 +300,6 @@ def sessions(self) -> AsyncSessionsClient:
             from hai_agents_local.sessions import LocalAsyncSessionsClient
 
             self._sessions = LocalAsyncSessionsClient(
-                client_wrapper=self._client_wrapper, runtime=self.local_runtime, auto_bridges=self._auto_bridges
+                client_wrapper=self._client_wrapper, hold=self._hold, auto_bridges=self._auto_bridges
             )
         return self._sessions

src/hai_agents_local/runtime/acquire.py

@@ -6,10 +6,15 @@
 import threading
 import typing
 
+import httpx
+
 from hai_agents.client import API_KEY_VAR, default_api_key
 
+from . import identity
+from .errors import LocalRuntimeError
 from .inference import Inference
-from .runtime import LocalRuntime
+from .process import probe_health
+from .runtime import CLIENT_TIMEOUT_S, LocalRuntime
 
 SHARED_RECIPE = "shared"
 RECIPE_ENV = "HAI_AGENT_RUNTIME_RECIPE"
@@ -50,6 +55,101 @@ async def acquire_runtime_async(
     return _claim(started, inference)
 
 
+class RuntimeHold:
+    """A client's runtime, and how to get it back once the process that spawned it has exited.
+
+    A runtime lives as long as its spawner (parent watchdog). A client that merely attached to it
+    is left talking to a closed port, so before new work it re-acquires: another runtime at the
+    same place, possibly under a new token, which the HTTP layer picks up through ``api_key``.
+    """
+
+    def __init__(
+        self,
+        runtime: LocalRuntime,
+        owned: bool,
+        *,
+        launch: typing.Optional[typing.Dict[str, typing.Any]] = None,
+    ) -> None:
+        self.runtime = runtime
+        self.owned = owned
+        self._launch = launch
+
+    @classmethod
+    def acquire(
+        cls,
+        runtime: typing.Optional[LocalRuntime],
+        *,
+        inference: typing.Optional[Inference],
+        local_options: typing.Optional[typing.Dict[str, typing.Any]],
+    ) -> "RuntimeHold":
+        acquired, owned = acquire_runtime(runtime, inference=inference, local_options=local_options)
+        return cls(acquired, owned, launch=_relaunch(runtime, inference, local_options))
+
+    @classmethod
+    async def acquire_async(
+        cls,
+        runtime: typing.Optional[LocalRuntime],
+        *,
+        inference: typing.Optional[Inference],
+        local_options: typing.Optional[typing.Dict[str, typing.Any]],
+    ) -> "RuntimeHold":
+        acquired, owned = await acquire_runtime_async(runtime, inference=inference, local_options=local_options)
+        return cls(acquired, owned, launch=_relaunch(runtime, inference, local_options))
+
+    def api_key(self) -> str:
+        return self.runtime.api_key
+
+    def http_client(self, timeout: typing.Optional[float] = None) -> httpx.Client:
+        return identity.http_client(
+            self.api_key, timeout=CLIENT_TIMEOUT_S if timeout is None else timeout, follow_redirects=True
+        )
+
+    def async_http_client(self, timeout: typing.Optional[float] = None) -> httpx.AsyncClient:
+        return identity.async_http_client(
+            self.api_key, timeout=CLIENT_TIMEOUT_S if timeout is None else timeout, follow_redirects=True
+        )
+
+    def lost(self) -> bool:
+        """The runtime this client attached to is gone: nothing answers, or a runtime under another token took its port."""
+        if self._launch is None or self.owned:
+            return False
+        try:
+            return probe_health(self.runtime.base_url, self.runtime.api_key) is None
+        except LocalRuntimeError:
+            return True
+
+    def reacquire(self) -> bool:
+        """Replace a lost runtime; True when the client now talks to another one."""
+        if not self.lost():
+            return False
+        self._replace(*acquire_runtime(None, **typing.cast(typing.Dict[str, typing.Any], self._launch)))
+        return True
+
+    async def reacquire_async(self) -> bool:
+        if not await asyncio.to_thread(self.lost):
+            return False
+        self._replace(*await acquire_runtime_async(None, **typing.cast(typing.Dict[str, typing.Any], self._launch)))
+        return True
+
+    def _replace(self, runtime: LocalRuntime, owned: bool) -> None:
+        if runtime.base_url != self.runtime.base_url:
+            raise LocalRuntimeError(
+                f"the replacement runtime listens at {runtime.base_url}, not {self.runtime.base_url}"
+            )
+        self.runtime, self.owned = runtime, owned
+
+
+def _relaunch(
+    runtime: typing.Optional[LocalRuntime],
+    inference: typing.Optional[Inference],
+    local_options: typing.Optional[typing.Dict[str, typing.Any]],
+) -> typing.Optional[typing.Dict[str, typing.Any]]:
+    """How to acquire again; None for a caller-provided runtime, whose lifetime is the caller's business."""
+    if runtime is not None:
+        return None
+    return {"inference": inference, "local_options": local_options}
+
+
 def _require_unconfigured(
     inference: typing.Optional[Inference], local_options: typing.Optional[typing.Dict[str, typing.Any]]
 ) -> None:

src/hai_agents_local/runtime/identity.py

@@ -13,6 +13,8 @@
 
 CHALLENGE_HEADER = "X-Hai-Runtime-Challenge"
 PROOF_HEADER = "X-Hai-Runtime-Proof"
+# A callable token follows a client whose runtime gets replaced.
+Token = typing.Union[str, typing.Callable[[], str]]
 
 
 def challenge() -> str:
@@ -31,31 +33,35 @@ def verify(token: str, request: httpx.Request, response: httpx.Response) -> None
         )
 
 
-def event_hooks(token: str) -> typing.Dict[str, typing.List[typing.Callable[..., None]]]:
+def _current(token: Token) -> str:
+    return token if isinstance(token, str) else token()
+
+
+def event_hooks(token: Token) -> typing.Dict[str, typing.List[typing.Callable[..., None]]]:
     def send_challenge(request: httpx.Request) -> None:
         request.headers[CHALLENGE_HEADER] = challenge()
 
     def check_proof(response: httpx.Response) -> None:
-        verify(token, response.request, response)
+        verify(_current(token), response.request, response)
 
     return {"request": [send_challenge], "response": [check_proof]}
 
 
-def async_event_hooks(token: str) -> typing.Dict[str, typing.List[typing.Callable[..., typing.Awaitable[None]]]]:
+def async_event_hooks(token: Token) -> typing.Dict[str, typing.List[typing.Callable[..., typing.Awaitable[None]]]]:
     async def send_challenge(request: httpx.Request) -> None:
         request.headers[CHALLENGE_HEADER] = challenge()
 
     async def check_proof(response: httpx.Response) -> None:
-        verify(token, response.request, response)
+        verify(_current(token), response.request, response)
 
     return {"request": [send_challenge], "response": [check_proof]}
 
 
-def http_client(token: str, **kwargs: typing.Any) -> httpx.Client:
+def http_client(token: Token, **kwargs: typing.Any) -> httpx.Client:
     """A proxy-free client that rejects any response not proven by the runtime holding ``token``."""
     return httpx.Client(trust_env=False, event_hooks=event_hooks(token), **kwargs)
 
 
-def async_http_client(token: str, **kwargs: typing.Any) -> httpx.AsyncClient:
+def async_http_client(token: Token, **kwargs: typing.Any) -> httpx.AsyncClient:
     """``http_client`` for asyncio callers."""
     return httpx.AsyncClient(trust_env=False, event_hooks=async_event_hooks(token), **kwargs)

src/hai_agents_local/sessions.py

@@ -28,6 +28,7 @@
 
 if typing.TYPE_CHECKING:
     from .runtime import LocalRuntime
+    from .runtime.acquire import RuntimeHold
 
 logger = logging.getLogger(__name__)
 
@@ -263,16 +264,25 @@ class _LocalSessionsState:
     """Bridge and session bookkeeping shared by the sync and async local sessions clients."""
 
     def __init__(
-        self, *, client_wrapper: typing.Any, runtime: typing.Optional[LocalRuntime] = None, auto_bridges: bool = True
+        self, *, client_wrapper: typing.Any, hold: typing.Optional[RuntimeHold] = None, auto_bridges: bool = True
     ) -> None:
         super().__init__(client_wrapper=client_wrapper)
-        self._runtime = runtime
+        self._hold = hold
+        self._client_wrapper = client_wrapper
         self._auto_bridges = auto_bridges
-        self._cancel_remote: RemoteCancel = functools.partial(_cancel_remote_session, client_wrapper, runtime)
         self._owned_bridges: typing.Dict[str, typing.List[str]] = {}
         # Sessions this client created on a local runtime; they never keep that runtime alive past close().
         self.own_session_ids: typing.Set[str] = set()
 
+    @property
+    def _runtime(self) -> typing.Optional[LocalRuntime]:
+        return None if self._hold is None else self._hold.runtime
+
+    @property
+    def _cancel_remote(self) -> RemoteCancel:
+        """Bound to the runtime of the moment, so an exit cancel reaches the runtime its session lives on."""
+        return functools.partial(_cancel_remote_session, self._client_wrapper, self._runtime)
+
     def _live_sessions(self) -> typing.List[str]:
         """Forget sessions whose bridges all stopped, since each such stop already ended or cancelled its session."""
         for session_id, bridge_ids in list(self._owned_bridges.items()):
@@ -309,6 +319,8 @@ def cancel_session(self, id: str, *, request_options: typing.Optional[RequestOpt
 
     @functools.wraps(SessionsClient.create_session)
     def create_session(self, **kwargs: typing.Any) -> typing.Any:
+        if self._hold is not None:
+            self._hold.reacquire()
         bridges = _localize(self._raw_client._client_wrapper, self._runtime, kwargs) if self._auto_bridges else []
         if bridges:
             _apply_runaway_budgets(kwargs)
@@ -354,6 +366,8 @@ async def cancel_session(self, id: str, *, request_options: typing.Optional[Requ
 
     @functools.wraps(AsyncSessionsClient.create_session)
     async def create_session(self, **kwargs: typing.Any) -> typing.Any:
+        if self._hold is not None:
+            await self._hold.reacquire_async()
         bridges = _localize(self._raw_client._client_wrapper, self._runtime, kwargs) if self._auto_bridges else []
         if bridges:
             _apply_runaway_budgets(kwargs)