liminal_server/agents/

retained.rs

//! Retained metadata and live delivery share one lock. Each watcher has a
//! real bounded library inbox, so its initial value is charged like a change.

use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, OnceLock, Weak};

use liminal::channel::{
    ChannelConfig, ChannelHandle, ChannelMode, ChannelSupervisor, ConnectionInboxBudget,
    InboxInstall, Schema, SubscriptionHandle,
};
use liminal::envelope::PublisherId;
use liminal::protocol::SchemaId;

use super::config::AgentsConfig;
use super::error::AgentError;
use crate::server::connection::services::{ConnectionSubscription, SubscriptionResource};

struct Watcher {
    subject: String,
    channel: ChannelHandle,
    overflow: Arc<AtomicBool>,
    wake: Option<liminal::channel::InboxNotifier>,
}

#[derive(Default)]
struct Kept {
    values: BTreeMap<String, serde_json::Value>,
    watchers: BTreeMap<u64, Watcher>,
}

/// Holds only server metadata. No agent message or working stream reaches it.
pub struct Retained {
    bounds: AgentsConfig,
    supervisor: ChannelSupervisor,
    publisher: PublisherId,
    kept: Mutex<Kept>,
    console: OnceLock<Weak<crate::console::ConsoleHub>>,
}

impl std::fmt::Debug for Retained {
    fn fmt(&self, out: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        out.debug_struct("Retained")
            .field("bounds", &self.bounds)
            .finish_non_exhaustive()
    }
}

impl Retained {
    /// Uses the existing console's signal, with no timer or second publisher.
    ///
    /// # Errors
    /// A console was already installed.
    pub fn bind_console(
        &self,
        console: &Arc<crate::console::ConsoleHub>,
    ) -> Result<(), AgentError> {
        self.console
            .set(Arc::downgrade(console))
            .map_err(|_| AgentError::Invalid("retained console already bound"))
    }
    /// Reads current metadata without consuming the watcher stream.
    ///
    /// # Errors
    /// The retained lock is poisoned.
    pub fn value(&self, subject: &str) -> Result<Option<serde_json::Value>, AgentError> {
        Ok(self
            .kept
            .lock()
            .map_err(|_| AgentError::Store("retained lock poisoned".into()))?
            .values
            .get(subject)
            .cloned())
    }
    /// Shares the server's channel scheduler; creates no worker or timer.
    #[must_use]
    pub fn new(bounds: AgentsConfig, supervisor: ChannelSupervisor, device: &str) -> Arc<Self> {
        Arc::new(Self {
            bounds,
            supervisor,
            publisher: PublisherId::new(format!("liminal.{device}")),
            kept: Mutex::new(Kept::default()),
            console: OnceLock::new(),
        })
    }

    /// Installs both presence and state atomically before delivering either.
    ///
    /// # Errors
    /// Metadata is oversized, a lock is poisoned, or a watcher cannot receive.
    /// A delivery failure does not erase the retained values.
    pub fn publish(&self, values: &[(String, serde_json::Value)]) -> Result<(), AgentError> {
        let mut encoded = BTreeMap::new();
        for (subject, value) in values {
            let bytes =
                serde_json::to_vec(value).map_err(|_| AgentError::Invalid("metadata encoding"))?;
            if bytes.len() > self.bounds.max_event_bytes {
                return Err(AgentError::Bound {
                    bound: "max_event_bytes",
                    limit: self.bounds.max_event_bytes,
                    count: bytes.len(),
                });
            }
            encoded.insert(subject.as_str(), bytes);
        }
        let mut kept = self
            .kept
            .lock()
            .map_err(|_| AgentError::Store("retained lock poisoned".into()))?;
        let changed = values
            .iter()
            .any(|(subject, value)| kept.values.get(subject) != Some(value));
        for (subject, value) in values {
            kept.values.insert(subject.clone(), value.clone());
        }
        for (id, watcher) in &kept.watchers {
            if values
                .iter()
                .any(|(subject, _)| subject == &watcher.subject)
            {
                let bytes = encoded
                    .get(watcher.subject.as_str())
                    .ok_or(AgentError::Invalid("retained encoding absent"))?;
                let delivery =
                    watcher
                        .channel
                        .publish_with_delivery(bytes, self.publisher.clone(), None);
                if delivery.is_err() || delivery.is_ok_and(|delivery| !delivery.is_admitted()) {
                    watcher.overflow.store(true, Ordering::Release);
                    if let Some(wake) = &watcher.wake {
                        wake();
                    }
                    tracing::warn!(watcher = id, "agent_retained_watcher_shed");
                }
            }
        }
        drop(kept);
        if changed {
            if let Some(console) = self.console.get().and_then(Weak::upgrade) {
                console
                    .agent_changed()
                    .map_err(|_| AgentError::Delivery("agent topology signal failed"))?;
            }
        }
        Ok(())
    }

