1"""Small JSON-lines client for the tinyagent helper."""23from __future__ import annotations45import asyncio6import base647import json8import logging9import os10from typing import Any, Awaitable, Callable111213class TinyRPC:14 def __init__(self, relay_url: str, key_env: str = "TINY_PRIVATE_KEY", cli_path: str = "tinyagent"):15 self.relay_url = relay_url16 self.key_env = key_env17 self.cli_path = cli_path18 self._proc: asyncio.subprocess.Process | None = None19 self._reader_task: asyncio.Task | None = None20 self._next_id = 021 self._pending: dict[int, asyncio.Future] = {}22 self._subscriptions: dict[str, Callable[[dict], Awaitable[None]]] = {}23 self._callback_tasks: set[asyncio.Task] = set()24 self._ready: dict[str, asyncio.Event] = {}25 self._write_lock = asyncio.Lock()2627 async def connect(self) -> None:28 if self._proc and self._proc.returncode is None:29 return30 env = dict(os.environ)31 # The key is inherited by the helper, never put in arguments or logs.32 self._proc = await asyncio.create_subprocess_exec(33 self.cli_path, "rpc", "--relay", self.relay_url, "--key-env", self.key_env,34 stdin=asyncio.subprocess.PIPE, stdout=asyncio.subprocess.PIPE,35 stderr=asyncio.subprocess.DEVNULL, env=env, limit=64 * 1024 * 1024,36 )37 self._reader_task = asyncio.create_task(self._read_loop())3839 async def close(self) -> None:40 if self._reader_task:41 self._reader_task.cancel()42 await asyncio.gather(self._reader_task, return_exceptions=True)43 for task in self._callback_tasks:44 task.cancel()45 await asyncio.gather(*self._callback_tasks, return_exceptions=True)46 self._callback_tasks.clear()47 if self._proc and self._proc.returncode is None:48 self._proc.terminate()49 await self._proc.wait()50 error = RuntimeError("tiny rpc closed")51 for future in self._pending.values():52 if not future.done():53 future.set_exception(error)54 self._pending.clear()55 self._proc = None5657 async def _read_loop(self) -> None:58 assert self._proc and self._proc.stdout59 try:60 async for line in self._proc.stdout:61 try:62 msg = json.loads(line)63 except (ValueError, UnicodeDecodeError):64 continue65 if msg.get("event") == "event":66 callback = self._subscriptions.get(str(msg.get("subscription", "")))67 if callback:68 task = asyncio.create_task(callback(msg.get("data") or {}))69 self._callback_tasks.add(task)70 task.add_done_callback(self._callback_done)71 continue72 if msg.get("event") == "connected":73 ready = self._ready.get(str(msg.get("subscription", "")))74 if ready:75 ready.set()76 continue77 request_id = msg.get("id")78 if isinstance(request_id, int) and request_id in self._pending:79 future = self._pending.pop(request_id)80 if future.done():81 continue82 if "error" in msg:83 future.set_exception(RuntimeError(str((msg["error"] or {}).get("message", "rpc error"))))84 else:85 future.set_result(msg.get("result"))86 finally:87 error = RuntimeError("tiny rpc exited")88 for future in self._pending.values():89 if not future.done():90 future.set_exception(error)91 self._pending.clear()9293 def _callback_done(self, task):94 self._callback_tasks.discard(task)95 if not task.cancelled() and task.exception():96 logging.getLogger(__name__).error("Tiny event handling failed", exc_info=task.exception())9798 async def call(self, method: str, params: dict[str, Any] | None = None, timeout: float = 30) -> Any:99 await self.connect()100 if not self._proc or not self._proc.stdin:101 raise RuntimeError("tiny rpc is not connected")102 self._next_id += 1103 request_id = self._next_id104 loop = asyncio.get_running_loop()105 future = loop.create_future()106 self._pending[request_id] = future107 payload = json.dumps({"id": request_id, "method": method, "params": params or {}}, separators=(",", ":"))108 try:109 async with self._write_lock:110 self._proc.stdin.write((payload + "\n").encode())111 await self._proc.stdin.drain()112 return await asyncio.wait_for(future, timeout)113 finally:114 self._pending.pop(request_id, None)115116 async def publish(self, event: dict) -> dict:117 return await self.call("publish", {"event": event})118119 async def query(self, filter: dict) -> list[dict]:120 return await self.call("query", {"filter": filter})121122 async def subscribe(self, subscription: str, filter: dict, callback: Callable[[dict], Awaitable[None]]) -> Any:123 self._subscriptions[subscription] = callback124 self._ready[subscription] = asyncio.Event()125 result = await self.call("subscribe", {"subscription": subscription, "filter": filter})126 await asyncio.wait_for(self._ready[subscription].wait(), 30)127 return result128129 async def unsubscribe(self, subscription: str) -> Any:130 self._subscriptions.pop(subscription, None)131 self._ready.pop(subscription, None)132 return await self.call("unsubscribe", {"subscription": subscription})133134 async def upload(self, room: str, filename: str, media_type: str, data: bytes) -> dict:135 return await self.call("upload", {"room": room, "filename": filename, "type": media_type,136 "data": base64.b64encode(data).decode()})137138 async def download(self, path: str) -> tuple[bytes, str]:139 result = await self.call("download", {"path": path})140 return base64.b64decode(result["data"]), result.get("type", "application/octet-stream")