enki_local/

mesh.rs

use async_trait::async_trait;
use dashmap::DashMap;
use enki_core::agent::{Agent, AgentContext};
use enki_core::error::{Error, Result};
use enki_core::mesh::Mesh;
use enki_core::message::Message;
use enki_core::request_queue::{MeshRequest, MeshResult, RequestQueue};
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::mpsc::{self, Sender};
use tokio::task::JoinHandle;
use tracing::{error, info, trace};

pub struct LocalMesh {
    name: String,
    agents: DashMap<String, Sender<Message>>,
    tasks: DashMap<String, JoinHandle<()>>,
    request_queue: Arc<RequestQueue>,
}

impl LocalMesh {
    pub fn new(name: impl Into<String>) -> Self {
        Self {
            name: name.into(),
            agents: DashMap::new(),
            tasks: DashMap::new(),
            request_queue: Arc::new(RequestQueue::new()),
        }
    }

    /// Get the request queue (for sharing with agent wrappers)
    pub fn request_queue(&self) -> Arc<RequestQueue> {
        self.request_queue.clone()
    }
}

#[async_trait]
impl Mesh for LocalMesh {
    async fn start(&self) -> Result<()> {
        // For local mesh, start might just be a signal,
        // but agents are started when added for now.
        Ok(())
    }

    async fn stop(&self) -> Result<()> {
        for entry in self.tasks.iter() {
            entry.value().abort();
        }
        self.tasks.clear();
        self.agents.clear();
        Ok(())
    }

    async fn add_agent(&self, mut agent: Box<dyn Agent + 'static>) -> Result<()> {
        let name = agent.name();
        if self.agents.contains_key(&name) {
            return Err(Error::MeshError(format!("Agent {} already exists", name)));
        }

        let (tx, mut rx) = mpsc::channel(100);
        self.agents.insert(name.clone(), tx);

        let mesh_name = self.name.clone();
        let agent_name = name.clone();
        let request_queue = self.request_queue.clone();

        // Spawn agent loop
        let handle = tokio::spawn(async move {
            let mut ctx = AgentContext::new(mesh_name, Some(request_queue));
            if let Err(e) = agent.on_start(&mut ctx).await {
                error!(agent = %agent_name, error = ?e, "Error starting agent");
                return;
            }

            while let Some(msg) = rx.recv().await {
                let now = chrono::Utc::now().timestamp_micros();
                let latency = (now - msg.created_at).max(0) as u64;
                trace!(
                    agent = %agent_name,
                    latency_us = latency,
                    msg_id = %msg.id,
                    from = %msg.from,
                    to = %msg.to,
                    content_type = %msg.content.type_name(),
                    correlation_id = ?msg.correlation_id,
                    "processing message"
                );

                let start = std::time::Instant::now();
                if let Err(e) = agent.on_message(msg, &mut ctx).await {
                    error!(agent = %agent_name, error = ?e, "Error processing message");
                }
                let duration = start.elapsed().as_micros() as u64;
                trace!(agent = %agent_name, duration_us = duration, "message processed");
            }

            if let Err(e) = agent.on_stop(&mut ctx).await {
                error!(agent = %agent_name, error = ?e, "Error stopping agent");
            }
        });

        self.tasks.insert(name, handle);
        Ok(())
    }

    async fn send(&self, message: Message) -> Result<()> {
        // Route based on message target
        // Clone target to avoid borrow checker issues when moving message
        let target = message.to.clone();

        match target {
            enki_core::message::MessageTarget::Agent(target_agent) => {
                if let Some(sender) = self.agents.get(&target_agent) {
                    info!(
                        mesh = %self.name,
                        from = %message.from,
                        to = %target_agent,
                        content_type = %message.content.type_name(),
                        msg_id = %message.id,
                        correlation_id = ?message.correlation_id,
                        "Sending message between agents"
                    );

                    sender.send(message).await.map_err(|_| {
                        Error::MeshError(format!("Failed to send to agent {}", target_agent))
                    })?;
                    Ok(())
                } else {
                    Err(Error::AgentNotFound(target_agent))
                }
            }
            enki_core::message::MessageTarget::Broadcast => {
                // Broadcast to all except sender
                for entry in self.agents.iter() {
                    if entry.key() != &message.from {
                        let _ = entry.value().send(message.clone()).await;
                    }
                }
                Ok(())
            }
            _ => Err(Error::MeshError(
                "Unsupported message target (Topic and Node routing not yet implemented)"
                    .to_string(),
            )),
        }
    }
}

