feat(sdk): compose local agents with Desktop and Workstation

.github/workflows/publish.yml

@@ -37,8 +37,33 @@ jobs:
           name: dist
           path: dist/
 
-  publish:
+  runtime-pin:
     needs: build
+    runs-on: macos-14
+    timeout-minutes: 10
+    steps:
+      - name: Install uv
+        uses: astral-sh/setup-uv@v10.1.0
+
+      - uses: actions/download-artifact@v8
+        with:
+          name: dist
+          path: dist/
+
+      - name: Pinned runtime downloads, verifies and serves the shared recipe
+        env:
+          HAI_AGENT_RUNTIME_FAKE: "1"
+        run: |
+          uv venv -q && uv pip install -q dist/*.whl
+          .venv/bin/python -c '
+          import tempfile
+          from hai_agents_local.runtime import LocalRuntime
+          LocalRuntime.ensure_started(required_recipe="shared", cache_dir=tempfile.mkdtemp(), port=18799,
+                                      spawn_env={"HAI_AGENT_RUNTIME_RECIPE": "shared"}).shutdown()
+          '
+
+  publish:
+    needs: [build, runtime-pin]
     runs-on: ubuntu-latest
     timeout-minutes: 10
     environment: release

README.md

@@ -67,6 +67,36 @@ print(result.answer)
 
 `result` is a `SessionRunResult`: `id`, `status`, `answer`, the accumulated `events`, and `final_changes`.
 
+## Local agents
+
+`Client.local()` and `await AsyncClient.local()` run the agent on this machine through a local agent runtime,
+started on demand or passed in with `runtime=...`. Each environment picks `host="user_device"` or `host="cloud"`.
+`Client()` keeps using the hosted Agents API. Closing the client stops a runtime it started, unless other clients still
+have active sessions there.
+
+```python
+from hai_agents import Client
+
+with Client.local(local_options={"binary_path": "/path/to/hai-agent-runtime"}) as client:
+    session = client.sessions.create_session(
+        agent={
+            "name": "local-example",
+            "description": "Local workstation example",
+            "instructions": "Answer the user's task using the workstation tools.",
+            "environments": [
+                {"id": "workstation", "kind": "workstation", "host": "user_device"}
+            ],
+        },
+        messages=[{"type": "user_message", "message": "Print hello using the shell."}],
+        max_steps=8,
+        max_time_s=120,
+    )
+    # Poll or steer the session here, before leaving the client context.
+```
+
+Inference stays hosted (`HAI_API_KEY`) unless you pass `inference=Inference.self_hosted(url, model=...)`
+(`from hai_agents_local.runtime import Inference`).
+
 ## How a session works
 
 A session is one run of an agent against a task. It moves through a small set of states: `pending`, `running`, and then a settled state such as `completed`, `idle`, `failed`, `timed_out`, or `interrupted`.

scripts/bump_runtime.py

@@ -0,0 +1,67 @@
+"""Pin a hai-agent-runtime release: its version plus a fresh sha256 for every platform in the pin file."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+from pathlib import Path
+
+# Stdlib only: runs in a bare checkout with no dependencies installed.
+PIN_FILE = Path(__file__).parents[1] / "src" / "hai_agents_local" / "runtime" / "pin.json"
+_VERSION_RE = re.compile(r"[0-9]+\.[0-9]+\.[0-9]+(?:[-+][0-9A-Za-z.-]+)?")
+_SHA256_RE = re.compile(r"[0-9a-f]{64}")
+_PLACEHOLDER_SHA256 = "0" * 64
+
+
+def apply_bump(pin: dict, version: str, shas: dict[str, str]) -> dict:
+    """`pin` moved to `version`; raises unless `shas` holds a real digest for exactly the pinned platforms."""
+    if not _VERSION_RE.fullmatch(version):
+        raise ValueError("runtime version must be a release version, e.g. 0.1.13")
+    for platform, sha in shas.items():
+        if not _SHA256_RE.fullmatch(sha) or sha == _PLACEHOLDER_SHA256:
+            raise ValueError(f"{platform}: {sha!r} is not a lowercase 64-char sha256")
+    published = set(pin["sha256"])
+    extra = shas.keys() - published
+    if extra:
+        raise ValueError(f"no manifest entry for platform(s): {sorted(extra)}")
+    # A platform left on its old digest would fail verification at the new version's URL.
+    missing = published - shas.keys()
+    if missing:
+        raise ValueError(f"missing sha for published platform(s): {sorted(missing)}")
+    return {**pin, "version": version, "sha256": {platform: shas[platform] for platform in pin["sha256"]}}
+
+
+def _parse_args(argv: list[str]) -> tuple[str, dict[str, str]]:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--version", required=True)
+    parser.add_argument(
+        "--sha",
+        action="append",
+        required=True,
+        metavar="PLATFORM=SHA256",
+        help="per-platform digest, e.g. darwin-arm64=<hex> (repeatable)",
+    )
+    args = parser.parse_args(argv)
+    shas: dict[str, str] = {}
+    for entry in args.sha:
+        platform, _, sha = entry.partition("=")
+        if not platform or not sha:
+            parser.error(f"--sha must be PLATFORM=SHA256, got {entry!r}")
+        if platform in shas:
+            parser.error(f"duplicate --sha for {platform}")
+        shas[platform] = sha.lower()
+    return args.version, shas
+
+
+def main(argv: list[str]) -> int:
+    version, shas = _parse_args(argv)
+    pin = apply_bump(json.loads(PIN_FILE.read_text(encoding="utf-8")), version, shas)
+    PIN_FILE.write_text(json.dumps(pin, indent=2) + "\n", encoding="utf-8")
+    print(f"bumped runtime to {version} ({', '.join(sorted(shas))})")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main(sys.argv[1:]))

src/hai_agents/client.py

@@ -7,6 +7,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import typing
 
 import typing_extensions
@@ -27,8 +28,57 @@
 from .sessions.client import AsyncSessionsClient, SessionsClient
 from .tools import ToolInput, as_tools
 
+if typing.TYPE_CHECKING:
+    from hai_agents_local.runtime import Inference, LocalRuntime
+
 
 class Client(BaseClient):
+    local_runtime: typing.Optional[LocalRuntime] = None
+    _owns_runtime = False
+    _auto_bridges = True
+
+    @classmethod
+    def local(
+        cls,
+        *,
+        runtime: typing.Optional[LocalRuntime] = None,
+        inference: typing.Optional[Inference] = None,
+        local_options: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        auto_bridges: bool = True,
+        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
+
+        runtime, owned = acquire_runtime(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))
+        except BaseException:
+            if owned:
+                runtime.shutdown()
+            raise
+        client.local_runtime, client._owns_runtime, client._auto_bridges = runtime, owned, auto_bridges
+        return client
+
+    def close(self) -> None:
+        """Stop sessions this client bridged; a local client also stops its runtime once no other client uses it."""
+        try:
+            if self._sessions is not None:
+                self._sessions.close()
+        finally:
+            if self.local_runtime is not None:
+                try:
+                    if self._owns_runtime:
+                        self.local_runtime.shutdown_if_idle(getattr(self._sessions, "own_session_ids", ()))
+                finally:
+                    self._client_wrapper.httpx_client.httpx_client.close()
+
+    def __enter__(self) -> Client:
+        return self
+
+    def __exit__(self, *exc: typing.Any) -> None:
+        self.close()
+
     def run_session(
         self,
         *,
@@ -81,11 +131,63 @@ def sessions(self) -> SessionsClient:
         if self._sessions is None:
             from hai_agents_local.sessions import LocalSessionsClient
 
-            self._sessions = LocalSessionsClient(client_wrapper=self._client_wrapper)
+            self._sessions = LocalSessionsClient(
+                client_wrapper=self._client_wrapper, runtime=self.local_runtime, auto_bridges=self._auto_bridges
+            )
         return self._sessions
 
 
 class AsyncClient(AsyncBaseClient):
+    local_runtime: typing.Optional[LocalRuntime] = None
+    _owns_runtime = False
+    _auto_bridges = True
+
+    @classmethod
+    async def local(
+        cls,
+        *,
+        runtime: typing.Optional[LocalRuntime] = None,
+        inference: typing.Optional[Inference] = None,
+        local_options: typing.Optional[typing.Dict[str, typing.Any]] = None,
+        auto_bridges: bool = True,
+        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
+
+        runtime, owned = await acquire_runtime_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)
+            )
+        except BaseException:
+            if owned:
+                await asyncio.to_thread(runtime.shutdown)
+            raise
+        client.local_runtime, client._owns_runtime, client._auto_bridges = runtime, owned, auto_bridges
+        return client
+
+    async def aclose(self) -> None:
+        """Stop sessions this client bridged; a local client also stops its runtime once no other client uses it."""
+        try:
+            if self._sessions is not None:
+                await self._sessions.aclose()
+        finally:
+            if self.local_runtime is not None:
+                try:
+                    if self._owns_runtime:
+                        await asyncio.to_thread(
+                            self.local_runtime.shutdown_if_idle, getattr(self._sessions, "own_session_ids", ())
+                        )
+                finally:
+                    await self._client_wrapper.httpx_client.httpx_client.aclose()
+
+    async def __aenter__(self) -> AsyncClient:
+        return self
+
+    async def __aexit__(self, *exc: typing.Any) -> None:
+        await self.aclose()
+
     async def run_session(
         self,
         *,
@@ -138,5 +240,7 @@ def sessions(self) -> AsyncSessionsClient:
         if self._sessions is None:
             from hai_agents_local.sessions import LocalAsyncSessionsClient
 
-            self._sessions = LocalAsyncSessionsClient(client_wrapper=self._client_wrapper)
+            self._sessions = LocalAsyncSessionsClient(
+                client_wrapper=self._client_wrapper, runtime=self.local_runtime, auto_bridges=self._auto_bridges
+            )
         return self._sessions

src/hai_agents_local/bridge.py

@@ -17,7 +17,8 @@
 import httpx
 
 from .config import default_base_url
-from .errors import RateLimitedError, SessionNotFoundError
+from .errors import ChannelClosedError, RateLimitedError, SessionNotFoundError
+from .runtime import identity
 from .transport import Command, CommandExchange, Json, deserialize_args, serialize_result
 
 logger = logging.getLogger(__name__)
@@ -61,6 +62,8 @@ class LocalBridge(ABC, Generic[DriverT]):
     environment_kind: ClassVar[str]
     startup_hint: ClassVar[str | None] = None
     """Appended to the manager's not-ready timeout error; names the common cause of a hung startup."""
+    verify_runtime: bool = False
+    """Reject responses not HMAC-proven with api_key; requires api_key to be the local runtime's token string."""
 
     def __init__(
         self,
@@ -110,13 +113,23 @@ def request_stop(self) -> None:
         """Signal the poll loop to stop; safe to call from a signal handler."""
         self._stop_event.set()
 
+    async def interrupt_driver(self) -> None:
+        """Stop run-owned work; drivers without owned processes need no special action."""
+
     async def run(self) -> None:
         """Serve commands until stopped; raises AuthError on a bad key."""
         # An asyncio.Event binds to the loop it is first awaited on; a restarted bridge runs on a new loop.
         self._stop_event = asyncio.Event()
+        options: dict[str, Any] = {
+            "headers": {"Accept": "application/json"},
+            "auth": _BearerAuth(self.api_key),
+            "follow_redirects": True,
+        }
         try:
-            async with httpx.AsyncClient(
-                headers={"Accept": "application/json"}, auth=_BearerAuth(self.api_key), follow_redirects=True
+            async with (
+                identity.async_http_client(self.api_key, **options)
+                if self.verify_runtime
+                else httpx.AsyncClient(**options)
             ) as client:
                 exchange = CommandExchange(client, self.base_url)
                 if not await self._open_channel(exchange):
@@ -195,6 +208,9 @@ async def _poll_loop(self, exchange: CommandExchange) -> None:
                 ):
                     # Instant empty polls are paced so a misbehaving server cannot cause a busy loop.
                     break
+            except ChannelClosedError:
+                logger.info("channel %s closed; the session ended", self.session_id)
+                return
             except SessionNotFoundError:
                 # Channel was garbage-collected server-side; recreate on the next iteration so
                 # rate limits and transient errors during recreation hit the handlers below.
@@ -256,7 +272,19 @@ async def _process_commands(self, exchange: CommandExchange, commands: list[Comm
                 self._results.move_to_end(cmd.command_uid)
                 result, error = self._results[cmd.command_uid]
             else:
-                result, error = await asyncio.to_thread(self._dispatch, cmd.name, cmd.args)
+                dispatch = asyncio.create_task(asyncio.to_thread(self._dispatch, cmd.name, cmd.args))
+                stopped = asyncio.create_task(self._stop_event.wait())
+                try:
+                    await asyncio.wait((dispatch, stopped), return_when=asyncio.FIRST_COMPLETED)
+                    if self._stop_event.is_set():
+                        await self.interrupt_driver()
+                        await dispatch
+                        return
+                    result, error = await dispatch
+                finally:
+                    stopped.cancel()
+                    with contextlib.suppress(asyncio.CancelledError):
+                        await stopped
                 self._results[cmd.command_uid] = (result, error)
                 while len(self._results) > RESULT_CACHE_SIZE:
                     self._results.popitem(last=False)

src/hai_agents_local/browser.py

@@ -117,6 +117,6 @@ def _ensure_local_chrome(self) -> None:
 
 def _debugger_listening(port: int) -> bool:
     try:
-        return httpx.get(f"http://127.0.0.1:{port}/json/version", timeout=2.0).status_code == 200
+        return httpx.get(f"http://127.0.0.1:{port}/json/version", timeout=2.0, trust_env=False).status_code == 200
     except httpx.HTTPError:
         return False

src/hai_agents_local/config.py

@@ -1,4 +1,4 @@
-"""Environment variables read by hai_agents.local."""
+"""Environment variables read by hai_agents_local."""
 
 from __future__ import annotations
 

src/hai_agents_local/desktop.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+import asyncio
 import sys
+import threading
 from typing import TYPE_CHECKING, Literal
 
 from .bridge import LocalBridge, TokenSource
@@ -28,6 +30,7 @@ def ensure_macos_input_permissions(prompt: bool = True) -> None:
     from ApplicationServices import AXIsProcessTrustedWithOptions, kAXTrustedCheckOptionPrompt
     from Quartz import CGPreflightScreenCaptureAccess, CGRequestScreenCaptureAccess
 
+    prompt = prompt and threading.current_thread() is threading.main_thread()
     missing = []
     if not AXIsProcessTrustedWithOptions({kAXTrustedCheckOptionPrompt: prompt}):
         missing.append("Accessibility (moves the mouse and types)")
@@ -98,6 +101,10 @@ def create_driver(self) -> DesktopDriverInterface:
             quality=self.quality,
         )
 
+    async def interrupt_driver(self) -> None:
+        if self._driver is not None:
+            await asyncio.to_thread(self._driver.close)
+
     def driver_interface(self) -> type:
         # Runtime import: hai-drivers is absent unless installed with hai-agents[desktop].
         from hai_drivers.desktop.interface import DesktopDriverInterface