Skip to main content

lc_a2a/
router.rs

1//! Skill-based chain routing (P2-4).
2//!
3//! `A2AServer` can be handed a [`SkillRouter`] so that incoming `tasks/send`
4//! requests carrying a `skillId` are dispatched to a different underlying
5//! chain than the default one, based on the skills advertised on the agent
6//! card. [`SkillMapRouter`] is the concrete static mapping shipped here.
7
8use std::collections::HashMap;
9use std::sync::Arc;
10
11use lc_chains::base::BaseChain;
12
13/// Resolves a chain for a requested skill id (P2-4).
14///
15/// Return `None` to fall back to the server's default chain. Implementations
16/// must be cheap and infallible — they run on every `tasks/send`.
17pub trait SkillRouter: Send + Sync {
18    /// The chain that should handle `skill_id`, if any.
19    fn chain_for(&self, skill_id: &str) -> Option<Arc<dyn BaseChain>>;
20}
21
22/// Static `skill_id -> chain` mapping.
23///
24/// The keys should match the skill ids advertised on the agent card so that
25/// clients can discover which skills are routable.
26#[derive(Default)]
27pub struct SkillMapRouter {
28    skills: HashMap<String, Arc<dyn BaseChain>>,
29}
30
31impl SkillMapRouter {
32    /// Create an empty router (all requests fall through to the default chain).
33    pub fn new() -> Self {
34        Self::default()
35    }
36
37    /// Register a chain for a skill id.
38    pub fn with_skill(mut self, skill_id: impl Into<String>, chain: Arc<dyn BaseChain>) -> Self {
39        self.skills.insert(skill_id.into(), chain);
40        self
41    }
42}
43
44impl SkillRouter for SkillMapRouter {
45    fn chain_for(&self, skill_id: &str) -> Option<Arc<dyn BaseChain>> {
46        self.skills.get(skill_id).cloned()
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use super::*;
53    use lc_chains::base::{ChainError, ChainResult};
54    use serde_json::Value;
55
56    /// A trivial chain that names itself, so routing can be observed.
57    struct NamedChain(String);
58
59    #[async_trait::async_trait]
60    impl BaseChain for NamedChain {
61        fn input_keys(&self) -> Vec<&str> {
62            vec!["input"]
63        }
64
65        fn output_keys(&self) -> Vec<&str> {
66            vec!["output"]
67        }
68
69        async fn invoke(&self, _inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
70            let mut out = HashMap::new();
71            out.insert("output".to_string(), Value::String(self.0.clone()));
72            Ok(out)
73        }
74
75        fn name(&self) -> &str {
76            &self.0
77        }
78    }
79
80    fn arc_named(name: &str) -> Arc<dyn BaseChain> {
81        Arc::new(NamedChain(name.to_string()))
82    }
83
84    #[test]
85    fn empty_router_falls_through() {
86        let router = SkillMapRouter::new();
87        assert!(router.chain_for("anything").is_none());
88    }
89
90    #[test]
91    fn routes_by_skill_id() {
92        let router = SkillMapRouter::new()
93            .with_skill("research", arc_named("research-chain"))
94            .with_skill("summarize", arc_named("summary-chain"));
95
96        assert!(router.chain_for("research").is_some());
97        assert!(router.chain_for("summarize").is_some());
98        // Unknown skill falls through to the default chain.
99        assert!(router.chain_for("nope").is_none());
100        // Distinct chains per skill.
101        let a = router.chain_for("research").unwrap();
102        let b = router.chain_for("summarize").unwrap();
103        assert_ne!(a.name(), b.name());
104    }
105}