myagent-ai 1.0.0

package/communication/channel.py

@@ -0,0 +1,576 @@
+"""
+communication/channel.py - Message Channel
+===========================================
+LocalChannel: same-machine agent-to-agent via asyncio queues.
+RemoteChannel: cross-machine communication via WebSocket (AICQ relay).
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import time
+import uuid
+import logging
+from dataclasses import dataclass, field, asdict
+from typing import Optional, List, Callable, Awaitable, Any
+
+from .crypto import (
+    derive_shared_key,
+    generate_x25519_keypair,
+    x25519_key_exchange,
+    encrypt_message,
+    decrypt_message,
+    sign,
+    verify,
+    serialize_public_key,
+    deserialize_public_key,
+)
+
+logger = logging.getLogger("myagent.communication")
+
+
+# =========================================================================
+# Message Data Model
+# =========================================================================
+
+@dataclass
+class Message:
+    """A communication message between agents."""
+
+    id: str = ""
+    from_agent: str = ""          # Sender agent_id
+    to_agent: str = ""            # Recipient agent_id
+    content: str = ""             # Plaintext content
+    timestamp: float = 0.0
+    encrypted: bool = False
+    signature: str = ""           # Hex signature (from sender)
+    msg_type: str = "text"        # text | handshake | heartbeat | ack
+
+    def __post_init__(self):
+        if not self.id:
+            self.id = uuid.uuid4().hex[:16]
+        if self.timestamp == 0.0:
+            self.timestamp = time.time()
+
+    def to_dict(self) -> dict:
+        d = asdict(self)
+        return d
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "Message":
+        return cls(**{k: v for k, v in data.items() if k in cls.__dataclass_fields__})
+
+    def to_json(self) -> str:
+        return json.dumps(self.to_dict(), ensure_ascii=False)
+
+
+# =========================================================================
+# Message Callback
+# =========================================================================
+
+MessageCallback = Callable[[Message], Awaitable[None]]
+
+
+# =========================================================================
+# Local Channel — Same-Machine Agent Communication
+# =========================================================================
+
+class LocalChannel:
+    """
+    LocalChannel enables encrypted message passing between agents on the
+    same machine using asyncio queues.
+
+    Each agent instance has its own ``LocalChannel``.  Messages are routed
+    through a shared broker (dict of queues) keyed by agent_id.
+    """
+
+    # Shared broker across all LocalChannel instances on this process
+    _queues: dict[str, asyncio.Queue] = {}
+    _callbacks: dict[str, List[MessageCallback]] = {}
+
+    def __init__(
+        self,
+        agent_id: str,
+        ed25519_private_key: bytes,
+        ed25519_public_key: bytes,
+    ):
+        self.agent_id = agent_id
+        self._ed_priv = ed25519_private_key
+        self._ed_pub = ed25519_public_key
+
+        # Generate standalone X25519 keypair for DH (separate from Ed25519)
+        self._x_priv, self._x_pub = generate_x25519_keypair()
+
+        # Ensure our queue exists
+        if agent_id not in LocalChannel._queues:
+            LocalChannel._queues[agent_id] = asyncio.Queue()
+        if agent_id not in LocalChannel._callbacks:
+            LocalChannel._callbacks[agent_id] = []
+
+        self._running = False
+        self._receiver_task: Optional[asyncio.Task] = None
+
+    # ------------------------------------------------------------------
+    # Derive a shared AES key with another agent
+    # ------------------------------------------------------------------
+
+    def _derive_key(self, their_ed_pub: bytes) -> bytes:
+        their_x_pub = ed25519_to_x25519_public(their_ed_pub)
+        shared_secret = None
+        # Simple DH: we need the actual x25519 exchange function
+        from .crypto import x25519_key_exchange
+        shared_secret = x25519_key_exchange(self._x_priv, their_x_pub)
+        return derive_shared_key(shared_secret)
+
+    # ------------------------------------------------------------------
+    # Send
+    # ------------------------------------------------------------------
+
+    async def send(
+        self,
+        to_agent: str,
+        content: str,
+        msg_type: str = "text",
+        their_public_key: Optional[bytes] = None,
+    ) -> Message:
+        """Send an encrypted message to a local peer."""
+        msg = Message(
+            from_agent=self.agent_id,
+            to_agent=to_agent,
+            content=content,
+            msg_type=msg_type,
+        )
+
+        # Sign the message
+        msg_data = json.dumps({
+            "id": msg.id,
+            "from": msg.from_agent,
+            "to": msg.to_agent,
+            "content": msg.content,
+            "timestamp": msg.timestamp,
+            "type": msg.msg_type,
+        }, sort_keys=True).encode("utf-8")
+        sig = sign(self._ed_priv, msg_data)
+        msg.signature = sig.hex()
+
+        # Note: encryption requires a pre-established shared key.
+        # In local mode, messages are already in-process and signed.
+        # For real E2EE, the shared key must be established via DH with
+        # the peer's X25519 public key (separate from Ed25519 identity).
+        msg.encrypted = False
+
+        # Put into recipient's queue
+        queue = LocalChannel._queues.get(to_agent)
+        if queue is None:
+            raise RuntimeError(f"LocalChannel: agent '{to_agent}' has no queue")
+
+        await queue.put(msg)
+        return msg
+
+    # ------------------------------------------------------------------
+    # Receive
+    # ------------------------------------------------------------------
+
+    def on_message(self, callback: MessageCallback):
+        """Register a callback for incoming messages."""
+        LocalChannel._callbacks.setdefault(self.agent_id, []).append(callback)
+
+    async def _dispatch(self, msg: Message):
+        """Dispatch a received message to registered callbacks."""
+        for cb in LocalChannel._callbacks.get(self.agent_id, []):
+            try:
+                await cb(msg)
+            except Exception as e:
+                logger.error(f"LocalChannel callback error: {e}", exc_info=True)
+
+    # ------------------------------------------------------------------
+    # Receiver loop
+    # ------------------------------------------------------------------
+
+    async def _receiver_loop(self):
+        """Background task: read from queue and dispatch."""
+        queue = LocalChannel._queues.get(self.agent_id)
+        if not queue:
+            return
+        while self._running:
+            try:
+                msg = await asyncio.wait_for(queue.get(), timeout=1.0)
+                await self._dispatch(msg)
+            except asyncio.TimeoutError:
+                continue
+            except Exception as e:
+                logger.error(f"LocalChannel receiver error: {e}")
+
+    def start(self):
+        """Start the receiver loop."""
+        if self._running:
+            return
+        self._running = True
+        self._receiver_task = asyncio.create_task(self._receiver_loop())
+
+    async def stop(self):
+        """Stop the receiver loop."""
+        self._running = False
+        if self._receiver_task:
+            self._receiver_task.cancel()
+            try:
+                await self._receiver_task
+            except asyncio.CancelledError:
+                pass
+            self._receiver_task = None
+
+    # ------------------------------------------------------------------
+    # Static: register a new agent queue
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def register_agent(agent_id: str):
+        """Register a new agent in the shared broker."""
+        if agent_id not in LocalChannel._queues:
+            LocalChannel._queues[agent_id] = asyncio.Queue()
+        if agent_id not in LocalChannel._callbacks:
+            LocalChannel._callbacks[agent_id] = []
+
+
+# =========================================================================
+# Remote Channel — Cross-Machine via WebSocket
+# =========================================================================
+
+class RemoteChannel:
+    """
+    RemoteChannel connects to the AICQ relay server via WebSocket for
+    cross-machine agent communication.
+
+    Features:
+    - Encrypted messaging via AES-256-GCM
+    - Offline message queue (SQLite)
+    - Automatic reconnection
+    - Heartbeat for online status
+    """
+
+    def __init__(
+        self,
+        agent_id: str,
+        ed25519_private_key: bytes,
+        ed25519_public_key: bytes,
+        server_url: str = "wss://aicq.online/ws",
+        queue_db_path: str = "",
+    ):
+        self.agent_id = agent_id
+        self._ed_priv = ed25519_private_key
+        self._ed_pub = ed25519_public_key
+        self.server_url = server_url
+        self._queue_db_path = queue_db_path
+
+        # Generate standalone X25519 keypair for DH
+        self._x_priv, self._x_pub = generate_x25519_keypair()
+
+        self._ws = None
+        self._running = False
+        self._connected = False
+        self._recv_task: Optional[asyncio.Task] = None
+        self._heartbeat_task: Optional[asyncio.Task] = None
+        self._callbacks: List[MessageCallback] = []
+        self._shared_keys: dict[str, bytes] = {}  # agent_id -> AES key
+        self._pending_handshakes: dict[str, asyncio.Future] = {}
+
+        # Offline queue
+        self._queue_conn = None
+
+    # ------------------------------------------------------------------
+    # Offline Queue (SQLite)
+    # ------------------------------------------------------------------
+
+    def _init_queue_db(self):
+        """Initialize the offline message queue SQLite database."""
+        if not self._queue_db_path:
+            return
+        import sqlite3
+        self._queue_conn = sqlite3.connect(self._queue_db_path)
+        self._queue_conn.execute("""
+            CREATE TABLE IF NOT EXISTS queued_messages (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                from_agent TEXT NOT NULL,
+                to_agent TEXT NOT NULL,
+                encrypted_content TEXT NOT NULL,
+                created_at REAL NOT NULL,
+                delivered INTEGER DEFAULT 0
+            )
+        """)
+        self._queue_conn.commit()
+
+    def _queue_message(self, from_agent: str, to_agent: str, encrypted_content: str):
+        """Queue a message for later delivery."""
+        if not self._queue_conn:
+            return
+        self._queue_conn.execute(
+            "INSERT INTO queued_messages (from_agent, to_agent, encrypted_content, created_at, delivered) VALUES (?, ?, ?, ?, 0)",
+            (from_agent, to_agent, encrypted_content, time.time()),
+        )
+        self._queue_conn.commit()
+
+    def _get_undelivered(self) -> list:
+        """Get undelivered messages."""
+        if not self._queue_conn:
+            return []
+        rows = self._queue_conn.execute(
+            "SELECT id, from_agent, to_agent, encrypted_content, created_at FROM queued_messages WHERE delivered = 0 ORDER BY created_at ASC"
+        ).fetchall()
+        return [{"id": r[0], "from": r[1], "to": r[2], "content": r[3], "created": r[4]} for r in rows]
+
+    def _mark_delivered(self, msg_id: int):
+        """Mark a message as delivered."""
+        if not self._queue_conn:
+            return
+        self._queue_conn.execute(
+            "UPDATE queued_messages SET delivered = 1 WHERE id = ?", (msg_id,)
+        )
+        self._queue_conn.commit()
+
+    # ------------------------------------------------------------------
+    # Key Management
+    # ------------------------------------------------------------------
+
+    def _get_or_derive_key(self, their_agent_id: str, their_x25519_pub: bytes) -> bytes:
+        """Get or derive a shared key with a peer using their X25519 public key."""
+        if their_agent_id in self._shared_keys:
+            return self._shared_keys[their_agent_id]
+        shared = x25519_key_exchange(self._x_priv, their_x25519_pub)
+        key = derive_shared_key(shared)
+        self._shared_keys[their_agent_id] = key
+        return key
+
+    # ------------------------------------------------------------------
+    # WebSocket Connection
+    # ------------------------------------------------------------------
+
+    async def connect(self):
+        """Connect to the AICQ relay server."""
+        try:
+            import websockets
+            self._ws = await websockets.connect(
+                self.server_url,
+                ping_interval=30,
+                ping_timeout=10,
+                close_timeout=5,
+            )
+            self._connected = True
+            logger.info(f"RemoteChannel: connected to {self.server_url}")
+
+            # Send auth/register message (includes both identity and DH key)
+            auth_msg = json.dumps({
+                "type": "register",
+                "agent_id": self.agent_id,
+                "public_key": serialize_public_key(self._ed_pub),
+                "dh_public_key": self._x_pub.hex(),
+            })
+            await self._ws.send(auth_msg)
+
+        except Exception as e:
+            logger.warning(f"RemoteChannel: connection failed: {e}")
+            self._connected = False
+
+    async def disconnect(self):
+        """Disconnect from the relay server."""
+        self._connected = False
+        if self._ws:
+            try:
+                await self._ws.close()
+            except Exception:
+                pass
+            self._ws = None
+
+    # ------------------------------------------------------------------
+    # Send
+    # ------------------------------------------------------------------
+
+    async def send(
+        self,
+        to_agent: str,
+        content: str,
+        their_ed_pub_hex: str = "",
+        their_x25519_pub_hex: str = "",
+        msg_type: str = "text",
+    ) -> Message:
+        """Send an encrypted message to a remote peer."""
+        msg = Message(
+            from_agent=self.agent_id,
+            to_agent=to_agent,
+            content=content,
+            msg_type=msg_type,
+        )
+
+        # Sign the message
+        msg_data = json.dumps({
+            "id": msg.id,
+            "from": msg.from_agent,
+            "to": msg.to_agent,
+            "content": msg.content,
+            "timestamp": msg.timestamp,
+            "type": msg.msg_type,
+        }, sort_keys=True).encode("utf-8")
+        sig = sign(self._ed_priv, msg_data)
+        msg.signature = sig.hex()
+
+        # Encrypt using X25519 DH shared key
+        if their_x25519_pub_hex:
+            try:
+                their_x_pub = bytes.fromhex(their_x25519_pub_hex)
+                shared_key = self._get_or_derive_key(to_agent, their_x_pub)
+                encrypted_payload = encrypt_message(shared_key, {
+                    "content": content,
+                    "timestamp": msg.timestamp,
+                    "msg_type": msg.msg_type,
+                    "msg_id": msg.id,
+                    "signature": msg.signature,
+                })
+                msg.content = encrypted_payload
+                msg.encrypted = True
+            except Exception as e:
+                logger.warning(f"RemoteChannel encryption failed: {e}")
+                msg.encrypted = False
+        else:
+            msg.encrypted = False
+
+        # Send via WebSocket or queue
+        wire_msg = json.dumps({
+            "type": msg_type,
+            "from": msg.from_agent,
+            "to": msg.to_agent,
+            "content": msg.content,
+            "encrypted": msg.encrypted,
+            "timestamp": msg.timestamp,
+            "msg_id": msg.id,
+            "signature": msg.signature,
+        })
+
+        if self._connected and self._ws:
+            try:
+                await self._ws.send(wire_msg)
+            except Exception as e:
+                logger.warning(f"RemoteChannel: send failed, queuing: {e}")
+                self._queue_message(msg.from_agent, msg.to_agent, wire_msg)
+                self._connected = False
+        else:
+            self._queue_message(msg.from_agent, msg.to_agent, wire_msg)
+
+        return msg
+
+    # ------------------------------------------------------------------
+    # Receive Loop
+    # ------------------------------------------------------------------
+
+    async def _receiver_loop(self):
+        """Background task: receive messages from WebSocket."""
+        while self._running:
+            if not self._connected:
+                try:
+                    await self.connect()
+                except Exception:
+                    await asyncio.sleep(5)
+                    continue
+
+            try:
+                raw = await self._ws.recv()
+                data = json.loads(raw)
+                await self._handle_incoming(data)
+            except Exception as e:
+                if self._running:
+                    logger.warning(f"RemoteChannel recv error: {e}")
+                    self._connected = False
+                await asyncio.sleep(1)
+
+    async def _handle_incoming(self, data: dict):
+        """Handle an incoming message from the relay."""
+        msg_type = data.get("type", "text")
+
+        if msg_type == "heartbeat_ack":
+            return
+
+        msg = Message(
+            id=data.get("msg_id", ""),
+            from_agent=data.get("from", ""),
+            to_agent=data.get("to", ""),
+            content=data.get("content", ""),
+            timestamp=data.get("timestamp", 0.0),
+            encrypted=data.get("encrypted", False),
+            signature=data.get("signature", ""),
+            msg_type=msg_type,
+        )
+
+        # Dispatch to callbacks
+        for cb in self._callbacks:
+            try:
+                await cb(msg)
+            except Exception as e:
+                logger.error(f"RemoteChannel callback error: {e}")
+
+    # ------------------------------------------------------------------
+    # Heartbeat
+    # ------------------------------------------------------------------
+
+    async def _heartbeat_loop(self):
+        """Periodically send heartbeat messages."""
+        while self._running:
+            await asyncio.sleep(30)
+            if self._connected and self._ws:
+                try:
+                    await self._ws.send(json.dumps({"type": "heartbeat"}))
+                except Exception:
+                    self._connected = False
+
+    # ------------------------------------------------------------------
+    # Retry Undelivered Messages
+    # ------------------------------------------------------------------
+
+    async def _retry_loop(self):
+        """Periodically retry sending queued messages."""
+        while self._running:
+            await asyncio.sleep(10)
+            if not self._connected or not self._ws:
+                continue
+            undelivered = self._get_undelivered()
+            for item in undelivered:
+                try:
+                    await self._ws.send(item["content"])
+                    self._mark_delivered(item["id"])
+                except Exception:
+                    break  # Stop retrying if send fails
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def on_message(self, callback: MessageCallback):
+        """Register a callback for incoming messages."""
+        self._callbacks.append(callback)
+
+    def start(self):
+        """Start the remote channel (receiver, heartbeat, retry loops)."""
+        if self._running:
+            return
+        self._running = True
+        self._init_queue_db()
+        self._recv_task = asyncio.create_task(self._receiver_loop())
+        self._heartbeat_task = asyncio.create_task(self._heartbeat_loop())
+        asyncio.create_task(self._retry_loop())
+
+    async def stop(self):
+        """Stop the remote channel."""
+        self._running = False
+        await self.disconnect()
+        if self._recv_task:
+            self._recv_task.cancel()
+            try:
+                await self._recv_task
+            except asyncio.CancelledError:
+                pass
+        if self._heartbeat_task:
+            self._heartbeat_task.cancel()
+            try:
+                await self._heartbeat_task
+            except asyncio.CancelledError:
+                pass
+        if self._queue_conn:
+            self._queue_conn.close()
+            self._queue_conn = None