agent_kernel/

registry.rs

//! Agent registry integration for MXP Nexus mesh discovery and heartbeats.

use std::collections::HashMap;
use std::fmt;
use std::io::ErrorKind;
use std::net::{SocketAddr, ToSocketAddrs};
use std::num::NonZeroUsize;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;

use agent_primitives::AgentManifest;
use async_trait::async_trait;
use mxp::protocol::Flags;
use mxp::transport::{SocketError, Transport, TransportConfig, TransportHandle};
use mxp::{Message, MessageType};
use thiserror::Error;
use tokio::task::JoinHandle;
use tokio::time::{MissedTickBehavior, sleep};
use tracing::{debug, info, warn};

use crate::registry_wire::{
    ErrorResponse, HeartbeatRequest, HeartbeatResponse, RegisterRequest, RegisterResponse,
};
use crate::{AgentState, SchedulerError, TaskScheduler};

/// Configuration for registration and heartbeat maintenance.
#[derive(Debug, Clone, Copy)]
pub struct RegistrationConfig {
    heartbeat_interval: Duration,
    initial_retry_delay: Duration,
    max_retry_delay: Duration,
    max_consecutive_failures: NonZeroUsize,
}

impl RegistrationConfig {
    /// Creates a new configuration.
    #[must_use]
    pub fn new(
        heartbeat_interval: Duration,
        initial_retry_delay: Duration,
        max_retry_delay: Duration,
        max_consecutive_failures: NonZeroUsize,
    ) -> Self {
        Self {
            heartbeat_interval,
            initial_retry_delay,
            max_retry_delay,
            max_consecutive_failures,
        }
    }

    /// Returns the heartbeat interval.
    #[must_use]
    pub const fn heartbeat_interval(self) -> Duration {
        self.heartbeat_interval
    }

    /// Returns the initial retry delay.
    #[must_use]
    pub const fn initial_retry_delay(self) -> Duration {
        self.initial_retry_delay
    }

    /// Returns the maximum retry delay.
    #[must_use]
    pub const fn max_retry_delay(self) -> Duration {
        self.max_retry_delay
    }

    /// Returns the limit on consecutive heartbeat failures before re-registration.
    #[must_use]
    pub const fn max_consecutive_failures(self) -> NonZeroUsize {
        self.max_consecutive_failures
    }

    /// Validates the configuration.
    ///
    /// # Errors
    ///
    /// Returns [`RegistryError::InvalidConfig`] when any duration is zero or the
    /// retry delay bounds are inconsistent.
    pub fn validate(self) -> RegistryResult<()> {
        if self.heartbeat_interval.is_zero() {
            return Err(RegistryError::InvalidConfig(
                "heartbeat interval must be greater than zero",
            ));
        }
        if self.initial_retry_delay.is_zero() {
            return Err(RegistryError::InvalidConfig(
                "initial retry delay must be greater than zero",
            ));
        }
        if self.max_retry_delay.is_zero() {
            return Err(RegistryError::InvalidConfig(
                "max retry delay must be greater than zero",
            ));
        }
        if self.initial_retry_delay > self.max_retry_delay {
            return Err(RegistryError::InvalidConfig(
                "initial retry delay cannot exceed max retry delay",
            ));
        }
        Ok(())
    }
}

impl Default for RegistrationConfig {
    fn default() -> Self {
        Self {
            heartbeat_interval: Duration::from_secs(10),
            initial_retry_delay: Duration::from_secs(1),
            max_retry_delay: Duration::from_secs(30),
            max_consecutive_failures: NonZeroUsize::new(3).expect("non-zero"),
        }
    }
}

/// Result alias for registry operations.
pub type RegistryResult<T> = Result<T, RegistryError>;

