1use std::collections::HashMap;
3use std::future::Future;
4use std::pin::Pin;
5use std::sync::Arc;
6
7use cruxx_core::prelude::{Agent, CruxCtx, CruxErr};
8use serde_json::Value;
9
10pub type BoxHandler = Arc<
12 dyn Fn(Value) -> Pin<Box<dyn Future<Output = Result<Value, CruxErr>> + Send>> + Send + Sync,
13>;
14
15pub type BoxAgentRunner = Arc<
19 dyn Fn(Value) -> Pin<Box<dyn Future<Output = Result<Value, CruxErr>> + Send>> + Send + Sync,
20>;
21
22pub struct HandlerRegistry {
24 handlers: HashMap<String, BoxHandler>,
25 agents: HashMap<String, BoxAgentRunner>,
26}
27
28impl HandlerRegistry {
29 pub fn new() -> Self {
30 Self {
31 handlers: HashMap::new(),
32 agents: HashMap::new(),
33 }
34 }
35
36 pub fn handler<F, Fut>(&mut self, name: impl Into<String>, f: F)
38 where
39 F: Fn(Value) -> Fut + Send + Sync + 'static,
40 Fut: Future<Output = Result<Value, CruxErr>> + Send + 'static,
41 {
42 let name = name.into();
43 self.handlers
44 .insert(name, Arc::new(move |v| Box::pin(f(v))));
45 }
46
47 pub fn agent<A>(&mut self, name: impl Into<String>)
51 where
52 A: Agent<Input = Value, Output = Value>,
53 {
54 let name_str = name.into();
55 let agent_name = name_str.clone();
56 self.agents.insert(
57 name_str,
58 Arc::new(move |input: Value| {
59 let n = agent_name.clone();
60 Box::pin(async move {
61 let mut ctx = CruxCtx::new(&n);
62 A::run(&mut ctx, input).await
63 }) as Pin<Box<dyn Future<Output = Result<Value, CruxErr>> + Send>>
64 }),
65 );
66 }
67
68 pub fn get_handler(&self, name: &str) -> Option<&BoxHandler> {
70 self.handlers.get(name)
71 }
72
73 pub fn get_agent(&self, name: &str) -> Option<&BoxAgentRunner> {
75 self.agents.get(name)
76 }
77}
78
79impl Default for HandlerRegistry {
80 fn default() -> Self {
81 Self::new()
82 }
83}