from __future__ import annotations import asyncio import json import os import uuid from collections.abc import AsyncIterator from pathlib import Path from typing import Any, Callable _JSONRPC_VERSION = "2.0" _NANOBOT_API_STREAM_LIMIT = 2 * 1024 * 1024 NANOBOT_API_SOCKET = Path( os.getenv("NANOBOT_API_SOCKET", str(Path.home() / ".nanobot" / "api.sock")) ).expanduser() class _NanobotApiError(RuntimeError): def __init__(self, code: int, message: str) -> None: super().__init__(message) self.code = code def _jsonrpc_request(request_id: str, method: str, params: dict[str, Any]) -> dict[str, Any]: return { "jsonrpc": _JSONRPC_VERSION, "id": request_id, "method": method, "params": params, } def _jsonrpc_notification(method: str, params: dict[str, Any]) -> dict[str, Any]: return { "jsonrpc": _JSONRPC_VERSION, "method": method, "params": params, } async def _open_nanobot_api_socket( socket_path: Path = NANOBOT_API_SOCKET, *, stream_limit: int = _NANOBOT_API_STREAM_LIMIT, ) -> tuple[asyncio.StreamReader, asyncio.StreamWriter]: if not socket_path.exists(): raise RuntimeError( f"Nanobot API socket not found at {socket_path}. " "Enable channels.api and start `nanobot gateway`." ) try: return await asyncio.open_unix_connection( path=str(socket_path), limit=stream_limit, ) except OSError as exc: raise RuntimeError(f"failed to connect to Nanobot API socket: {exc}") from exc async def _send_nanobot_api_request( method: str, params: dict[str, Any], *, timeout_seconds: float, socket_path: Path = NANOBOT_API_SOCKET, stream_limit: int = _NANOBOT_API_STREAM_LIMIT, ) -> Any: request_id = str(uuid.uuid4()) reader, writer = await _open_nanobot_api_socket(socket_path, stream_limit=stream_limit) try: writer.write( ( json.dumps(_jsonrpc_request(request_id, method, params), ensure_ascii=False) + "\n" ).encode("utf-8") ) await writer.drain() loop = asyncio.get_running_loop() deadline = loop.time() + timeout_seconds while True: remaining = deadline - loop.time() if remaining <= 0: raise RuntimeError(f"timed out waiting for Nanobot API response to {method}") try: line = await asyncio.wait_for(reader.readline(), timeout=remaining) except ValueError as exc: raise RuntimeError( "Nanobot API response exceeded the configured stream limit" ) from exc if not line: raise RuntimeError("Nanobot API socket closed before responding") try: message = json.loads(line.decode("utf-8", errors="replace")) except json.JSONDecodeError: continue if not isinstance(message, dict): continue if message.get("jsonrpc") != _JSONRPC_VERSION: continue if "method" in message: continue if str(message.get("id", "")).strip() != request_id: continue if "error" in message: error = message.get("error", {}) if isinstance(error, dict): raise _NanobotApiError( int(error.get("code", -32000)), str(error.get("message", "unknown Nanobot API error")), ) raise _NanobotApiError(-32000, str(error)) return message.get("result") finally: writer.close() await writer.wait_closed() async def _run_nanobot_agent_turn( content: str, *, chat_id: str, metadata: dict[str, Any] | None = None, timeout_seconds: float = 90.0, notification_to_event: Callable[[dict[str, Any]], dict[str, Any] | None], socket_path: Path = NANOBOT_API_SOCKET, stream_limit: int = _NANOBOT_API_STREAM_LIMIT, ) -> dict[str, Any]: final_message = "" saw_activity = False async for typed_event in _stream_nanobot_agent_turn( content, chat_id=chat_id, metadata=metadata, timeout_seconds=timeout_seconds, notification_to_event=notification_to_event, socket_path=socket_path, stream_limit=stream_limit, ): event_type = typed_event.get("type") if event_type == "message": saw_activity = True if ( not bool(typed_event.get("is_progress", False)) and typed_event.get("role") == "nanobot" ): final_message = str(typed_event.get("content", "")) continue if event_type == "agent_state": saw_activity = True return { "status": "ok", "message": final_message, "completed": saw_activity, } async def _stream_nanobot_agent_turn( content: str, *, chat_id: str, metadata: dict[str, Any] | None = None, timeout_seconds: float = 60.0, idle_grace_seconds: float = 1.0, notification_to_event: Callable[[dict[str, Any]], dict[str, Any] | None], socket_path: Path = NANOBOT_API_SOCKET, stream_limit: int = _NANOBOT_API_STREAM_LIMIT, ) -> AsyncIterator[dict[str, Any]]: reader, writer = await _open_nanobot_api_socket(socket_path, stream_limit=stream_limit) saw_final_message = False saw_activity = False try: writer.write( ( json.dumps( _jsonrpc_notification( "message.send", { "content": content, "chat_id": chat_id, "metadata": metadata or {}, }, ), ensure_ascii=False, ) + "\n" ).encode("utf-8") ) await writer.drain() loop = asyncio.get_running_loop() deadline = loop.time() + timeout_seconds while True: remaining = deadline - loop.time() if remaining <= 0: raise RuntimeError("Timed out waiting for a Nanobot response.") timeout = min(remaining, idle_grace_seconds) if saw_final_message else remaining try: line = await asyncio.wait_for(reader.readline(), timeout=timeout) except asyncio.TimeoutError: if saw_final_message: break raise RuntimeError("Timed out waiting for Nanobot to start responding") except ValueError as exc: raise RuntimeError( "Nanobot API response exceeded the configured stream limit" ) from exc if not line: break try: obj = json.loads(line.decode("utf-8", errors="replace").strip()) except json.JSONDecodeError: continue if not isinstance(obj, dict): continue typed_event = notification_to_event(obj) if typed_event is None: continue event_type = typed_event.get("type") if event_type == "message": saw_activity = True if ( not bool(typed_event.get("is_progress", False)) and typed_event.get("role") == "nanobot" ): saw_final_message = True yield typed_event if ( saw_final_message and event_type == "agent_state" and str(typed_event.get("state", "")).strip().lower() == "idle" ): break if not saw_activity: raise RuntimeError("Timed out waiting for Nanobot to start responding") finally: writer.close() await writer.wait_closed() NanobotApiError = _NanobotApiError send_nanobot_api_request = _send_nanobot_api_request run_nanobot_agent_turn = _run_nanobot_agent_turn stream_nanobot_agent_turn = _stream_nanobot_agent_turn __all__ = [ "NANOBOT_API_SOCKET", "NanobotApiError", "run_nanobot_agent_turn", "send_nanobot_api_request", "stream_nanobot_agent_turn", ]