1"""Small JSON-lines client for the tinyagent helper."""23from __future__ import annotations45import asyncio6import base647import json8import logging9import os10from typing import Any, Awaitable, Callable1112# The helper receives only what it needs: a lookup path, a home for TLS and13# temp files, its signing key and the proxy settings of the gateway.14_HELPER_ENV = ("PATH", "HOME", "TMPDIR", "HTTP_PROXY", "HTTPS_PROXY", "NO_PROXY",15 "http_proxy", "https_proxy", "no_proxy")161718def scoped_secret(name: str, default: str = "", *, strict: bool = False) -> str:19 """Resolve a TINY_* setting through Hermes's profile secret scope.2021 A multiplexed gateway serves several profiles from one process and keeps each22 profile's ``.env`` in a context-local scope rather than in ``os.environ``; reading23 through ``get_secret`` hands a profile its own key. Outside Hermes, or on a Hermes24 without the scope API, the process environment is the source. ``strict`` lets the25 fail-closed error of an unscoped read in multiplex mode propagate; passive probes26 use the default instead.27 """28 try:29 from agent.secret_scope import get_secret30 except ImportError:31 return os.environ.get(name, default)32 try:33 value = get_secret(name, default)34 except RuntimeError:35 if strict:36 raise37 return default38 return default if value is None else str(value)394041class TinyRPC:42 def __init__(self, relay_url: str, key_env: str = "TINY_PRIVATE_KEY", cli_path: str = "tinyagent",43 on_exit: Callable[["TinyRPC"], Any] | None = None):44 self.relay_url = relay_url45 self.key_env = key_env46 self.cli_path = cli_path47 self.on_exit = on_exit48 self._proc: asyncio.subprocess.Process | None = None49 self._reader_task: asyncio.Task | None = None50 self._next_id = 051 self._pending: dict[int, asyncio.Future] = {}52 self._subscriptions: dict[str, Callable[[dict], Awaitable[None]]] = {}53 self._callback_tasks: set[asyncio.Task] = set()54 self._ready: dict[str, asyncio.Event] = {}55 self._write_lock = asyncio.Lock()56 self._closing = False57 self._exited = False5859 def helper_env(self) -> dict[str, str]:60 """Return the scoped environment handed to the helper process.6162 The signing key is resolved through the active profile scope and reaches the63 helper only through its environment variable."""64 env = {name: value for name, value in os.environ.items() if name in _HELPER_ENV}65 key = scoped_secret(self.key_env, "", strict=True)66 if key:67 env[self.key_env] = key68 return env697071 def alive(self) -> bool:72 return self._proc is not None and self._proc.returncode is None and not self._exited7374 async def connect(self) -> None:75 if self._proc is not None:76 if self.alive:77 return78 # A dead helper is not respawned in place: its subscriptions are gone,79 # so the owner must rebuild the session through on_exit.80 raise RuntimeError("tiny rpc exited")81 # The key is inherited by the helper, never put in arguments or logs.82 self._proc = await asyncio.create_subprocess_exec(83 self.cli_path, "rpc", "--relay", self.relay_url, "--key-env", self.key_env,84 stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE,85 stderr=asyncio.subprocess.DEVNULL, env=self.helper_env(), limit=64 * 1024 * 1024,86 )87 self._reader_task = asyncio.create_task(self._read_loop())8889 async def close(self) -> None:90 self._closing = True91 if self._reader_task:92 self._reader_task.cancel()93 await asyncio.gather(self._reader_task, return_exceptions=True)94 if not self._exited:95 # A deliberate close ends in-flight event handling; turns started by a helper96 # that died on its own run to completion and are redelivered if they fail.97 for task in self._callback_tasks:98 task.cancel()99 await asyncio.gather(*self._callback_tasks, return_exceptions=True)100 self._callback_tasks.clear()101 if self._proc and self._proc.returncode is None:102 self._proc.terminate()103 await self._proc.wait()104 self._fail_pending(RuntimeError("tiny rpc closed"))105 self._proc = None106107 def _fail_pending(self, error: Exception) -> None:108 for future in self._pending.values():109 if not future.done():110 future.set_exception(error)111 self._pending.clear()112113 async def _read_loop(self) -> None:114 assert self._proc and self._proc.stdout115 died = False116 try:117 try:118 await self._pump()119 except asyncio.CancelledError:120 raise121 except Exception:122 logging.getLogger(__name__).exception("Tiny helper stream failed")123 # EOF or a broken pipe: the helper exited (or was closed) rather than being cancelled.124 died = not self._closing125 finally:126 self._fail_pending(RuntimeError("tiny rpc exited"))127 if died and not self._exited:128 # The helper went away on its own; tell the owner exactly once.129 self._exited = True130 if self.on_exit is not None:131 try:132 result = self.on_exit(self)133 if asyncio.iscoroutine(result):134 await result135 except Exception:136 logging.getLogger(__name__).exception("Tiny helper exit handler failed")137138 async def _pump(self) -> None:139 assert self._proc and self._proc.stdout140 async for line in self._proc.stdout:141 try:142 msg = json.loads(line)143 except (ValueError, UnicodeDecodeError):144 continue145 if msg.get("event") == "event":146 callback = self._subscriptions.get(str(msg.get("subscription", "")))147 if callback:148 task = asyncio.create_task(callback(msg.get("data") or {}))149 self._callback_tasks.add(task)150 task.add_done_callback(self._callback_done)151 continue152 if msg.get("event") == "connected":153 ready = self._ready.get(str(msg.get("subscription", "")))154 if ready:155 ready.set()156 continue157 request_id = msg.get("id")158 if isinstance(request_id, int) and request_id in self._pending:159 future = self._pending.pop(request_id)160 if future.done():161 continue162 if "error" in msg:163 future.set_exception(RuntimeError(str((msg["error"] or {}).get("message", "rpc error"))))164 else:165 future.set_result(msg.get("result"))166167 def _callback_done(self, task):168 self._callback_tasks.discard(task)169 if not task.cancelled() and task.exception():170 logging.getLogger(__name__).error("Tiny event handling failed", exc_info=task.exception())171172 async def call(self, method: str, params: dict[str, Any] | None = None, timeout: float = 30) -> Any:173 await self.connect()174 if not self._proc or not self._proc.stdin:175 raise RuntimeError("tiny rpc is not connected")176 self._next_id += 1177 request_id = self._next_id178 loop = asyncio.get_running_loop()179 future = loop.create_future()180 self._pending[request_id] = future181 payload = json.dumps({"id": request_id, "method": method, "params": params or {}}, separators=(",", ":"))182 try:183 async with self._write_lock:184 self._proc.stdin.write((payload + "\n").encode())185 await self._proc.stdin.drain()186 return await asyncio.wait_for(future, timeout)187 finally:188 self._pending.pop(request_id, None)189190 async def publish(self, event: dict) -> dict:191 return await self.call("publish", {"event": event})192193 async def query(self, filter: dict) -> list[dict]:194 return await self.call("query", {"filter": filter})195196 async def subscribe(self, subscription: str, filter: dict, callback: Callable[[dict], Awaitable[None]]) -> Any:197 self._subscriptions[subscription] = callback198 self._ready[subscription] = asyncio.Event()199 result = await self.call("subscribe", {"subscription": subscription, "filter": filter})200 await asyncio.wait_for(self._ready[subscription].wait(), 30)201 return result202203 async def unsubscribe(self, subscription: str) -> Any:204 self._subscriptions.pop(subscription, None)205 self._ready.pop(subscription, None)206 return await self.call("unsubscribe", {"subscription": subscription})207208 async def upload(self, room: str, filename: str, media_type: str, data: bytes) -> dict:209 return await self.call("upload", {"room": room, "filename": filename, "type": media_type,210 "data": base64.b64encode(data).decode()})211212 async def download(self, path: str) -> tuple[bytes, str]:213 result = await self.call("download", {"path": path})214 return base64.b64decode(result["data"]), result.get("type", "application/octet-stream")