"""Transport-only Unix-socket client for TAMQ integrations.""" from __future__ import annotations import asyncio import json from dataclasses import dataclass from pathlib import Path from typing import Any from .config import socket_path from .protocol import ( DELIVERY_RELIABILITY_CAPABILITY, IDEMPOTENT_SEND_CAPABILITY, PROTOCOL_VERSION, ) class TamqClientError(RuntimeError): """The local TAMQ service rejected or could not complete an operation.""" class TamqProtocolError(TamqClientError): """The service and client do not share a compatible protocol contract.""" class TamqTargetUnavailable(TamqClientError): """No unambiguous live endpoint owns the requested repository.""" @dataclass(frozen=True) class WakeRequest: lease_id: str target_repo: str prompt: str trigger_id: str source_repo: str = "coordination-engine" endpoint_id: str | None = None target_agent: str | None = None @dataclass(frozen=True) class WakeReceipt: lease_id: str trigger_id: str endpoint_id: str message_id: str state: str deduplicated: bool class TamqClient: """Small async client with no tmux/control-mode dependency.""" def __init__( self, path: Path | None = None, *, client_id: str = "coordination-engine", protocol: str = PROTOCOL_VERSION, timeout: float = 5.0, ): self.path = path or socket_path() self.client_id = client_id self.protocol = protocol self.timeout = timeout async def request(self, operation: str, **fields: Any) -> dict: payload = {"op": operation, "protocol": self.protocol, **fields} try: reader, writer = await asyncio.wait_for( asyncio.open_unix_connection(str(self.path), limit=1024 * 1024), timeout=self.timeout, ) try: writer.write((json.dumps(payload, sort_keys=True) + "\n").encode()) await asyncio.wait_for(writer.drain(), timeout=self.timeout) line = await asyncio.wait_for(reader.readline(), timeout=self.timeout) finally: writer.close() await writer.wait_closed() except (OSError, asyncio.TimeoutError) as exc: raise TamqClientError(f"TAMQ service unavailable: {exc}") from exc try: response = json.loads(line) except (json.JSONDecodeError, UnicodeDecodeError) as exc: raise TamqProtocolError("TAMQ returned an invalid response") from exc if not response.get("ok"): error = str(response.get("error") or "operation failed") if error.startswith("incompatible protocol"): raise TamqProtocolError(error) raise TamqClientError(error) return response async def negotiate(self) -> dict: response = await self.request("ping") server_protocol = str(response.get("protocol", "")) if server_protocol.split(".")[0] != self.protocol.split(".")[0]: raise TamqProtocolError( f"incompatible protocol: client {self.protocol}, server {server_protocol}" ) required = {DELIVERY_RELIABILITY_CAPABILITY, IDEMPOTENT_SEND_CAPABILITY} missing = required.difference(response.get("capabilities", [])) if missing: raise TamqProtocolError( "TAMQ lacks required capabilities: " + ", ".join(sorted(missing)) ) return response async def send( self, *, sender_repo: str, target_repo: str, body: str, idempotency_key: str, endpoint_id: str | None = None, correlation_id: str | None = None, metadata: dict | None = None, provenance: str = "coordination_engine", ) -> dict: fields: dict[str, Any] = { "sender_repo": sender_repo, "target_repo": target_repo, "body": body, "client_id": self.client_id, "idempotency_key": idempotency_key, "provenance": provenance, } if endpoint_id is not None: fields["endpoint_id"] = endpoint_id if correlation_id is not None: fields["correlation_id"] = correlation_id if metadata is not None: fields["metadata"] = metadata return await self.request("send", **fields) async def message(self, message_id: str) -> dict: return (await self.request("message", message_id=message_id))["message"] async def history( self, *, target_repo: str | None = None, state: str | None = None ) -> list[dict]: fields = { key: value for key, value in {"target_repo": target_repo, "state": state}.items() if value is not None } return (await self.request("history", **fields))["messages"] async def endpoints(self) -> list[dict]: return (await self.request("endpoints"))["endpoints"] async def acknowledge(self, message_id: str) -> dict: return await self.request("ack", message_id=message_id) async def retry(self, message_id: str) -> dict: return await self.request("retry", message_id=message_id) class CoordinationEngineAdapter: """Map coordination-engine wake requests onto durable TAMQ messages.""" def __init__(self, client: TamqClient | None = None): self.client = client or TamqClient() async def _resolve_endpoint(self, target_repo: str, requested: str | None) -> str: endpoints = await self.client.endpoints() matching = [ endpoint for endpoint in endpoints if target_repo in json.loads(endpoint["repos"]) and (requested is None or endpoint["endpoint_id"] == requested) ] if not matching: raise TamqTargetUnavailable( f"no live TAMQ endpoint is attached to {target_repo}" ) if len(matching) > 1: raise TamqTargetUnavailable( f"multiple TAMQ endpoints are attached to {target_repo}; endpoint_id is required" ) return str(matching[0]["endpoint_id"]) async def wake(self, request: WakeRequest) -> WakeReceipt: await self.client.negotiate() endpoint_id = await self._resolve_endpoint( request.target_repo, request.endpoint_id ) response = await self.client.send( sender_repo=request.source_repo, target_repo=request.target_repo, body=request.prompt, idempotency_key=request.lease_id, endpoint_id=endpoint_id, correlation_id=request.trigger_id, metadata={ "lease_id": request.lease_id, "trigger_id": request.trigger_id, **( {"target_agent": request.target_agent} if request.target_agent is not None else {} ), }, ) return WakeReceipt( lease_id=request.lease_id, trigger_id=request.trigger_id, endpoint_id=endpoint_id, message_id=response["message_id"], state=response["state"], deduplicated=bool(response.get("deduplicated")), ) async def receipt(self, message_id: str) -> dict: return await self.client.message(message_id) async def acknowledge(self, message_id: str) -> dict: return await self.client.acknowledge(message_id) async def retry_failed(self, message_id: str) -> dict: return await self.client.retry(message_id)