qae_kernel/

registry.rs

// SPDX-License-Identifier: BUSL-1.1
//! Runtime registry for constraint channels.
//!
//! Channels can be added, removed, and reconfigured without recompilation.
//! This is the single mechanism used by ALL domain adapters — finance, agentic,
//! and future domains.

use crate::constraint::ConstraintChannel;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

/// Runtime registry for constraint channels.
pub struct ConstraintRegistry {
    channels: BTreeMap<String, RegisteredChannel>,
}

/// A channel registered in the runtime registry.
pub struct RegisteredChannel {
    /// The constraint channel implementation.
    pub channel: Box<dyn ConstraintChannel>,
    /// Whether this channel is currently enabled.
    pub enabled: bool,
    /// Per-channel configuration.
    pub config: ChannelConfig,
}

/// Per-channel configuration — overrides global CertifierConfig thresholds.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChannelConfig {
    /// Per-channel weight (default 1.0). Higher = more influence on bottleneck.
    pub weight: f64,
    /// Optional per-channel override thresholds.
    pub safe_threshold: Option<f64>,
    pub caution_threshold: Option<f64>,
    pub block_threshold: Option<f64>,
    /// Human-readable description (shown in UI).
    pub description: String,
    /// Domain tag (e.g., "finance", "agentic", "custom").
    pub domain_tag: String,
}

impl Default for ChannelConfig {
    fn default() -> Self {
        Self {
            weight: 1.0,
            safe_threshold: None,
            caution_threshold: None,
            block_threshold: None,
            description: String::new(),
            domain_tag: "custom".to_string(),
        }
    }
}

/// Summary of a registered channel (for UI rendering).
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChannelSummary {
    pub id: String,
    pub name: String,
    pub enabled: bool,
    pub config: ChannelConfig,
    pub dimensions: Vec<String>,
}

impl ConstraintRegistry {
    pub fn new() -> Self {
        Self {
            channels: BTreeMap::new(),
        }
    }

    /// Register a new constraint channel at runtime.
    pub fn register(
        &mut self,
        id: &str,
        channel: Box<dyn ConstraintChannel>,
        config: ChannelConfig,
    ) {
        self.channels.insert(
            id.to_string(),
            RegisteredChannel {
                channel,
                enabled: true,
                config,
            },
        );
    }

    /// Remove a channel by id.
    pub fn unregister(&mut self, id: &str) -> bool {
        self.channels.remove(id).is_some()
    }

    /// Enable/disable a channel without removing it.
    pub fn set_enabled(&mut self, id: &str, enabled: bool) -> bool {
        if let Some(ch) = self.channels.get_mut(id) {
            ch.enabled = enabled;
            true
        } else {
            false
        }
    }

    /// Update a channel's config at runtime.
    pub fn update_config(&mut self, id: &str, config: ChannelConfig) -> bool {
        if let Some(ch) = self.channels.get_mut(id) {
            ch.config = config;
            true
        } else {
            false
        }
    }

    /// Return all enabled channels (used by SafetyCertifier).
    pub fn active_channels(&self) -> Vec<(&str, &RegisteredChannel)> {
        self.channels
            .iter()
            .filter(|(_, c)| c.enabled)
            .map(|(id, c)| (id.as_str(), c))
            .collect()
    }

    /// Number of registered channels.
    pub fn len(&self) -> usize {
        self.channels.len()
    }

    /// Whether the registry is empty.
    pub fn is_empty(&self) -> bool {
        self.channels.is_empty()
    }

    /// Number of active (enabled) channels.
    pub fn active_count(&self) -> usize {
        self.channels.values().filter(|c| c.enabled).count()
    }

    /// List all channels with their config (for UI rendering).
    pub fn list_all(&self) -> Vec<ChannelSummary> {
        self.channels
            .iter()
            .map(|(id, rc)| ChannelSummary {
                id: id.clone(),
                name: rc.channel.name().to_string(),
                enabled: rc.enabled,
                config: rc.config.clone(),
                dimensions: rc.channel.dimension_names(),
            })
            .collect()
    }

    /// Export all channel definitions as JSON (for constraint management API).
    pub fn export(&self) -> Vec<ChannelSummary> {
        self.list_all()
    }

