from __future__ import annotations import asyncio import json import re import time from dataclasses import dataclass, field from pathlib import Path from typing import Any, Callable, Protocol from urllib.parse import urlparse import httpx from .downstream import DownstreamProtocolError, make_timeout from .hermes import HermesAdapter from .models import ( DiscoveredNode, HermesRun, HermesRunRequest, NodeView, NodeKillSwitchResponse, PairNodeRequest, PairingReceipt, TrustedNode, new_id, ) MDNS_SERVICE_TYPE = "_omni-node._tcp.local." MDNS_PROTOCOL_VERSION = "1.0" TRUST_STORE_VERSION = 1 class PairingClientProtocol(Protocol): async def claim(self, request: PairNodeRequest) -> PairingReceipt: ... class DiscoveryProtocol(Protocol): async def discover(self) -> list[DiscoveredNode]: ... class AmbiguousNodeTarget(ValueError): def __init__(self, node_ids: list[str]) -> None: super().__init__("Prompt names more than one trusted node") self.node_ids = node_ids class TrustStore: """Local allowlist. Pairing codes are intentionally never persisted.""" def __init__(self, path: Path) -> None: self.path = Path(path) self._lock = asyncio.Lock() self._nodes: dict[str, TrustedNode] | None = None self._kill_switch_engaged = False async def load(self) -> list[TrustedNode]: async with self._lock: if self._nodes is None: self._nodes, self._kill_switch_engaged = await asyncio.to_thread( self._read_sync ) return list(self._nodes.values()) async def put(self, node: TrustedNode) -> None: await self.load() async with self._lock: assert self._nodes is not None self._nodes[node.node_id] = node await asyncio.to_thread( self._write_sync, self._nodes, self._kill_switch_engaged ) async def remove(self, node_id: str) -> bool: await self.load() async with self._lock: assert self._nodes is not None removed = self._nodes.pop(node_id, None) is not None if removed: await asyncio.to_thread( self._write_sync, self._nodes, self._kill_switch_engaged ) return removed async def kill_switch_engaged(self) -> bool: await self.load() return self._kill_switch_engaged async def set_kill_switch(self, engaged: bool) -> None: await self.load() async with self._lock: assert self._nodes is not None self._kill_switch_engaged = engaged await asyncio.to_thread(self._write_sync, self._nodes, engaged) def _read_sync(self) -> tuple[dict[str, TrustedNode], bool]: if not self.path.exists(): return {}, False payload = json.loads(self.path.read_text(encoding="utf-8")) if payload.get("version") != TRUST_STORE_VERSION: raise RuntimeError("Unsupported trusted-node store version") nodes = [TrustedNode.model_validate(item) for item in payload.get("nodes", [])] return {node.node_id: node for node in nodes}, bool( payload.get("kill_switch_engaged", False) ) def _write_sync(self, nodes: dict[str, TrustedNode], kill_switch_engaged: bool) -> None: self.path.parent.mkdir(parents=True, exist_ok=True) payload = { "version": TRUST_STORE_VERSION, "kill_switch_engaged": kill_switch_engaged, "nodes": [node.model_dump(mode="json") for node in nodes.values()], } temporary = self.path.with_suffix(self.path.suffix + ".tmp") temporary.write_text( json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8" ) temporary.replace(self.path) class MdnsDiscovery: """Bounded mDNS service lookup. It never scans addresses or pairs a node.""" def __init__(self, timeout_s: float = 0.75) -> None: self.timeout_s = timeout_s async def discover(self) -> list[DiscoveredNode]: return await asyncio.to_thread(self._discover_sync) def _discover_sync(self) -> list[DiscoveredNode]: try: from zeroconf import ServiceBrowser, ServiceListener, Zeroconf except ImportError: return [] found: dict[str, DiscoveredNode] = {} class Listener(ServiceListener): def add_service(self, zeroconf, service_type: str, name: str) -> None: info = zeroconf.get_service_info(service_type, name, timeout=500) candidate = discovered_node_from_mdns_info( info, service_type=service_type, service_name=name ) if candidate is not None: found[candidate.node_id] = candidate def update_service(self, zeroconf, service_type: str, name: str) -> None: self.add_service(zeroconf, service_type, name) def remove_service(self, zeroconf, service_type: str, name: str) -> None: return None zeroconf = Zeroconf() browser = ServiceBrowser(zeroconf, MDNS_SERVICE_TYPE, Listener()) try: time.sleep(self.timeout_s) finally: browser.cancel() zeroconf.close() return list(found.values()) def discovered_node_from_mdns_info( info: Any, *, service_type: str = MDNS_SERVICE_TYPE, service_name: str, ) -> DiscoveredNode | None: """Parse the minimal advertisement without treating TXT as trust data.""" if info is None or service_type != MDNS_SERVICE_TYPE: return None try: properties = { ( key.decode("utf-8", "replace") if isinstance(key, bytes) else str(key) ): ( value.decode("utf-8", "replace") if isinstance(value, bytes) else str(value) ) for key, value in info.properties.items() } port = int(info.port) advertised_port = int(properties.get("api_port", "")) addresses = list(info.parsed_addresses()) except (AttributeError, TypeError, ValueError): return None if ( properties.get("protocol_version") != MDNS_PROTOCOL_VERSION or not 1 <= port <= 65_535 or advertised_port != port or not addresses ): return None node_id = properties.get("node_id", "").strip() if not node_id: return None host = addresses[0] if ":" in host: host = f"[{host}]" try: return DiscoveredNode( node_id=node_id, name=service_name.removesuffix(f".{service_type}"), endpoint=f"https://{host}:{port}", # Fingerprint and capabilities must come from authenticated, # out-of-band pairing/capability exchange, never from mDNS TXT. certificate_fingerprint_sha256=None, capabilities=set(), ) except ValueError: return None class PairingClient: """Claims a node only through explicit one-time-code + TLS client identity.""" def __init__( self, *, conductor_id: str, connect_timeout_s: float = 1.0, request_timeout_s: float = 10.0, ) -> None: self.conductor_id = conductor_id self.connect_timeout_s = connect_timeout_s self.request_timeout_s = request_timeout_s async def claim(self, request: PairNodeRequest) -> PairingReceipt: endpoint = _validated_https_origin(request.endpoint) ca_path, cert_path, key_path = _validated_tls_paths(request) async with httpx.AsyncClient( verify=str(ca_path), cert=(str(cert_path), str(key_path)), timeout=make_timeout(self.connect_timeout_s, self.request_timeout_s), follow_redirects=False, ) as client: response = await client.post( f"{endpoint}/v1/pairing/claim", json={ "one_time_code": request.one_time_code.get_secret_value(), "conductor_id": self.conductor_id, "requested_permissions": sorted(request.requested_permissions), }, headers={"Cache-Control": "no-store"}, ) response.raise_for_status() receipt = PairingReceipt.model_validate(response.json()) if receipt.certificate_fingerprint_sha256.lower() != ( request.expected_fingerprint_sha256.lower() ): raise DownstreamProtocolError("Pairing certificate fingerprint mismatch") if not _same_origin(endpoint, receipt.endpoint): raise DownstreamProtocolError("Pairing response changed the node origin") return receipt class RemoteNodeRegistry: def __init__( self, *, trust_store: TrustStore, pairing: PairingClientProtocol, discovery: DiscoveryProtocol, enabled: bool = True, discovery_enabled: bool = True, ) -> None: self.trust_store = trust_store self.pairing = pairing self.discovery = discovery self.enabled = enabled self.discovery_enabled = discovery_enabled async def initialize(self) -> None: if self.enabled: await self.trust_store.load() async def discovered(self) -> list[DiscoveredNode]: if not self.enabled or not self.discovery_enabled: return [] return await self.discovery.discover() async def pair(self, request: PairNodeRequest) -> TrustedNode: if not self.enabled: raise ValueError("Remote nodes are disabled by OMNI_REMOTE_NODES_ENABLED") ca_path, cert_path, key_path = _validated_tls_paths(request) receipt = await self.pairing.claim(request) node = TrustedNode( node_id=receipt.node_id, name=receipt.name, endpoint=_validated_https_origin(receipt.endpoint), certificate_fingerprint_sha256=receipt.certificate_fingerprint_sha256.lower(), capabilities=receipt.capabilities, granted_permissions=receipt.granted_permissions, ca_cert_path=ca_path, client_cert_path=cert_path, client_key_path=key_path, ) await self.trust_store.put(node) return node async def unpair(self, node_id: str) -> bool: return await self.trust_store.remove(node_id) async def trusted(self) -> list[TrustedNode]: if not self.enabled: return [] return await self.trust_store.load() async def kill_switch_engaged(self) -> bool: return await self.trust_store.kill_switch_engaged() async def set_kill_switch(self, engaged: bool) -> None: await self.trust_store.set_kill_switch(engaged) async def views( self, metrics: dict[str, tuple[str, float | None]] | None = None ) -> list[NodeView]: metrics = metrics or {} return [ NodeView( node_id=node.node_id, name=node.name, endpoint=node.endpoint, capabilities=node.capabilities, granted_permissions=node.granted_permissions, enabled=node.enabled, status=metrics.get(node.node_id, ("unknown", None))[0], # type: ignore[arg-type] latency_ms=metrics.get(node.node_id, ("unknown", None))[1], ) for node in await self.trusted() ] @dataclass(slots=True) class _Provider: provider_id: str adapter: Any capabilities: set[str] permissions: set[str] remote: bool name: str status: str = "unknown" latency_ms: float | None = None last_probe: float = 0.0 @dataclass(slots=True) class _RunBinding: provider_id: str downstream_run_id: str session_id: str class HermesProviderPool: """Trusted-provider selection with health, affinity and submit-time failover.""" # The direct conductor -> Hermes connection has no run-bound sidecar # ticket. Privileged resources therefore require a paired provider even # when that provider represents this same computer. LOCAL_CAPABILITIES = {"hermes.run", "code", "tools", "internet.search"} LOCAL_PERMISSIONS = {"tool.execute", "network.search"} PERMISSION_CAPABILITIES = { "desktop.observe": "computer.desktop", "desktop.control": "computer.desktop", "files.read": "files", "files.write": "files", "tool.execute": "tools", "webcam.read": "webcam", "network.search": "internet.search", } def __init__( self, *, local: Any, registry: RemoteNodeRegistry, health_ttl_s: float = 5.0, connect_timeout_s: float = 0.75, request_timeout_s: float = 20.0, stream_timeout_s: float = 300.0, remote_adapter_factory: Callable[[TrustedNode], Any] | None = None, ) -> None: self.local = local self.registry = registry self.health_ttl_s = health_ttl_s self.connect_timeout_s = connect_timeout_s self.request_timeout_s = request_timeout_s self.stream_timeout_s = stream_timeout_s self.remote_adapter_factory = remote_adapter_factory self._providers: dict[str, _Provider] = { "local": _Provider( "local", local, set(self.LOCAL_CAPABILITIES), set(self.LOCAL_PERMISSIONS), False, "Local computer", ) } self._remote_clients: dict[str, httpx.AsyncClient] = {} self._factory_adapters: dict[str, Any] = {} self._affinity: dict[str, str] = {} self._runs: dict[str, _RunBinding] = {} self._lock = asyncio.Lock() async def _sync_remote_providers(self) -> None: await self.registry.initialize() trusted = ( {} if await self.registry.kill_switch_engaged() else {node.node_id: node for node in await self.registry.trusted() if node.enabled} ) async with self._lock: stale = [ provider_id for provider_id, provider in self._providers.items() if provider.remote and provider_id not in trusted ] for provider_id in stale: self._providers.pop(provider_id, None) client = self._remote_clients.pop(provider_id, None) if client is not None: await client.aclose() adapter = self._factory_adapters.pop(provider_id, None) if adapter is not None: await adapter.close() for node_id, node in trusted.items(): if node_id in self._providers: self._providers[node_id].capabilities = set(node.capabilities) self._providers[node_id].permissions = set(node.granted_permissions) self._providers[node_id].name = node.name continue if self.remote_adapter_factory is not None: adapter = self.remote_adapter_factory(node) self._factory_adapters[node_id] = adapter else: client = httpx.AsyncClient( verify=str(node.ca_cert_path), cert=(str(node.client_cert_path), str(node.client_key_path)), timeout=make_timeout( self.connect_timeout_s, self.request_timeout_s, self.stream_timeout_s, ), follow_redirects=False, ) self._remote_clients[node_id] = client adapter = HermesAdapter(node.endpoint, client=client) self._providers[node_id] = _Provider( provider_id=node_id, adapter=adapter, capabilities=set(node.capabilities), permissions=set(node.granted_permissions), remote=True, name=node.name, ) async def refresh_health(self, *, force: bool = False) -> None: await self._sync_remote_providers() now = time.monotonic() candidates = [ provider for provider in self._providers.values() if force or now - provider.last_probe >= self.health_ttl_s ] async def probe(provider: _Provider) -> None: started = time.perf_counter() try: result = await provider.adapter.health() provider.status = "ready" if result.get("status") == "ready" else "unavailable" except Exception: provider.status = "unavailable" provider.latency_ms = (time.perf_counter() - started) * 1000 provider.last_probe = time.monotonic() await asyncio.gather(*(probe(provider) for provider in candidates)) async def resolve_target(self, prompt: str) -> str | None: """Resolve only exact trusted node IDs/names; never fuzzy-match a machine.""" matches: set[str] = set() for node in await self.registry.trusted(): if not node.enabled: continue for candidate in {node.node_id, node.name}: if re.search( rf"(? 1: raise AmbiguousNodeTarget(sorted(matches)) return next(iter(matches), None) async def create_run(self, request: HermesRunRequest, *, idempotency_key: str) -> HermesRun: await self.refresh_health() required = {"hermes.run"} if request.needs_vision: required.add("context.vision") required.update(request.required_capabilities) requested_permissions = set(request.requested_permissions) required.update( self.PERMISSION_CAPABILITIES[permission] for permission in requested_permissions if permission in self.PERMISSION_CAPABILITIES ) target = request.target_node.casefold() if request.target_node else None eligible = [ provider for provider in self._providers.values() if provider.status == "ready" and required.issubset(provider.capabilities) and requested_permissions.issubset(provider.permissions) and ( target is None or provider.provider_id.casefold() == target or provider.name.casefold() == target ) ] affine = self._affinity.get(request.session_id) eligible.sort( key=lambda provider: ( 0 if provider.provider_id == affine else 1, provider.latency_ms if provider.latency_ms is not None else float("inf"), provider.provider_id, ) ) if not eligible: raise DownstreamProtocolError("No healthy trusted Hermes provider is eligible") errors: list[str] = [] for provider in eligible: try: run = await provider.adapter.create_run( request, idempotency_key=idempotency_key ) except Exception as exc: provider.status = "unavailable" # A submit-time transport failure must not remain hidden behind # the normal health TTL. The next idempotent request should # immediately probe this provider and allow a recovered node # back into the eligible set. provider.last_probe = 0.0 errors.append(f"{provider.provider_id}:{type(exc).__name__}") continue pool_run_id = new_id("run") self._runs[pool_run_id] = _RunBinding( provider_id=provider.provider_id, downstream_run_id=run.run_id, session_id=request.session_id, ) self._affinity[request.session_id] = provider.provider_id return HermesRun(run_id=pool_run_id, provider_id=provider.provider_id) raise DownstreamProtocolError( "All eligible Hermes providers rejected the idempotent run: " + ",".join(errors) ) async def events(self, run_id: str, *, after_event_id: str | None = None): binding = self._binding(run_id) provider = self._providers[binding.provider_id] async for event in provider.adapter.events( binding.downstream_run_id, after_event_id=after_event_id ): yield event async def stop(self, run_id: str) -> None: binding = self._binding(run_id) provider = self._providers.get(binding.provider_id) if provider is None: return await provider.adapter.stop(binding.downstream_run_id) async def approve(self, run_id: str, choice: str) -> None: binding = self._binding(run_id) provider = self._providers[binding.provider_id] await provider.adapter.approve(binding.downstream_run_id, choice) async def health(self) -> dict[str, Any]: await self.refresh_health() trusted_ids = {node.node_id for node in await self.registry.trusted()} return { "status": "ready" if any(p.status == "ready" for p in self._providers.values()) else "unavailable", "selection": "session-affinity-then-lowest-latency-eligible", "remote_kill_switch_engaged": await self.registry.kill_switch_engaged(), "providers": { provider.provider_id: { "status": provider.status, "latency_ms": provider.latency_ms, "capabilities": sorted(provider.capabilities), "permissions": sorted(provider.permissions), "trusted": not provider.remote or provider.provider_id in trusted_ids, } for provider in self._providers.values() }, } async def node_metrics(self) -> dict[str, tuple[str, float | None]]: await self.refresh_health() return { provider.provider_id: (provider.status, provider.latency_ms) for provider in self._providers.values() if provider.remote } async def set_kill_switch(self, engaged: bool) -> NodeKillSwitchResponse: cancelled: list[str] = [] if engaged: for pool_run_id, binding in list(self._runs.items()): provider = self._providers.get(binding.provider_id) if provider is None or not provider.remote: continue try: await provider.adapter.stop(binding.downstream_run_id) except Exception: pass cancelled.append(pool_run_id) self._runs.pop(pool_run_id, None) self._affinity = { session_id: provider_id for session_id, provider_id in self._affinity.items() if provider_id == "local" } await self.registry.set_kill_switch(engaged) await self._sync_remote_providers() return NodeKillSwitchResponse(engaged=engaged, cancelled_run_ids=cancelled) async def unpair(self, node_id: str) -> bool: provider = self._providers.get(node_id) if provider is not None and provider.remote: for pool_run_id, binding in list(self._runs.items()): if binding.provider_id != node_id: continue try: await provider.adapter.stop(binding.downstream_run_id) except Exception: pass self._runs.pop(pool_run_id, None) self._affinity = { session_id: provider_id for session_id, provider_id in self._affinity.items() if provider_id != node_id } removed = await self.registry.unpair(node_id) await self._sync_remote_providers() return removed def _binding(self, run_id: str) -> _RunBinding: binding = self._runs.get(run_id) if binding is None: raise DownstreamProtocolError("Unknown provider run ID") return binding async def close(self) -> None: await self.local.close() for client in list(self._remote_clients.values()): await client.aclose() self._remote_clients.clear() for adapter in list(self._factory_adapters.values()): await adapter.close() self._factory_adapters.clear() def _validated_https_origin(value: str) -> str: parsed = urlparse(value) if parsed.scheme != "https" or not parsed.hostname or parsed.username or parsed.password: raise ValueError("Remote nodes require an HTTPS endpoint without URL credentials") return value.rstrip("/") def _validated_tls_paths(request: PairNodeRequest) -> tuple[Path, Path, Path]: paths = tuple( Path(path).expanduser().resolve() for path in (request.ca_cert_path, request.client_cert_path, request.client_key_path) ) for path in paths: if not path.is_file(): raise ValueError(f"Required local TLS file does not exist: {path}") return paths # type: ignore[return-value] def _same_origin(left: str, right: str) -> bool: a, b = urlparse(left), urlparse(right) return (a.scheme, a.hostname, a.port or 443) == (b.scheme, b.hostname, b.port or 443)