From fe71af0073b23f8f82a793f19101b156a4a4ff33 Mon Sep 17 00:00:00 2001 From: abonneth Date: Thu, 8 Oct 2026 15:48:39 +0200 Subject: [PATCH] An attached client re-acquires its runtime once the process that spawned it exits Co-authored-by: Cursor --- src/hai_agents/client.py | 54 +++++++----- src/hai_agents_local/runtime/acquire.py | 102 ++++++++++++++++++++++- src/hai_agents_local/runtime/identity.py | 18 ++-- src/hai_agents_local/sessions.py | 20 ++++- tests/test_runtime_placement.py | 44 +++++++++- 5 files changed, 204 insertions(+), 34 deletions(-) diff --git a/src/hai_agents/client.py b/src/hai_agents/client.py index c00d0a9..e4818b3 100644 --- a/src/hai_agents/client.py +++ b/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 diff --git a/src/hai_agents_local/runtime/acquire.py b/src/hai_agents_local/runtime/acquire.py index 9040baf..c5c6957 100644 --- a/src/hai_agents_local/runtime/acquire.py +++ b/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: diff --git a/src/hai_agents_local/runtime/identity.py b/src/hai_agents_local/runtime/identity.py index c55736a..19c870a 100644 --- a/src/hai_agents_local/runtime/identity.py +++ b/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) diff --git a/src/hai_agents_local/sessions.py b/src/hai_agents_local/sessions.py index b866425..deb9b31 100644 --- a/src/hai_agents_local/sessions.py +++ b/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) diff --git a/tests/test_runtime_placement.py b/tests/test_runtime_placement.py index f260692..44431e7 100644 --- a/tests/test_runtime_placement.py +++ b/tests/test_runtime_placement.py @@ -20,8 +20,8 @@ class RuntimeServer(ThreadingHTTPServer): """A loopback stand-in for the runtime: answers every response with an HMAC proof keyed by `proof_token`.""" - def __init__(self, token): - super().__init__(("127.0.0.1", 0), _RuntimeHandler) + def __init__(self, token, port=0): + super().__init__(("127.0.0.1", port), _RuntimeHandler) self.token = token self.proof_token = token self.requests = [] @@ -43,6 +43,9 @@ def do_GET(self): def do_DELETE(self): self._route() + def do_POST(self): + self._route() + def _route(self): server = self.server server.requests.append({name.lower() for name in self.headers}) @@ -196,6 +199,43 @@ def test_runtimes_for_different_inference_never_share(tmp_path, runtime_server): assert LocalRuntime.ensure_started(spawn_env={}, **start).serves == hosted +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [False, True]) +async def test_attached_client_reacquires_a_runtime_its_spawner_took_along( + tmp_path, runtime_server, monkeypatch, asynchronous +): + from hai_agents.core.api_error import ApiError + + port = runtime_server.port + write_owner_only(token_file_path(port, cache_dir=tmp_path), runtime_server.token) + monkeypatch.setattr(LocalRuntime, "ensure_started", lambda **_: LocalRuntime.attach(port=port, cache_dir=tmp_path)) + client = ( + await AsyncClient.local(local_options={"cache_dir": tmp_path}) + if asynchronous + else Client.local(local_options={"cache_dir": tmp_path}) + ) + assert not client._owns_runtime and client.local_runtime.api_key == "local-token" + + runtime_server.shutdown() + runtime_server.server_close() + replacement = RuntimeServer("fresh-token", port=port) + threading.Thread(target=replacement.serve_forever, daemon=True).start() + write_owner_only(token_file_path(port, cache_dir=tmp_path), replacement.token) + try: + # The replacement answers 401 to any other bearer, so a 404 proves the client switched tokens. + with pytest.raises(ApiError) as failed: + if asynchronous: + await client.sessions.create_session(agent="h/agent", messages="hi") + else: + client.sessions.create_session(agent="h/agent", messages="hi") + assert failed.value.status_code == 404 + assert client.local_runtime.api_key == "fresh-token" + finally: + await client.aclose() if asynchronous else client.close() + replacement.shutdown() + replacement.server_close() + + @pytest.mark.parametrize("proof_token", ["local-token", "squatter-token", None]) def test_only_the_runtime_holding_the_token_ever_receives_it(tmp_path, runtime_server, proof_token): write_owner_only(token_file_path(runtime_server.port, cache_dir=tmp_path), "local-token")