    /// Check if a channel exists.
    pub fn contains(&self, id: &str) -> bool {
        self.channels.contains_key(id)
    }
}

impl Default for ConstraintRegistry {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::KernelResult;

    struct FixedChannel {
        name: String,
        margin: f64,
    }

    impl ConstraintChannel for FixedChannel {
        fn name(&self) -> &str {
            &self.name
        }
        fn evaluate(&self, _state: &[f64]) -> KernelResult<f64> {
            Ok(self.margin)
        }
        fn dimension_names(&self) -> Vec<String> {
            vec!["x".into()]
        }
    }

    fn make_channel(name: &str, margin: f64) -> Box<dyn ConstraintChannel> {
        Box::new(FixedChannel {
            name: name.to_string(),
            margin,
        })
    }

    fn make_config(domain: &str) -> ChannelConfig {
        ChannelConfig {
            domain_tag: domain.to_string(),
            description: format!("{} channel", domain),
            ..Default::default()
        }
    }

    #[test]
    fn register_and_list() {
        let mut reg = ConstraintRegistry::new();
        reg.register("ch1", make_channel("ch1", 0.8), make_config("test"));
        reg.register("ch2", make_channel("ch2", 0.5), make_config("test"));

        assert_eq!(reg.len(), 2);
        assert_eq!(reg.active_count(), 2);

        let summaries = reg.list_all();
        assert_eq!(summaries.len(), 2);
        assert_eq!(summaries[0].id, "ch1");
        assert_eq!(summaries[1].id, "ch2");
    }

    #[test]
    fn unregister_channel() {
        let mut reg = ConstraintRegistry::new();
        reg.register("ch1", make_channel("ch1", 0.8), make_config("test"));

        assert!(reg.unregister("ch1"));
        assert!(!reg.unregister("ch1")); // already removed
        assert_eq!(reg.len(), 0);
    }

    #[test]
    fn enable_disable_channel() {
        let mut reg = ConstraintRegistry::new();
        reg.register("ch1", make_channel("ch1", 0.8), make_config("test"));

        assert_eq!(reg.active_count(), 1);

        reg.set_enabled("ch1", false);
        assert_eq!(reg.active_count(), 0);
        assert_eq!(reg.len(), 1); // still registered

        reg.set_enabled("ch1", true);
        assert_eq!(reg.active_count(), 1);
    }

    #[test]
    fn update_config() {
        let mut reg = ConstraintRegistry::new();
        reg.register("ch1", make_channel("ch1", 0.8), make_config("test"));

        let new_config = ChannelConfig {
            weight: 2.0,
            safe_threshold: Some(0.7),
            ..make_config("updated")
        };
        assert!(reg.update_config("ch1", new_config));

        let summaries = reg.list_all();
        assert!((summaries[0].config.weight - 2.0).abs() < f64::EPSILON);
        assert_eq!(summaries[0].config.domain_tag, "updated");
    }

    #[test]
    fn active_channels_excludes_disabled() {
        let mut reg = ConstraintRegistry::new();
        reg.register("ch1", make_channel("ch1", 0.8), make_config("test"));
        reg.register("ch2", make_channel("ch2", 0.5), make_config("test"));

        reg.set_enabled("ch1", false);
        let active = reg.active_channels();
        assert_eq!(active.len(), 1);
        assert_eq!(active[0].0, "ch2");
    }

    #[test]
    fn contains_check() {
        let mut reg = ConstraintRegistry::new();
        reg.register("ch1", make_channel("ch1", 0.8), make_config("test"));

        assert!(reg.contains("ch1"));
        assert!(!reg.contains("ch2"));
    }

    #[test]
    fn deterministic_ordering() {
        let mut reg = ConstraintRegistry::new();
        reg.register("z_channel", make_channel("z", 0.1), make_config("test"));
        reg.register("a_channel", make_channel("a", 0.9), make_config("test"));
        reg.register("m_channel", make_channel("m", 0.5), make_config("test"));

        let summaries = reg.list_all();
        assert_eq!(summaries[0].id, "a_channel");
        assert_eq!(summaries[1].id, "m_channel");
        assert_eq!(summaries[2].id, "z_channel");
    }
}