tmux-amq/src/tamq/client.py

"""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)