impl LocalMesh {
    /// Broadcast a message to all agents in the mesh
    pub async fn broadcast(
        &self,
        message: Message,
        exclude: Option<&str>,
    ) -> Result<Vec<Result<()>>> {
        let mut results = Vec::new();
        let recipient_count = self.agents.len() - if exclude.is_some() { 1 } else { 0 };

        info!(
            mesh = %self.name,
            from = %message.from,
            recipients = recipient_count,
            excluded = ?exclude,
            "Broadcasting message to agents"
        );

        for entry in self.agents.iter() {
            let agent_name = entry.key();

            // Skip excluded agent if specified
            if let Some(excluded) = exclude {
                if agent_name == excluded {
                    continue;
                }
            }

            // Clone message for each agent
            let mut msg = message.clone();
            msg.to = enki_core::message::MessageTarget::Agent(agent_name.clone());

            trace!(
                mesh = %self.name,
                from = %message.from,
                to = %agent_name,
                msg_id = %msg.id,
                "Broadcasting to individual agent"
            );

            let result = entry.value().send(msg).await.map_err(|_| {
                Error::MeshError(format!("Failed to broadcast to agent {}", agent_name))
            });

            results.push(result);
        }

        Ok(results)
    }

    /// Submit a request (fire-and-forget). Returns request ID.
    pub async fn submit(&self, target: &str, payload: String) -> Result<String> {
        let request_id = self.request_queue.submit(target, payload.clone());

        info!(
            mesh = %self.name,
            target = %target,
            request_id = %request_id,
            payload_size = payload.len(),
            "Submitting request to agent"
        );

        let msg = Message::text("mesh", payload)
            .to(enki_core::message::MessageTarget::Agent(target.to_string()))
            .correlation_id(&request_id)
            .build();

        self.send(msg).await?;
        Ok(request_id)
    }

    /// Get pending requests
    pub fn get_pending(&self) -> Vec<MeshRequest> {
        self.request_queue.get_pending()
    }

    /// Check if there are pending requests
    pub fn has_pending(&self) -> bool {
        self.request_queue.has_pending()
    }

    /// Get available results (removes them from queue)
    pub fn get_results(&self) -> Vec<MeshResult> {
        self.request_queue.take_results()
    }

    /// Peek at results without removing
    pub fn peek_results(&self) -> Vec<MeshResult> {
        self.request_queue.peek_results()
    }

    /// Send reminders for stale requests
    pub async fn send_reminders(&self, older_than_secs: f64) -> Result<Vec<String>> {
        let stale = self
            .request_queue
            .get_stale(Duration::from_secs_f64(older_than_secs));
        let mut reminded = Vec::new();

        for req in stale {
            let reminder_msg = Message::text(
                "mesh",
                format!("REMINDER: Please complete request {}", req.id),
            )
            .to(enki_core::message::MessageTarget::Agent(req.target.clone()))
            .correlation_id(&req.id)
            .build();

            if self.send(reminder_msg).await.is_ok() {
                self.request_queue.increment_reminder(&req.id);
                reminded.push(req.id);
            }
        }

        Ok(reminded)
    }

    /// Wait for a specific result with auto-reminders
    pub async fn wait_for(
        &self,
        request_id: &str,
        timeout_secs: f64,
        reminder_interval_secs: f64,
    ) -> Result<MeshResult> {
        let timeout = Duration::from_secs_f64(timeout_secs);
        let reminder_interval = Duration::from_secs_f64(reminder_interval_secs);
        let deadline = std::time::Instant::now() + timeout;
        let mut last_reminder = std::time::Instant::now();

        loop {
            // Check if result is available
            if let Some(result) = self.request_queue.take_result(request_id) {
                return Ok(result);
            }

            // Check if request still exists
            let pending = self.request_queue.get_pending();
            let request = pending.iter().find(|r| r.id == request_id);
            if request.is_none() {
                return Err(Error::MeshError(format!(
                    "Request {} not found",
                    request_id
                )));
            }

            // Check timeout
            if std::time::Instant::now() >= deadline {
                return Err(Error::MeshError(format!(
                    "Request {} timed out",
                    request_id
                )));
            }

            // Send reminder if interval passed
            if last_reminder.elapsed() >= reminder_interval {
                let _ = self.send_reminders(reminder_interval_secs).await;
                last_reminder = std::time::Instant::now();
            }

            // Wait briefly
            tokio::time::sleep(Duration::from_millis(100)).await;
        }
    }

    /// Collect all results, blocking until all pending complete
    pub async fn collect_results(&self, reminder_interval_secs: f64) -> Vec<MeshResult> {
        let reminder_interval = Duration::from_secs_f64(reminder_interval_secs);
        let mut last_reminder = std::time::Instant::now();
        let mut all_results = Vec::new();

        while self.request_queue.has_pending() {
            // Collect any available results
            let results = self.request_queue.take_results();
            all_results.extend(results);

            // Send reminders if interval passed
            if last_reminder.elapsed() >= reminder_interval {
                let _ = self.send_reminders(reminder_interval_secs).await;
                last_reminder = std::time::Instant::now();
            }

            // Wait briefly before checking again
            tokio::time::sleep(Duration::from_millis(100)).await;
        }

        // Get any final results
        let results = self.request_queue.take_results();
        all_results.extend(results);

        all_results
    }
}