/// Errors surfaced by registry integration.
#[derive(Debug, Error)]
pub enum RegistryError {
    /// Registration configuration was invalid.
    #[error("invalid registration configuration: {0}")]
    InvalidConfig(&'static str),
    /// Scheduler rejected a task submission.
    #[error(transparent)]
    Scheduler(#[from] SchedulerError),
    /// Registry backend failure.
    #[error("registry backend error: {reason}")]
    Backend {
        /// Human-readable context provided by the backend.
        reason: String,
    },
}

impl RegistryError {
    /// Convenience helper to construct backend errors.
    #[must_use]
    pub fn backend(reason: impl Into<String>) -> Self {
        Self::Backend {
            reason: reason.into(),
        }
    }
}

/// MXP-backed registry client that speaks directly to the MXP Nexus registry service.
#[derive(Debug)]
pub struct MxpRegistryClient {
    handle: TransportHandle,
    registry_addr: SocketAddr,
    agent_endpoint: SocketAddr,
}

impl MxpRegistryClient {
    /// Establishes a registry client using the provided endpoint configuration.
    ///
    /// # Errors
    ///
    /// Returns [`RegistryError::Backend`] if the transport cannot be bound.
    pub fn connect(
        registry_addr: impl ToSocketAddrs,
        agent_endpoint: SocketAddr,
        transport_config: Option<TransportConfig>,
    ) -> RegistryResult<Self> {
        let registry_addr = registry_addr
            .to_socket_addrs()
            .map_err(|err| {
                RegistryError::backend(format!("failed to resolve registry endpoint: {err:?}"))
            })?
            .next()
            .ok_or_else(|| RegistryError::backend("registry endpoint resolved to no address"))?;

        let config = transport_config.unwrap_or_else(default_transport_config);
        let transport = Transport::new(config);
        let local_bind: SocketAddr = "0.0.0.0:0".parse().map_err(|err| {
            RegistryError::backend(format!("invalid bind address configuration: {err:?}"))
        })?;
        let handle = transport
            .bind(local_bind)
            .map_err(|err| RegistryError::backend(format!("transport bind failed: {err:?}")))?;

        Ok(Self {
            handle,
            registry_addr,
            agent_endpoint,
        })
    }

    fn agent_id(manifest: &AgentManifest) -> String {
        manifest.id().to_string()
    }

    fn manifest_to_register_request(&self, manifest: &AgentManifest) -> RegisterRequest {
        let capabilities = manifest
            .capabilities()
            .iter()
            .map(|cap| cap.id().as_str().to_string())
            .collect::<Vec<_>>();

        let mut metadata = HashMap::new();
        metadata.insert("version".to_string(), manifest.version().to_string());
        if let Some(description) = manifest.description() {
            metadata.insert("description".to_string(), description.to_string());
        }
        if !manifest.tags().is_empty() {
            metadata.insert(
                "tags".to_string(),
                serde_json::to_string(manifest.tags()).unwrap_or_default(),
            );
        }

        RegisterRequest {
            id: Self::agent_id(manifest),
            name: manifest.name().to_string(),
            capabilities,
            address: self.agent_endpoint,
            metadata,
        }
    }

    fn send_request_blocking(
        handle: &TransportHandle,
        registry_addr: SocketAddr,
        message: &Message,
    ) -> RegistryResult<Message> {
        let encoded = message.encode();
        let message_id = message.message_id();

        handle
            .send(&encoded, registry_addr)
            .map_err(|err| RegistryError::backend(format!("send failed: {err:?}")))?;

        let mut buffer = handle.acquire_buffer();
        let response = loop {
            match handle.receive(&mut buffer) {
                Ok((_len, _addr)) => {
                    let payload = buffer.as_slice().to_vec();
                    match Message::decode(payload) {
                        Ok(response) => {
                            if response.message_id() == message_id {
                                break response;
                            }
                        }
                        Err(err) => {
                            return Err(RegistryError::backend(format!(
                                "failed to decode registry response: {err:?}"
                            )));
                        }
                    }
                }
                Err(SocketError::Io(err))
                    if matches!(
                        err.kind(),
                        ErrorKind::WouldBlock | ErrorKind::TimedOut | ErrorKind::Interrupted
                    ) =>
                {
                    if err.kind() == ErrorKind::Interrupted {
                        debug!("registry receive interrupted; retrying");
                        continue;
                    }
                    return Err(RegistryError::backend(
                        "timed out waiting for registry response",
                    ));
                }
                Err(SocketError::Io(err)) => {
                    return Err(RegistryError::backend(format!(
                        "registry receive failed: {err:?}"
                    )));
                }
            }
        };

        Ok(response)
    }

    async fn send_request(&self, message: Message) -> RegistryResult<Message> {
        let handle = self.handle.clone();
        let registry_addr = self.registry_addr;
        tokio::task::spawn_blocking(move || {
            Self::send_request_blocking(&handle, registry_addr, &message)
        })
        .await
        .map_err(|err| RegistryError::backend(format!("registry task join error: {err:?}")))?
    }

    fn handle_error_message(message: &Message) -> RegistryResult<()> {
        let payload =
            serde_json::from_slice::<ErrorResponse>(message.payload()).map_err(|err| {
                RegistryError::backend(format!("failed to parse registry error payload: {err:?}"))
            })?;
        Err(RegistryError::backend(payload.error))
    }
}

fn default_transport_config() -> TransportConfig {
    TransportConfig {
        buffer_size: 16 * 1024,
        max_buffers: 256,
        read_timeout: Some(Duration::from_secs(5)),
        write_timeout: Some(Duration::from_secs(5)),
        #[cfg(feature = "debug-tools")]
        pcap_send_path: None,
        #[cfg(feature = "debug-tools")]
        pcap_recv_path: None,
    }
}

/// Trait implemented by discovery/registry backends.
#[async_trait]
pub trait AgentRegistry: Send + Sync {
    /// Registers an agent manifest with the mesh.
    async fn register(&self, manifest: &AgentManifest) -> RegistryResult<()>;

    /// Sends a heartbeat for an already registered agent.
    async fn heartbeat(&self, manifest: &AgentManifest) -> RegistryResult<()>;

    /// Removes the agent from the registry.
    async fn deregister(&self, manifest: &AgentManifest) -> RegistryResult<()>;
}

pub(crate) struct RegistrationController {
    registry: Arc<dyn AgentRegistry>,
    manifest: Arc<AgentManifest>,
    config: RegistrationConfig,
    shutdown: Arc<AtomicBool>,
    worker: Option<JoinHandle<()>>,
}

impl fmt::Debug for RegistrationController {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("RegistrationController")
            .field("registry", &"dyn AgentRegistry")
            .field("manifest", &self.manifest.id())
            .field("config", &self.config)
            .field("shutdown", &self.shutdown.load(Ordering::Relaxed))
            .field("worker", &self.worker.is_some())
            .finish()
    }
}

impl RegistrationController {
    pub(crate) fn new(
        registry: Arc<dyn AgentRegistry>,
        manifest: AgentManifest,
        config: RegistrationConfig,
    ) -> Self {
        Self {
            registry,
            manifest: Arc::new(manifest),
            config,
            shutdown: Arc::new(AtomicBool::new(false)),
            worker: None,
        }
    }

