klieo-ops 3.5.0

//! `KvSupervisor` default impl on top of `klieo_core::KvStore`.

use crate::supervisor::trait_::{Status, Supervisor, SupervisorError, SupervisorEvent};
use crate::types::{AgentId, AgentMeta, KillReason, KillTrigger, RuntimeState};
use async_trait::async_trait;
use bytes::Bytes;
use chrono::{DateTime, Duration, Utc};
use dashmap::DashMap;
use futures_core::stream::BoxStream;
use klieo_core::KvStore;
use std::sync::Arc;
use tokio::sync::broadcast;
use tokio_stream::wrappers::BroadcastStream;
use tokio_stream::StreamExt;
use tracing;

const BUCKET_HEARTBEAT: &str = "ops.supervisor.heartbeat";
const BUCKET_STATE: &str = "ops.supervisor.state";
const KEY_RUNTIME_STATE: &str = "runtime_state";

/// Default supervisor backed by a `KvStore` for heartbeat TTL and runtime
/// state. Events broadcast to all `watch` subscribers via a Tokio broadcast
/// channel.
pub struct KvSupervisor {
    kv: Arc<dyn KvStore>,
    heartbeat_ttl: Duration,
    metas: DashMap<AgentId, AgentMeta>,
    events: broadcast::Sender<SupervisorEvent>,
}

impl KvSupervisor {
    /// Build with safe defaults (5 s heartbeat TTL).
    #[must_use]
    pub fn with_defaults(kv: Arc<dyn KvStore>) -> Self {
        Self::new(kv, Duration::seconds(5))
    }

    /// Build with explicit TTL.
    #[must_use]
    pub fn new(kv: Arc<dyn KvStore>, heartbeat_ttl: Duration) -> Self {
        let (tx, _rx) = broadcast::channel(256);
        Self {
            kv,
            heartbeat_ttl,
            metas: DashMap::new(),
            events: tx,
        }
    }

    fn emit(&self, ev: SupervisorEvent) {
        if let Err(err) = self.events.send(ev) {
            // No live receivers — operators that subscribed via watch() may
            // miss this event. Surface to tracing rather than swallow.
            tracing::warn!(
                event = ?err.0,
                "supervisor event has no live receivers; subscriber may have dropped"
            );
        }
    }

    async fn read_runtime_state(&self) -> RuntimeState {
        match self.kv.get(BUCKET_STATE, KEY_RUNTIME_STATE).await {
            Ok(Some(entry)) => {
                let raw = std::str::from_utf8(&entry.value).unwrap_or("");
                match raw {
                    "draining" => RuntimeState::Draining,
                    "halted" => RuntimeState::Halted,
                    "running" => RuntimeState::Running,
                    _ => {
                        // Unrecognised payload — fail CLOSED per spec §3.5.
                        tracing::warn!(
                            raw,
                            "unrecognised runtime_state payload; defaulting to Halted"
                        );
                        RuntimeState::Halted
                    }
                }
            }
            Ok(None) => RuntimeState::Running,
            // Fail-CLOSED per spec §3.5: assume Halted if we cannot read.
            Err(_) => RuntimeState::Halted,
        }
    }

    async fn write_runtime_state(&self, s: RuntimeState) -> Result<(), SupervisorError> {
        self.kv
            .put(BUCKET_STATE, KEY_RUNTIME_STATE, Bytes::from(s.to_string()))
            .await
            .map(|_| ())
            .map_err(|e| SupervisorError::Storage(e.to_string()))
    }
}

#[async_trait]
impl Supervisor for KvSupervisor {
    async fn register(&self, agent: AgentId, meta: AgentMeta) -> Result<(), SupervisorError> {
        self.metas.insert(agent.clone(), meta);
        self.emit(SupervisorEvent::Registered(agent));
        Ok(())
    }

