klieo_workflow/
registry.rs

1//! Compile-time authz allow-list + runtime handles for workflow refs.
2
3use klieo_core::llm::LlmClient;
4use klieo_flows::Flow;
5use std::collections::{HashMap, HashSet};
6use std::sync::Arc;
7
8/// The set of primitives a workflow is permitted to reference. Anything
9/// not registered here is rejected at compile time — this IS the authz
10/// boundary for no-code workflows.
11#[derive(Default, Clone)]
12pub struct Registry {
13    models: HashMap<String, Arc<dyn LlmClient>>,
14    tools: HashSet<String>,
15    subflows: HashMap<String, Arc<dyn Flow>>,
16}
17
18impl Registry {
19    /// Empty registry.
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    /// Register an LLM client under `id`. Later registrations overwrite.
25    pub fn with_model(mut self, id: impl Into<String>, llm: Arc<dyn LlmClient>) -> Self {
26        self.models.insert(id.into(), llm);
27        self
28    }
29
30    /// Allow a tool id. The tool impl itself is resolved at runtime via
31    /// the context's `ToolInvoker`.
32    pub fn with_tool(mut self, id: impl Into<String>) -> Self {
33        self.tools.insert(id.into());
34        self
35    }
36
37    /// Register a prebuilt sub-flow under `id`.
38    pub fn with_subflow(mut self, id: impl Into<String>, flow: Arc<dyn Flow>) -> Self {
39        self.subflows.insert(id.into(), flow);
40        self
41    }
42
43    /// Resolve a model handle.
44    pub fn model(&self, id: &str) -> Option<Arc<dyn LlmClient>> {
45        self.models.get(id).cloned()
46    }
47
48    /// Whether a tool id is allowed.
49    pub fn allows_tool(&self, id: &str) -> bool {
50        self.tools.contains(id)
51    }
52
53    /// Resolve a sub-flow handle.
54    pub fn subflow(&self, id: &str) -> Option<Arc<dyn Flow>> {
55        self.subflows.get(id).cloned()
56    }
57
58    /// The id of every registered model, in unspecified order. The palette
59    /// of models a workflow may name in an agent node.
60    pub fn model_ids(&self) -> Vec<&str> {
61        self.models.keys().map(String::as_str).collect()
62    }
63
64    /// The id of every allowed tool, in unspecified order. The palette of
65    /// tools a workflow may invoke.
66    pub fn tool_ids(&self) -> Vec<&str> {
67        self.tools.iter().map(String::as_str).collect()
68    }
69
70    /// The id of every registered sub-flow, in unspecified order. The
71    /// palette of sub-flows a workflow may reference.
72    pub fn subflow_ids(&self) -> Vec<&str> {
73        self.subflows.keys().map(String::as_str).collect()
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80    use async_trait::async_trait;
81    use klieo_core::agent::AgentContext;
82    use klieo_flows::FlowError;
83    use serde_json::Value;
84
85    struct NoopFlow;
86
87    #[async_trait]
88    impl Flow for NoopFlow {
89        fn name(&self) -> &str {
90            "noop"
91        }
92        async fn run(&self, _ctx: AgentContext, input: Value) -> Result<Value, FlowError> {
93            Ok(input)
94        }
95    }
96
97    #[test]
98    fn tool_allow_list_distinguishes_registered_from_absent() {
99        let r = Registry::new().with_tool("lookup");
100        assert!(r.allows_tool("lookup"));
101        assert!(!r.allows_tool("delete_everything"));
102    }
103
104    #[test]
105    fn absent_model_resolves_to_none() {
106        let r = Registry::new();
107        assert!(r.model("ghost").is_none());
108    }
109
110    #[test]
111    fn subflow_lookup_distinguishes_registered_from_absent() {
112        let r = Registry::new().with_subflow("fraud", Arc::new(NoopFlow));
113        assert!(r.subflow("fraud").is_some());
114        assert!(r.subflow("unknown").is_none());
115    }
116
117    #[test]
118    fn enumerates_registered_palette_ids() {
119        let llm = crate::test_support::dummy_llm();
120        let r = Registry::new()
121            .with_model("m1", llm.clone())
122            .with_model("m2", llm)
123            .with_tool("lookup")
124            .with_subflow("fraud", Arc::new(NoopFlow));
125        let mut models = r.model_ids();
126        models.sort_unstable();
127        assert_eq!(models, vec!["m1", "m2"]);
128        assert_eq!(r.tool_ids(), vec!["lookup"]);
129        assert_eq!(r.subflow_ids(), vec!["fraud"]);
130    }
131
132    #[test]
133    fn empty_registry_enumerates_empty() {
134        let r = Registry::new();
135        assert!(r.model_ids().is_empty());
136        assert!(r.tool_ids().is_empty());
137        assert!(r.subflow_ids().is_empty());
138    }
139}
klieo_workflow/registry.rs

klieo_workflow/
registry.rs
