enki_core/

agent.rs

use crate::error::Result;
use crate::message::Message;
use crate::request_queue::RequestQueue;
use async_trait::async_trait;
use std::sync::Arc;

/// The core trait that defines an autonomous agent.
///
/// Agents are the primary unit of computation in Enki. Each agent:
/// - Has a unique name for identification and message routing
/// - Receives lifecycle callbacks (`on_start`, `on_stop`)
/// - Processes messages via `on_message` or `on_generic_message`
/// - Optionally exposes tools via `tool_invoker`
///
/// # Implementing an Agent
///
/// ```rust,no_run
/// use enki_core::{Agent, AgentContext, Message};
/// use enki_core::error::Result;
/// use async_trait::async_trait;
///
/// struct CounterAgent {
///     count: u32,
/// }
///
/// #[async_trait]
/// impl Agent for CounterAgent {
///     fn name(&self) -> String {
///         "counter".to_string()
///     }
///
///     async fn on_start(&mut self, _ctx: &mut AgentContext) -> Result<()> {
///         println!("Counter starting with count: {}", self.count);
///         Ok(())
///     }
///
///     async fn on_message(&mut self, msg: Message, ctx: &mut AgentContext) -> Result<()> {
///         self.count += 1;
///         if let Some(id) = msg.correlation_id() {
///             ctx.report_result(&id, format!("Count: {}", self.count));
///         }
///         Ok(())
///     }
/// }
/// ```
#[async_trait]
pub trait Agent: Send + Sync {
    /// Returns the unique name of this agent.
    ///
    /// This name is used for message routing and identification within a mesh.
    fn name(&self) -> String;

    /// Called when the agent starts.
    ///
    /// Use this for initialization logic. The mesh is running but other agents
    /// may not have started yet.
    async fn on_start(&mut self, _ctx: &mut AgentContext) -> Result<()> {
        Ok(())
    }

    /// Called when a binary message is received.
    ///
    /// # Arguments
    /// * `msg` - The received message with binary payload
    /// * `ctx` - The agent context for mesh interaction
    async fn on_message(&mut self, _msg: Message, _ctx: &mut AgentContext) -> Result<()> {
        Ok(())
    }

    /// Handle a generic string message and return a response.
    ///
    /// Default implementation echoes the content back.
    async fn on_generic_message(
        &mut self,
        msg: crate::message::GenericMessage,
        _ctx: &mut AgentContext,
    ) -> Result<crate::message::GenericResponse> {
        Ok(crate::message::GenericResponse::new(msg.content))
    }

    /// Called when the agent is stopping.
    ///
    /// Use this for cleanup logic.
    async fn on_stop(&mut self, _ctx: &mut AgentContext) -> Result<()> {
        Ok(())
    }

    /// Get the tool invoker for this agent (if it has actions).
    ///
    /// Return `Some` if this agent exposes tools that can be called by LLMs.
    fn tool_invoker(&self) -> Option<&crate::action::ToolInvoker> {
        None
    }

    /// Get mutable tool invoker for dynamic tool registration.
    fn tool_invoker_mut(&mut self) -> Option<&mut crate::action::ToolInvoker> {
        None
    }
}

/// Runtime context provided to agents during lifecycle and message handling.
///
/// The context gives agents access to mesh information and utilities
/// for interacting with the request/response system.
///
/// # Example
///
/// ```rust,no_run
/// # use enki_core::{Agent, AgentContext, Message};
/// # use enki_core::error::Result;
/// # use async_trait::async_trait;
/// # struct MyAgent;
/// # #[async_trait] impl Agent for MyAgent {
/// #     fn name(&self) -> String { "my".into() }
/// async fn on_message(&mut self, msg: Message, ctx: &mut AgentContext) -> Result<()> {
///     println!("Running in mesh: {}", ctx.mesh_name);
///     if let Some(id) = msg.correlation_id() {
///         ctx.report_result(&id, "Done!".to_string());
///     }
///     Ok(())
/// }
/// # }
/// ```
use serde::{Deserialize, Serialize};

/// Events specific to agent execution flow
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum AgentEvent {
    MessageSent {
        from: String,
        to: String,
        content: String,
    },
    ToolCalled {
        agent: String,
        tool: String,
        input: String,
    },
    ToolResult {
        agent: String,
        tool: String,
        result: String,
    },
    // Generic event for extensibility
    Custom {
        kind: String,
        data: String,
    },
}

pub struct AgentContext {
    /// The name of the mesh this agent is running in.
    pub mesh_name: String,
    request_queue: Option<Arc<RequestQueue>>,
    listeners: Vec<Arc<dyn Fn(&AgentEvent) + Send + Sync>>,
}

impl AgentContext {
    /// Create a new agent context.
    ///
    /// # Arguments
    /// * `mesh_name` - Name of the mesh
    /// * `request_queue` - Optional request queue for request/response pattern
    pub fn new(mesh_name: String, request_queue: Option<Arc<RequestQueue>>) -> Self {
        Self {
            mesh_name,
            request_queue,
            listeners: Vec::new(),
        }
    }

    /// Report a result for a pending request.
    ///
    /// Call this to send a response back to the requester when handling
    /// a message with a correlation ID.
    pub fn report_result(&self, request_id: &str, response: String) {
        if let Some(queue) = &self.request_queue {
            queue.complete(request_id, response);
        }
    }

    /// Check if this context has a request queue available.
    pub fn has_request_queue(&self) -> bool {
        self.request_queue.is_some()
    }

    /// Register an event listener.
    pub fn listen(&mut self, callback: impl Fn(&AgentEvent) + Send + Sync + 'static) {
        self.listeners.push(Arc::new(callback));
    }

    /// Emit an event to all listeners.
    pub fn emit(&self, event: AgentEvent) {
        for listener in &self.listeners {
            listener(&event);
        }
    }
}