    pub(crate) fn on_state_change(
        &mut self,
        state: AgentState,
        scheduler: &TaskScheduler,
    ) -> RegistryResult<()> {
        match state {
            AgentState::Ready | AgentState::Active => {
                self.ensure_worker(scheduler)?;
            }
            AgentState::Retiring | AgentState::Terminated => {
                self.shutdown.store(true, Ordering::Release);
                self.spawn_deregister(scheduler)?;
                if let Some(handle) = self.worker.take() {
                    handle.abort();
                }
            }
            _ => {}
        }

        Ok(())
    }

    fn ensure_worker(&mut self, scheduler: &TaskScheduler) -> RegistryResult<()> {
        if self.worker.is_some() {
            return Ok(());
        }

        self.config.validate()?;

        let registry = Arc::clone(&self.registry);
        let manifest = Arc::clone(&self.manifest);
        let shutdown = Arc::clone(&self.shutdown);
        let config = self.config;

        let handle = scheduler.spawn(async move {
            run_registration_loop(registry, manifest, shutdown, config).await;
        })?;

        self.worker = Some(handle);
        Ok(())
    }

    fn spawn_deregister(&self, scheduler: &TaskScheduler) -> RegistryResult<()> {
        let registry = Arc::clone(&self.registry);
        let manifest = Arc::clone(&self.manifest);
        scheduler.spawn(async move {
            if let Err(err) = registry.deregister(&manifest).await {
                warn!(?err, "agent deregistration failed");
            } else {
                info!(agent_id = %manifest.id(), "agent deregistered");
            }
        })?;
        Ok(())
    }
}

async fn run_registration_loop(
    registry: Arc<dyn AgentRegistry>,
    manifest: Arc<AgentManifest>,
    shutdown: Arc<AtomicBool>,
    config: RegistrationConfig,
) {
    let mut retry_delay = config.initial_retry_delay();

    loop {
        if shutdown.load(Ordering::Acquire) {
            break;
        }

        match registry.register(&manifest).await {
            Ok(()) => {
                info!(agent_id = %manifest.id(), "agent registered with mesh");
                retry_delay = config.initial_retry_delay();
                if !run_heartbeat_loop(
                    Arc::clone(&registry),
                    Arc::clone(&manifest),
                    Arc::clone(&shutdown),
                    config,
                )
                .await
                {
                    continue;
                }
                break;
            }
            Err(err) => {
                warn!(?err, "agent registration failed; retrying");
                sleep(retry_delay).await;
                retry_delay = (retry_delay * 2).min(config.max_retry_delay());
            }
        }
    }
}

async fn run_heartbeat_loop(
    registry: Arc<dyn AgentRegistry>,
    manifest: Arc<AgentManifest>,
    shutdown: Arc<AtomicBool>,
    config: RegistrationConfig,
) -> bool {
    let mut failures: usize = 0;
    let mut interval = tokio::time::interval(config.heartbeat_interval());
    interval.set_missed_tick_behavior(MissedTickBehavior::Delay);

    while !shutdown.load(Ordering::Acquire) {
        interval.tick().await;
        if shutdown.load(Ordering::Acquire) {
            break;
        }

        match registry.heartbeat(&manifest).await {
            Ok(()) => {
                failures = 0;
            }
            Err(err) => {
                failures += 1;
                warn!(?err, failures, "heartbeat failure");
                if failures >= config.max_consecutive_failures().get() {
                    warn!(
                        failures,
                        "heartbeat failure threshold reached; attempting re-registration"
                    );
                    return false;
                }
            }
        }
    }

    true
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::AtomicUsize;

    use agent_primitives::{AgentId, Capability, CapabilityId};

    struct MockRegistry {
        registers: Arc<AtomicUsize>,
        heartbeats: Arc<AtomicUsize>,
        deregistrations: Arc<AtomicUsize>,
    }

    #[async_trait]
    impl AgentRegistry for MockRegistry {
        async fn register(&self, _manifest: &AgentManifest) -> RegistryResult<()> {
            self.registers.fetch_add(1, Ordering::SeqCst);
            Ok(())
        }

        async fn heartbeat(&self, _manifest: &AgentManifest) -> RegistryResult<()> {
            self.heartbeats.fetch_add(1, Ordering::SeqCst);
            Ok(())
        }

        async fn deregister(&self, _manifest: &AgentManifest) -> RegistryResult<()> {
            self.deregistrations.fetch_add(1, Ordering::SeqCst);
            Ok(())
        }
    }

    fn manifest() -> AgentManifest {
        let capability = Capability::builder(CapabilityId::new("mock.cap").unwrap())
            .name("Mock")
            .unwrap()
            .version("1.0.0")
            .unwrap()
            .add_scope("read:mock")
            .unwrap()
            .build()
            .unwrap();

        AgentManifest::builder(AgentId::random())
            .name("mock-agent")
            .unwrap()
            .version("0.1.0")
            .unwrap()
            .capabilities(vec![capability])
            .build()
            .unwrap()
    }

    #[tokio::test]
    async fn lifecycle_starts_and_stops_heartbeat() {
        let registry = Arc::new(MockRegistry {
            registers: Arc::new(AtomicUsize::new(0)),
            heartbeats: Arc::new(AtomicUsize::new(0)),
            deregistrations: Arc::new(AtomicUsize::new(0)),
        });

        let manifest = manifest();
        let config = RegistrationConfig::new(
            Duration::from_millis(10),
            Duration::from_millis(5),
            Duration::from_millis(20),
            NonZeroUsize::new(3).unwrap(),
        );

        let mut controller = RegistrationController::new(registry.clone(), manifest, config);
        let scheduler = TaskScheduler::default();

        controller
            .on_state_change(AgentState::Ready, &scheduler)
            .unwrap();

        tokio::time::sleep(Duration::from_millis(40)).await;

        assert!(registry.registers.load(Ordering::SeqCst) >= 1);
        assert!(registry.heartbeats.load(Ordering::SeqCst) >= 1);

        controller
            .on_state_change(AgentState::Retiring, &scheduler)
            .unwrap();
        tokio::time::sleep(Duration::from_millis(20)).await;
        assert!(registry.deregistrations.load(Ordering::SeqCst) >= 1);
    }
}

#[async_trait]
impl AgentRegistry for MxpRegistryClient {
    async fn register(&self, manifest: &AgentManifest) -> RegistryResult<()> {
        let request = self.manifest_to_register_request(manifest);
        let payload = serde_json::to_vec(&request)
            .map_err(|err| RegistryError::backend(format!("encode register payload: {err:?}")))?;
        let message = Message::new(MessageType::AgentRegister, payload);
        let response = self.send_request(message).await?;

        match response.message_type() {
            Some(MessageType::Response) => {
                let ack = serde_json::from_slice::<RegisterResponse>(response.payload()).map_err(
                    |err| {
                        RegistryError::backend(format!("parse register response failed: {err:?}"))
                    },
                )?;
                if ack.success {
                    debug!(agent_id = ack.agent_id, "registry registration acked");
                    Ok(())
                } else {
                    Err(RegistryError::backend(format!(
                        "registry rejected registration: {}",
                        ack.message
                    )))
                }
            }
            Some(MessageType::Error) => Self::handle_error_message(&response),
            other => Err(RegistryError::backend(format!(
                "unexpected message type {other:?} for register response"
            ))),
        }
    }