    async fn heartbeat(&self, agent: AgentId) -> Result<(), SupervisorError> {
        if !self.metas.contains_key(&agent) {
            return Err(SupervisorError::UnknownAgent(agent));
        }
        let expiry: DateTime<Utc> = Utc::now() + self.heartbeat_ttl;
        let bytes = Bytes::from(expiry.to_rfc3339());
        self.kv
            .put(BUCKET_HEARTBEAT, &agent.0, bytes)
            .await
            .map_err(|e| SupervisorError::Storage(e.to_string()))?;
        self.emit(SupervisorEvent::Heartbeat(agent));
        Ok(())
    }

    async fn report(&self, agent: AgentId, _status: Status) -> Result<(), SupervisorError> {
        if !self.metas.contains_key(&agent) {
            return Err(SupervisorError::UnknownAgent(agent));
        }
        Ok(())
    }

    async fn watch(&self) -> BoxStream<'static, SupervisorEvent> {
        let rx = self.events.subscribe();
        let stream = BroadcastStream::new(rx).filter_map(|res| res.ok());
        Box::pin(stream)
    }

    async fn trip_kill_switch(
        &self,
        reason: KillReason,
        _trigger: KillTrigger,
    ) -> Result<(), SupervisorError> {
        self.write_runtime_state(RuntimeState::Halted).await?;
        self.emit(SupervisorEvent::KillSwitchTripped { reason: reason.0 });
        Ok(())
    }

    async fn runtime_state(&self) -> RuntimeState {
        self.read_runtime_state().await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::supervisor::trait_::Supervisor;
    use bytes::Bytes;
    use klieo_bus_memory::MemoryBus;

    fn make_meta(id: &str) -> AgentMeta {
        AgentMeta {
            id: AgentId(id.into()),
            role: "test-role".into(),
            version: "0.0.0".into(),
            identity_pubkey: [0u8; 32],
            expected_step_p99: None,
        }
    }

    #[tokio::test]
    async fn unknown_runtime_state_payload_is_halted() {
        let bus = MemoryBus::new();
        let kv: Arc<dyn klieo_core::KvStore> = bus.kv.clone();
        let sup = KvSupervisor::with_defaults(kv.clone());

        kv.put(BUCKET_STATE, KEY_RUNTIME_STATE, Bytes::from("xyzzy"))
            .await
            .expect("put ok");

        assert_eq!(
            sup.read_runtime_state().await,
            RuntimeState::Halted,
            "unrecognised payload must default to Halted (fail-CLOSED per spec §3.5)"
        );
    }

    #[tokio::test]
    async fn register_then_heartbeat_succeeds() {
        let bus = MemoryBus::new();
        let kv: Arc<dyn klieo_core::KvStore> = bus.kv.clone();
        let sup = KvSupervisor::with_defaults(kv.clone());
        let id = AgentId("agent-hb".into());
        sup.register(id.clone(), make_meta("agent-hb"))
            .await
            .expect("register ok");
        let result = sup.heartbeat(id.clone()).await;
        assert!(result.is_ok(), "heartbeat after register must succeed");

        // Verify the heartbeat actually wrote an expiry entry to the KV store.
        let entry = kv
            .get(BUCKET_HEARTBEAT, &id.0)
            .await
            .expect("kv.get must not fail");
        assert!(
            entry.is_some(),
            "heartbeat must write expiry entry to KV bucket '{BUCKET_HEARTBEAT}' under key '{}'",
            id.0
        );
        let entry_owned = entry.unwrap();
        let value = std::str::from_utf8(&entry_owned.value).expect("entry value must be utf-8");
        assert!(
            DateTime::parse_from_rfc3339(value).is_ok(),
            "heartbeat KV value must be a valid RFC-3339 datetime, got: {value}"
        );
    }

    #[tokio::test]
    async fn heartbeat_for_unregistered_agent_returns_err() {
        let bus = MemoryBus::new();
        let kv: Arc<dyn klieo_core::KvStore> = bus.kv.clone();
        let sup = KvSupervisor::with_defaults(kv);
        let id = AgentId("ghost-agent".into());
        let err = sup
            .heartbeat(id.clone())
            .await
            .expect_err("must error for unknown agent");
        assert!(
            matches!(err, SupervisorError::UnknownAgent(ref a) if a == &id),
            "expected UnknownAgent, got {err:?}"
        );
    }
}