klieo-workflow 3.4.0

//! Compile-time authz allow-list + runtime handles for workflow refs.

use klieo_core::llm::LlmClient;
use klieo_flows::Flow;
use std::collections::{HashMap, HashSet};
use std::sync::Arc;

/// The set of primitives a workflow is permitted to reference. Anything
/// not registered here is rejected at compile time — this IS the authz
/// boundary for no-code workflows.
#[derive(Default, Clone)]
pub struct Registry {
    models: HashMap<String, Arc<dyn LlmClient>>,
    tools: HashSet<String>,
    subflows: HashMap<String, Arc<dyn Flow>>,
}

impl Registry {
    /// Empty registry.
    pub fn new() -> Self {
        Self::default()
    }

    /// Register an LLM client under `id`. Later registrations overwrite.
    pub fn with_model(mut self, id: impl Into<String>, llm: Arc<dyn LlmClient>) -> Self {
        self.models.insert(id.into(), llm);
        self
    }

    /// Allow a tool id. The tool impl itself is resolved at runtime via
    /// the context's `ToolInvoker`.
    pub fn with_tool(mut self, id: impl Into<String>) -> Self {
        self.tools.insert(id.into());
        self
    }

    /// Register a prebuilt sub-flow under `id`.
    pub fn with_subflow(mut self, id: impl Into<String>, flow: Arc<dyn Flow>) -> Self {
        self.subflows.insert(id.into(), flow);
        self
    }

    /// Resolve a model handle.
    pub fn model(&self, id: &str) -> Option<Arc<dyn LlmClient>> {
        self.models.get(id).cloned()
    }

    /// Whether a tool id is allowed.
    pub fn allows_tool(&self, id: &str) -> bool {
        self.tools.contains(id)
    }

    /// Resolve a sub-flow handle.
    pub fn subflow(&self, id: &str) -> Option<Arc<dyn Flow>> {
        self.subflows.get(id).cloned()
    }

    /// The id of every registered model, in unspecified order. The palette
    /// of models a workflow may name in an agent node.
    pub fn model_ids(&self) -> Vec<&str> {
        self.models.keys().map(String::as_str).collect()
    }

    /// The id of every allowed tool, in unspecified order. The palette of
    /// tools a workflow may invoke.
    pub fn tool_ids(&self) -> Vec<&str> {
        self.tools.iter().map(String::as_str).collect()
    }

    /// The id of every registered sub-flow, in unspecified order. The
    /// palette of sub-flows a workflow may reference.
    pub fn subflow_ids(&self) -> Vec<&str> {
        self.subflows.keys().map(String::as_str).collect()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;
    use klieo_core::agent::AgentContext;
    use klieo_flows::FlowError;
    use serde_json::Value;

    struct NoopFlow;

    #[async_trait]
    impl Flow for NoopFlow {
        fn name(&self) -> &str {
            "noop"
        }
        async fn run(&self, _ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
            Ok(input)
        }
    }

    #[test]
    fn tool_allow_list_distinguishes_registered_from_absent() {
        let r = Registry::new().with_tool("lookup");
        assert!(r.allows_tool("lookup"));
        assert!(!r.allows_tool("delete_everything"));
    }

    #[test]
    fn absent_model_resolves_to_none() {
        let r = Registry::new();
        assert!(r.model("ghost").is_none());
    }

    #[test]
    fn subflow_lookup_distinguishes_registered_from_absent() {
        let r = Registry::new().with_subflow("fraud", Arc::new(NoopFlow));
        assert!(r.subflow("fraud").is_some());
        assert!(r.subflow("unknown").is_none());
    }

    #[test]
    fn enumerates_registered_palette_ids() {
        let llm = crate::test_support::dummy_llm();
        let r = Registry::new()
            .with_model("m1", llm.clone())
            .with_model("m2", llm)
            .with_tool("lookup")
            .with_subflow("fraud", Arc::new(NoopFlow));
        let mut models = r.model_ids();
        models.sort_unstable();
        assert_eq!(models, vec!["m1", "m2"]);
        assert_eq!(r.tool_ids(), vec!["lookup"]);
        assert_eq!(r.subflow_ids(), vec!["fraud"]);
    }

    #[test]
    fn empty_registry_enumerates_empty() {
        let r = Registry::new();
        assert!(r.model_ids().is_empty());
        assert!(r.tool_ids().is_empty());
        assert!(r.subflow_ids().is_empty());
    }
}