    async fn heartbeat(&self, manifest: &AgentManifest) -> RegistryResult<()> {
        let request = HeartbeatRequest {
            agent_id: Self::agent_id(manifest),
        };
        let payload = serde_json::to_vec(&request)
            .map_err(|err| RegistryError::backend(format!("encode heartbeat payload: {err:?}")))?;
        let message = Message::new(MessageType::AgentHeartbeat, payload);
        let response = self.send_request(message).await?;

        match response.message_type() {
            Some(MessageType::Response) => {
                let ack = serde_json::from_slice::<HeartbeatResponse>(response.payload()).map_err(
                    |err| {
                        RegistryError::backend(format!("parse heartbeat response failed: {err:?}"))
                    },
                )?;
                if ack.success && !ack.needs_register {
                    Ok(())
                } else if ack.needs_register {
                    Err(RegistryError::backend("registry requested re-registration"))
                } else {
                    Err(RegistryError::backend(
                        ack.message
                            .unwrap_or_else(|| "heartbeat rejected".to_string()),
                    ))
                }
            }
            Some(MessageType::Error) => Self::handle_error_message(&response),
            other => Err(RegistryError::backend(format!(
                "unexpected message type {other:?} for heartbeat response"
            ))),
        }
    }

    async fn deregister(&self, manifest: &AgentManifest) -> RegistryResult<()> {
        let request = HeartbeatRequest {
            agent_id: Self::agent_id(manifest),
        };
        let payload = serde_json::to_vec(&request)
            .map_err(|err| RegistryError::backend(format!("encode deregister payload: {err:?}")))?;
        let mut message = Message::new(MessageType::AgentHeartbeat, payload);
        message.set_flags(message.flags().with(Flags::FINAL));
        let response = self.send_request(message).await?;

        match response.message_type() {
            Some(MessageType::Response) => {
                let ack = serde_json::from_slice::<HeartbeatResponse>(response.payload()).map_err(
                    |err| {
                        RegistryError::backend(format!("parse deregister response failed: {err:?}"))
                    },
                )?;
                if ack.success {
                    Ok(())
                } else {
                    Err(RegistryError::backend(
                        ack.message
                            .unwrap_or_else(|| "deregister failed".to_string()),
                    ))
                }
            }
            Some(MessageType::Error) => Self::handle_error_message(&response),
            other => Err(RegistryError::backend(format!(
                "unexpected message type {other:?} for deregister response"
            ))),
        }
    }
}