myagent-ai 1.0.0

package/communication/manager.py

@@ -0,0 +1,397 @@
+"""
+communication/manager.py - Communication Manager
+================================================
+Central manager for agent-to-agent encrypted communication.
+Orchestrates crypto keys, peer store, local/remote channels, and offline queue.
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from pathlib import Path
+from typing import Optional, List, Callable, Awaitable
+
+from .crypto import (
+    generate_ed25519_keypair,
+    generate_agent_id,
+    serialize_public_key,
+    serialize_private_key,
+    deserialize_public_key,
+    deserialize_private_key,
+    decrypt_message,
+    verify,
+    generate_x25519_keypair,
+    x25519_key_exchange,
+    derive_shared_key,
+)
+from .peer import Peer, PeerStore
+from .channel import (
+    Message,
+    MessageCallback,
+    LocalChannel,
+    RemoteChannel,
+)
+
+logger = logging.getLogger("myagent.communication")
+
+
+class CommunicationManager:
+    """
+    Manages all aspects of agent-to-agent encrypted communication.
+
+    Responsibilities:
+      - Initialize/load Ed25519 identity keys
+      - Manage the peer store (add/remove/list)
+      - Run local channel (same-machine messaging)
+      - Run remote channel (cross-machine via AICQ WebSocket)
+      - Route incoming messages to registered callbacks
+      - Offline message queue with automatic retry
+    """
+
+    def __init__(
+        self,
+        config,
+        data_dir: str | Path = "",
+    ):
+        """
+        Args:
+            config: A ``CommunicationConfig`` dataclass instance.
+            data_dir: Base data directory for persistent storage.
+        """
+        self.config = config
+        self._data_dir = Path(data_dir) if data_dir else Path.home() / ".myagent" / "data"
+
+        # Identity keys (Ed25519)
+        self._ed_priv: bytes = b""
+        self._ed_pub: bytes = b""
+        self.agent_id: str = ""
+
+        # DH keys (X25519, standalone from Ed25519)
+        self._x_priv: bytes = b""
+        self._x_pub: bytes = b""
+
+        # Components
+        self._peer_store: Optional[PeerStore] = None
+        self._local_channel: Optional[LocalChannel] = None
+        self._remote_channel: Optional[RemoteChannel] = None
+
+        # State
+        self._running = False
+        self._message_callbacks: List[MessageCallback] = []
+        self._received_messages: List[Message] = []  # In-memory recent messages
+        self._max_stored_messages = 500
+
+        # Paths
+        self._peers_path = self._data_dir / "communication" / "peers.json"
+        self._queue_db_path = self._data_dir / "communication" / "message_queue.db"
+
+    # ==================================================================
+    # Initialization
+    # ==================================================================
+
+    def initialize(self):
+        """Load or generate identity keys, initialize peer store."""
+        # Ensure directories
+        self._peers_path.parent.mkdir(parents=True, exist_ok=True)
+        self._queue_db_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Load or generate Ed25519 identity keys
+        if self.config.agent_id and self.config.private_key:
+            self._ed_priv = deserialize_private_key(self.config.private_key)
+            self._ed_pub = deserialize_public_key(self.config.agent_id)
+            self.agent_id = generate_agent_id(self._ed_pub)
+        else:
+            self._ed_priv, self._ed_pub = generate_ed25519_keypair()
+            self.agent_id = generate_agent_id(self._ed_pub)
+            self.config.agent_id = serialize_public_key(self._ed_pub)
+            self.config.private_key = serialize_private_key(self._ed_priv)
+            logger.info(f"CommunicationManager: generated new identity, agent_id={self.agent_id}")
+
+        # Generate standalone X25519 keypair for DH
+        self._x_priv, self._x_pub = generate_x25519_keypair()
+
+        # Initialize peer store
+        self._peer_store = PeerStore(self._peers_path)
+
+        logger.info(
+            f"CommunicationManager initialized: agent_id={self.agent_id}, "
+            f"peers={self._peer_store.peer_count()}, "
+            f"has_crypto={True}"
+        )
+
+    async def start(self):
+        """Start communication channels (must be called from async context)."""
+        if not self.config.enabled:
+            logger.info("CommunicationManager: disabled by config")
+            return
+
+        if self._running:
+            return
+
+        self._running = True
+
+        # Initialize local channel
+        LocalChannel.register_agent(self.agent_id)
+        self._local_channel = LocalChannel(
+            agent_id=self.agent_id,
+            ed25519_private_key=self._ed_priv,
+            ed25519_public_key=self._ed_pub,
+        )
+        self._local_channel.on_message(self._on_local_message)
+        self._local_channel.start()
+
+        # Initialize remote channel if server URL configured
+        if self.config.server_url:
+            self._remote_channel = RemoteChannel(
+                agent_id=self.agent_id,
+                ed25519_private_key=self._ed_priv,
+                ed25519_public_key=self._ed_pub,
+                server_url=self.config.server_url,
+                queue_db_path=str(self._queue_db_path),
+            )
+            self._remote_channel.on_message(self._on_remote_message)
+            self._remote_channel.start()
+
+        logger.info("CommunicationManager started")
+
+    async def stop(self):
+        """Stop all communication channels."""
+        self._running = False
+        if self._local_channel:
+            await self._local_channel.stop()
+        if self._remote_channel:
+            await self._remote_channel.stop()
+        logger.info("CommunicationManager stopped")
+
+    # ==================================================================
+    # Message Handlers
+    # ==================================================================
+
+    async def _on_local_message(self, msg: Message):
+        """Handle a message received via LocalChannel."""
+        await self._dispatch_message(msg)
+
+    async def _on_remote_message(self, msg: Message):
+        """Handle a message received via RemoteChannel."""
+        # If encrypted, attempt to decrypt
+        if msg.encrypted and msg.from_agent:
+            peer = self._peer_store.get_peer(msg.from_agent) if self._peer_store else None
+            if peer and peer.public_key:
+                try:
+                    # Use the peer's stored X25519 public key for decryption
+                    their_x_pub_hex = peer.permissions.get("x25519_public_key", "")
+                    if their_x_pub_hex:
+                        their_x_pub = bytes.fromhex(their_x_pub_hex)
+                        shared = x25519_key_exchange(self._x_priv, their_x_pub)
+                        shared_key = derive_shared_key(shared)
+                        payload = decrypt_message(shared_key, msg.content)
+                        their_ed_pub = deserialize_public_key(peer.public_key)
+                        if payload.get("signature"):
+                            sig = bytes.fromhex(payload["signature"])
+                            if verify(their_ed_pub, msg.content.encode("utf-8"), sig):
+                                msg.content = payload.get("content", msg.content)
+                                msg.encrypted = False
+                            else:
+                                logger.warning(f"Remote message from {msg.from_agent}: signature verification failed")
+                        else:
+                            msg.content = payload.get("content", msg.content)
+                            msg.encrypted = False
+                    else:
+                        logger.debug(f"No X25519 key for peer {msg.from_agent}, cannot decrypt")
+                except Exception as e:
+                    logger.warning(f"Remote message decrypt failed: {e}")
+
+        await self._dispatch_message(msg)
+
+    async def _dispatch_message(self, msg: Message):
+        """Store message and dispatch to callbacks."""
+        # Store in memory
+        self._received_messages.append(msg)
+        if len(self._received_messages) > self._max_stored_messages:
+            self._received_messages = self._received_messages[-self._max_stored_messages:]
+
+        # Notify callbacks
+        for cb in self._message_callbacks:
+            try:
+                await cb(msg)
+            except Exception as e:
+                logger.error(f"Message callback error: {e}", exc_info=True)
+
+    # ==================================================================
+    # Send Message
+    # ==================================================================
+
+    async def send_message(
+        self,
+        to_agent: str,
+        content: str,
+        msg_type: str = "text",
+    ) -> Message:
+        """
+        Send an encrypted message to another agent.
+
+        Args:
+            to_agent: Recipient agent ID.
+            content: Plaintext message content.
+            msg_type: Message type (text, heartbeat, ack, etc.)
+
+        Returns:
+            The sent Message object.
+
+        Raises:
+            RuntimeError: If the peer is not found.
+        """
+        if not self._running or not self.config.enabled:
+            raise RuntimeError("Communication is not enabled")
+
+        peer = self._peer_store.get_peer(to_agent) if self._peer_store else None
+        if not peer:
+            raise RuntimeError(f"Peer '{to_agent}' not found")
+
+        their_pub = deserialize_public_key(peer.public_key)
+
+        # Try local first, then remote
+        msg = None
+        # Check if the peer is on this machine (has a local queue)
+        if to_agent in LocalChannel._queues:
+            msg = await self._local_channel.send(
+                to_agent=to_agent,
+                content=content,
+                msg_type=msg_type,
+                their_public_key=their_pub,
+            )
+        elif self._remote_channel:
+            their_x_pub_hex = peer.permissions.get("x25519_public_key", "")
+            msg = await self._remote_channel.send(
+                to_agent=to_agent,
+                content=content,
+                their_ed_pub_hex=peer.public_key,
+                their_x25519_pub_hex=their_x_pub_hex,
+                msg_type=msg_type,
+            )
+        else:
+            # No channel available — queue locally
+            msg = Message(
+                from_agent=self.agent_id,
+                to_agent=to_agent,
+                content=content,
+                msg_type=msg_type,
+            )
+            # Store in offline queue manually
+            import json
+            wire = json.dumps({
+                "type": msg_type,
+                "from": msg.from_agent,
+                "to": msg.to_agent,
+                "content": content,
+                "encrypted": False,
+                "timestamp": msg.timestamp,
+                "msg_id": msg.id,
+            })
+            if self._remote_channel:
+                self._remote_channel._queue_message(msg.from_agent, msg.to_agent, wire)
+
+        return msg
+
+    # ==================================================================
+    # Peer Management
+    # ==================================================================
+
+    def add_peer(
+        self,
+        agent_id: str,
+        public_key: str,
+        display_name: str = "",
+    ) -> bool:
+        """
+        Add a known peer.
+
+        Args:
+            agent_id: The peer's short agent ID (or full hex public key).
+            public_key: The peer's Ed25519 public key in hex.
+            display_name: Optional human-readable name.
+
+        Returns:
+            True if added, False if already exists.
+        """
+        if not self._peer_store:
+            return False
+
+        peer = Peer(
+            agent_id=agent_id,
+            public_key=public_key,
+            display_name=display_name or agent_id,
+            added_at=time.time(),
+        )
+        return self._peer_store.add_peer(peer)
+
+    def remove_peer(self, agent_id: str) -> bool:
+        """Remove a peer by agent_id."""
+        if not self._peer_store:
+            return False
+        return self._peer_store.remove_peer(agent_id)
+
+    def get_peer(self, agent_id: str) -> Optional[Peer]:
+        """Get a peer by agent_id."""
+        if not self._peer_store:
+            return None
+        return self._peer_store.get_peer(agent_id)
+
+    def get_peer_list(self) -> List[Peer]:
+        """Return all known peers."""
+        if not self._peer_store:
+            return []
+        return self._peer_store.list_peers()
+
+    # ==================================================================
+    # Message History
+    # ==================================================================
+
+    def get_messages(
+        self,
+        limit: int = 50,
+        from_agent: str = "",
+        to_agent: str = "",
+    ) -> List[dict]:
+        """Get recent received messages with optional filtering."""
+        msgs = list(self._received_messages)
+        if from_agent:
+            msgs = [m for m in msgs if m.from_agent == from_agent]
+        if to_agent:
+            msgs = [m for m in msgs if m.to_agent == to_agent]
+        msgs = msgs[-limit:]
+        return [m.to_dict() for m in msgs]
+
+    # ==================================================================
+    # Callbacks
+    # ==================================================================
+
+    def on_message(self, callback: MessageCallback):
+        """Register a callback for incoming messages."""
+        self._message_callbacks.append(callback)
+
+    # ==================================================================
+    # Status
+    # ==================================================================
+
+    def get_status(self) -> dict:
+        """Return current communication status."""
+        remote_connected = False
+        if self._remote_channel:
+            remote_connected = self._remote_channel._connected
+
+        return {
+            "enabled": self.config.enabled,
+            "running": self._running,
+            "agent_id": self.agent_id,
+            "public_key": serialize_public_key(self._ed_pub) if self._ed_pub else "",
+            "x25519_public_key": self._x_pub.hex() if self._x_pub else "",
+            "peers_count": self._peer_store.peer_count() if self._peer_store else 0,
+            "online_peers": len(self._peer_store.get_online_peers()) if self._peer_store else 0,
+            "remote_connected": remote_connected,
+            "server_url": self.config.server_url,
+            "local_channel": self._local_channel is not None,
+            "remote_channel": self._remote_channel is not None,
+            "messages_received": len(self._received_messages),
+        }

package/communication/peer.py

@@ -0,0 +1,156 @@
+"""
+communication/peer.py - Peer Management
+========================================
+Manages known peers (other agents) with persistent JSON storage.
+"""
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import dataclass, field, asdict
+from pathlib import Path
+from typing import Optional, List, Dict
+
+
+@dataclass
+class Peer:
+    """Represents a known remote/local agent peer."""
+
+    agent_id: str                          # Short hex ID (16 chars)
+    public_key: str                        # Ed25519 public key (hex, 64 chars)
+    display_name: str = ""                 # Human-readable name
+    added_at: float = 0.0                  # Unix timestamp
+    permissions: Dict[str, bool] = field(default_factory=lambda: {
+        "send_messages": True,
+        "receive_messages": True,
+    })
+    last_seen: float = 0.0                 # Last online timestamp
+    online: bool = False                   # Current online status
+
+    def to_dict(self) -> dict:
+        return {
+            "agent_id": self.agent_id,
+            "public_key": self.public_key,
+            "display_name": self.display_name,
+            "added_at": self.added_at,
+            "permissions": self.permissions,
+            "last_seen": self.last_seen,
+            "online": self.online,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "Peer":
+        return cls(
+            agent_id=data.get("agent_id", ""),
+            public_key=data.get("public_key", ""),
+            display_name=data.get("display_name", ""),
+            added_at=data.get("added_at", 0.0),
+            permissions=data.get("permissions", {
+                "send_messages": True,
+                "receive_messages": True,
+            }),
+            last_seen=data.get("last_seen", 0.0),
+            online=data.get("online", False),
+        )
+
+
+class PeerStore:
+    """Persistent store for known peers backed by a JSON file."""
+
+    def __init__(self, store_path: str | Path):
+        self._path = Path(store_path)
+        self._peers: Dict[str, Peer] = {}
+        self._load()
+
+    # ------------------------------------------------------------------
+    # Persistence
+    # ------------------------------------------------------------------
+
+    def _load(self):
+        """Load peers from disk."""
+        if not self._path.exists():
+            return
+        try:
+            data = json.loads(self._path.read_text(encoding="utf-8"))
+            for peer_data in data.get("peers", []):
+                p = Peer.from_dict(peer_data)
+                if p.agent_id:
+                    self._peers[p.agent_id] = p
+        except (json.JSONDecodeError, IOError) as e:
+            # Corrupt file — start fresh
+            self._peers = {}
+
+    def _save(self):
+        """Save peers to disk."""
+        if self._path.suffix:
+            # Path is a file (e.g., /tmp/peers.json)
+            self._path.parent.mkdir(parents=True, exist_ok=True)
+        else:
+            # Path is a directory, use peers.json inside
+            self._path.mkdir(parents=True, exist_ok=True)
+            self._path = self._path / "peers.json"
+        data = {
+            "peers": [p.to_dict() for p in self._peers.values()],
+            "updated_at": time.time(),
+        }
+        self._path.write_text(
+            json.dumps(data, indent=2, ensure_ascii=False),
+            encoding="utf-8",
+        )
+
+    # ------------------------------------------------------------------
+    # CRUD
+    # ------------------------------------------------------------------
+
+    def add_peer(self, peer: Peer) -> bool:
+        """Add a peer. Returns False if already exists."""
+        if peer.agent_id in self._peers:
+            return False
+        if peer.added_at == 0.0:
+            peer.added_at = time.time()
+        self._peers[peer.agent_id] = peer
+        self._save()
+        return True
+
+    def remove_peer(self, agent_id: str) -> bool:
+        """Remove a peer by agent_id. Returns False if not found."""
+        if agent_id not in self._peers:
+            return False
+        del self._peers[agent_id]
+        self._save()
+        return True
+
+    def get_peer(self, agent_id: str) -> Optional[Peer]:
+        """Get a peer by agent_id."""
+        return self._peers.get(agent_id)
+
+    def update_peer(self, agent_id: str, **fields) -> bool:
+        """Update fields on an existing peer. Returns False if not found."""
+        peer = self._peers.get(agent_id)
+        if not peer:
+            return False
+        for k, v in fields.items():
+            if hasattr(peer, k):
+                setattr(peer, k, v)
+        self._save()
+        return True
+
+    def list_peers(self) -> List[Peer]:
+        """Return all peers."""
+        return list(self._peers.values())
+
+    def peer_count(self) -> int:
+        """Return the number of stored peers."""
+        return len(self._peers)
+
+    def set_online(self, agent_id: str, online: bool):
+        """Mark a peer as online/offline."""
+        if agent_id in self._peers:
+            self._peers[agent_id].online = online
+            if online:
+                self._peers[agent_id].last_seen = time.time()
+            self._save()
+
+    def get_online_peers(self) -> List[Peer]:
+        """Return peers currently marked as online."""
+        return [p for p in self._peers.values() if p.online]