    /// Subscribes under the same lock used by writers: the current value is
    /// always queued first and every later change follows it. The supplied
    /// wire inbox budget is installed before that first enqueue.
    ///
    /// # Errors
    /// Unknown subject, too many watchers, a refused schema, or failed delivery.
    pub fn subscribe(
        self: &Arc<Self>,
        id: u64,
        subject: &str,
        schemas: &[SchemaId],
        install: Option<InboxInstall>,
    ) -> Result<ConnectionSubscription, AgentError> {
        let schema_bytes = b"{}";
        let protocol_schema = SchemaId::from_schema_bytes(schema_bytes);
        if !schemas.is_empty() && !schemas.contains(&protocol_schema) {
            return Err(AgentError::Invalid("retained subscription schema"));
        }
        let mut kept = self
            .kept
            .lock()
            .map_err(|_| AgentError::Store("retained lock poisoned".into()))?;
        let value = kept
            .values
            .get(subject)
            .ok_or_else(|| AgentError::Unknown(subject.into()))?
            .clone();
        if kept.watchers.len() >= self.bounds.max_watchers {
            return Err(AgentError::Bound {
                bound: "max_watchers",
                limit: self.bounds.max_watchers,
                count: kept.watchers.len(),
            });
        }
        if kept.watchers.contains_key(&id) {
            return Err(AgentError::Invalid("duplicate watcher id"));
        }
        let install = install.unwrap_or_else(|| InboxInstall {
            budget: ConnectionInboxBudget::new(self.bounds.watcher_bytes),
            depth_cap: self.bounds.watcher_depth,
            notifier: None,
            capacity: None,
        });
        let wake = install.notifier.clone();
        let channel = ChannelHandle::with_supervisor(
            ChannelConfig::new(
                format!("agent-watch-{id}"),
                Schema::new(serde_json::json!({}))
                    .map_err(|_| AgentError::Invalid("metadata schema"))?,
                ChannelMode::Ephemeral,
            ),
            self.supervisor.clone(),
        );
        let subscription = channel
            .subscribe_with_install(install)
            .map_err(|_| AgentError::Delivery("retained subscribe refused"))?;
        let initial =
            serde_json::to_vec(&value).map_err(|_| AgentError::Invalid("metadata encoding"))?;
        let delivery = channel
            .publish_with_delivery(&initial, self.publisher.clone(), None)
            .map_err(|_| AgentError::Delivery("retained initial value refused"))?;
        if !delivery.is_admitted() {
            return Err(AgentError::Delivery("retained initial inbox full"));
        }
        let overflow = Arc::new(AtomicBool::new(false));
        kept.watchers.insert(
            id,
            Watcher {
                subject: subject.into(),
                channel: channel.clone(),
                overflow: Arc::clone(&overflow),
                wake,
            },
        );
        drop(kept);
        Ok(ConnectionSubscription::new(
            id,
            protocol_schema,
            Box::new(Resource {
                id,
                channel,
                subscription,
                owner: Arc::downgrade(self),
                overflow,
                closed: false,
            }),
        ))
    }

    fn detach(&self, id: u64) -> Result<(), AgentError> {
        self.kept
            .lock()
            .map_err(|_| AgentError::Store("retained lock poisoned".into()))?
            .watchers
            .remove(&id);
        Ok(())
    }
}

struct Resource {
    id: u64,
    channel: ChannelHandle,
    subscription: SubscriptionHandle,
    owner: Weak<Retained>,
    overflow: Arc<AtomicBool>,
    closed: bool,
}

impl std::fmt::Debug for Resource {
    fn fmt(&self, out: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        out.debug_struct("AgentWatch")
            .field("id", &self.id)
            .finish_non_exhaustive()
    }
}

impl Resource {
    fn release(&self) -> Result<(), crate::ServerError> {
        if let Some(owner) = self.owner.upgrade() {
            owner.detach(self.id)?;
        }
        Ok(())
    }
    fn finish(&mut self) -> Result<(), crate::ServerError> {
        if self.closed {
            return Ok(());
        }
        self.closed = true;
        let unsubscribed = self.channel.unsubscribe(&self.subscription);
        let closed = self.channel.close();
        let released = self.release();
        if unsubscribed.is_err() {
            tracing::error!(watcher = self.id, "agent_watch_unsubscribe_failed");
        }
        if closed.is_err() {
            tracing::error!(watcher = self.id, "agent_watch_actor_close_failed");
        }
        released?;
        if unsubscribed.is_err() {
            return Err(AgentError::Delivery("watch unsubscribe failed").into());
        }
        if closed.is_err() {
            return Err(AgentError::Delivery("watch actor close failed").into());
        }
        Ok(())
    }
}

impl SubscriptionResource for Resource {
    fn unsubscribe(mut self: Box<Self>) -> Result<(), crate::ServerError> {
        self.finish()
    }
    fn try_next(&mut self) -> Option<liminal::envelope::Envelope> {
        if let Ok(value) = self.subscription.try_next() {
            value
        } else {
            self.overflow.store(true, Ordering::Release);
            tracing::error!(watcher = self.id, "agent_watch_read_failed");
            None
        }
    }
    fn try_next_positioned(
        &mut self,
    ) -> Result<Option<(liminal::envelope::Envelope, Option<u64>)>, crate::ServerError> {
        self.subscription
            .try_next_positioned()
            .map_err(|_| AgentError::Delivery("watch read failed").into())
    }
    fn has_pending(&self) -> bool {
        self.overflow.load(Ordering::Acquire) || self.subscription.has_pending()
    }
    fn is_overflowed(&self) -> bool {
        self.overflow.load(Ordering::Acquire) || self.subscription.is_overflowed()
    }
    fn subscriber_pid(&self) -> Option<u64> {
        Some(self.subscription.subscriber_pid())
    }
}

impl Drop for Resource {
    fn drop(&mut self) {
        if self.finish().is_err() {
            tracing::error!("agent_watch_release_failed");
        }
    }
}