Skip to main content

switchyard_libsy/algorithms/
subagent.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Delegated sub-agent routing around an arbitrary parent algorithm.
5
6use std::sync::Arc;
7
8use switchyard_protocol::{Category, Metadata, Request};
9
10use super::fall_through::FallThrough;
11use super::util::affinity::{AffinityRouter, ClassifyTrigger};
12use super::util::subagent::SubagentGate;
13use crate::algorithms::llm_class::DefaultCategoryClassifier;
14use crate::core::algorithm::{Algorithm, Driver};
15use crate::core::classifier::Classifier;
16use crate::core::state::State;
17use crate::{LibsyError, Result, RoutingOutcome};
18
19/// Runtime components for delegated sub-agent routing.
20pub struct SubagentRouterConfig {
21    /// Classifier invoked for delegated work according to `classify_trigger`.
22    pub classifier: Arc<dyn Classifier<State>>,
23    /// Child model category used when `classifier` abstains.
24    pub default_target: Category,
25    /// Controls whether each child is classified once or on every request.
26    pub classify_trigger: ClassifyTrigger,
27    /// Unsupported for child routing because child identity must come from harness metadata.
28    pub message_hash_fallback: bool,
29}
30
31impl SubagentRouterConfig {
32    /// Routes all delegated work to the first model in the sub-agent `Any` category.
33    pub fn fixed_target() -> Self {
34        Self {
35            classifier: Arc::new(DefaultCategoryClassifier(Category::Any)),
36            default_target: Category::Any,
37            classify_trigger: ClassifyTrigger::EveryRequest,
38            message_hash_fallback: false,
39        }
40    }
41}
42
43/// Routes delegated work independently while preserving the parent algorithm for other traffic.
44pub struct SubagentRouter {
45    parent: Arc<dyn Algorithm>,
46    subagent: FallThrough<State>,
47}
48
49impl SubagentRouter {
50    /// Wraps `parent` with the configured delegated-work route.
51    ///
52    /// # Errors
53    ///
54    /// Returns an error when the affinity settings cannot identify delegated children safely.
55    pub fn new(parent: Arc<dyn Algorithm>, config: SubagentRouterConfig) -> Result<Self> {
56        if config.message_hash_fallback {
57            return Err(LibsyError::AlgorithmError {
58                message: "sub-agent routing cannot use message_hash_fallback".to_string(),
59            });
60        }
61
62        let mut subagent = match config.classify_trigger {
63            ClassifyTrigger::EveryRequest => FallThrough::new_with_state().with_name("subagent"),
64            ClassifyTrigger::NewSession => {
65                let affinity = Arc::new(AffinityRouter::for_subagents());
66                FallThrough::new_with_state()
67                    .with_name("subagent")
68                    .with_processor(affinity.clone())
69                    .with_classifier(affinity)
70            }
71            ClassifyTrigger::UserTurn => {
72                return Err(LibsyError::AlgorithmError {
73                    message: "sub-agent routing cannot use classify_trigger = user_turn"
74                        .to_string(),
75                });
76            }
77        };
78        subagent = subagent
79            .with_classifier(Arc::new(SubagentGate::new(config.classifier)))
80            .with_classifier(Arc::new(DefaultCategoryClassifier(config.default_target)));
81
82        Ok(Self { parent, subagent })
83    }
84}
85
86#[async_trait::async_trait]
87impl Algorithm for SubagentRouter {
88    fn name(&self) -> &str {
89        self.parent.name()
90    }
91
92    async fn route(self: Arc<Self>, driver: Driver, request: Request) -> Result<RoutingOutcome> {
93        if request
94            .metadata
95            .as_ref()
96            .is_some_and(Metadata::is_subagent_work)
97        {
98            // Delegated work routes over the sub-agent's own models, never the parent's.
99            self.subagent.execute(driver.for_subagent()?, request).await
100        } else {
101            self.parent.clone().route(driver, request).await
102        }
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use std::sync::Arc;
109    use std::sync::atomic::{AtomicUsize, Ordering};
110
111    use async_trait::async_trait;
112    use switchyard_protocol::{Category, Metadata, ModelId, Request, Response, text_request};
113
114    use super::{SubagentRouter, SubagentRouterConfig};
115    use crate::algorithms::passthrough::Passthrough;
116    use crate::core::classifier::{Classification, Classifier, Score};
117    use crate::core::testing::{echo, test_drive_with_models};
118    use crate::{ClassifyTrigger, Driver, RuntimeModels, State};
119
120    struct ScriptedClassifier {
121        calls: AtomicUsize,
122    }
123
124    #[async_trait]
125    impl Classifier<State> for ScriptedClassifier {
126        async fn score(
127            &self,
128            _state: &mut State,
129            _request: &mut Request,
130            driver: &Driver,
131        ) -> crate::Result<(Classification, Option<Response>)> {
132            let category = match self.calls.fetch_add(1, Ordering::Relaxed) {
133                0 => Some(Category::Capable),
134                1 => Some(Category::Efficient),
135                _ => None,
136            };
137            let scores = match category {
138                Some(category) => vec![Score {
139                    confidence: 1.0,
140                    target: driver.first_model_for(&category)?.clone(),
141                    category: Some(category),
142                }],
143                None => Vec::new(),
144            };
145            Ok((Classification::Scores(scores), None))
146        }
147    }
148
149    fn request(metadata: Option<Metadata>) -> Request {
150        Request {
151            llm_request: text_request(Some("auto".to_string()), "hi"),
152            raw_request: None,
153            metadata,
154        }
155    }
156
157    fn child(agent_id: &str) -> Request {
158        request(Some(Metadata {
159            session_id: Some("session-1".to_string()),
160            agent_id: Some(agent_id.to_string()),
161            is_subagent: true,
162            is_delegated_work: true,
163            ..Metadata::default()
164        }))
165    }
166
167    fn configured(classifier: Arc<dyn Classifier<State>>) -> crate::Result<Arc<SubagentRouter>> {
168        Ok(Arc::new(SubagentRouter::new(
169            Arc::new(Passthrough),
170            SubagentRouterConfig {
171                classifier,
172                default_target: Category::Capable,
173                classify_trigger: ClassifyTrigger::NewSession,
174                message_hash_fallback: false,
175            },
176        )?))
177    }
178
179    #[tokio::test]
180    async fn routes_parent_and_children_with_affinity_and_default() -> crate::Result<()> {
181        let classifier = Arc::new(ScriptedClassifier {
182            calls: AtomicUsize::new(0),
183        });
184        let router = configured(classifier.clone())?;
185
186        // The parent and its children route over separate model groups.
187        let models = RuntimeModels::new([(Category::Any, vec![ModelId::from("parent")])].into())
188            .with_subagent(
189                [
190                    (
191                        Category::Any,
192                        vec![ModelId::from("worker"), ModelId::from("reviewer")],
193                    ),
194                    (Category::Capable, vec![ModelId::from("worker")]),
195                    (Category::Efficient, vec![ModelId::from("reviewer")]),
196                ]
197                .into(),
198            );
199        let (selected_parent, _) =
200            test_drive_with_models(router.clone(), request(None), models.clone(), echo()).await?;
201        let (first, _) =
202            test_drive_with_models(router.clone(), child("child-1"), models.clone(), echo())
203                .await?;
204        let (same_child, _) =
205            test_drive_with_models(router.clone(), child("child-1"), models.clone(), echo())
206                .await?;
207        let (sibling, _) =
208            test_drive_with_models(router.clone(), child("child-2"), models.clone(), echo())
209                .await?;
210        let (defaulted, _) =
211            test_drive_with_models(router.clone(), child("child-3"), models.clone(), echo())
212                .await?;
213        let maintenance = request(Some(Metadata {
214            session_id: Some("session-1".to_string()),
215            agent_id: Some("child-1".to_string()),
216            is_subagent: true,
217            is_delegated_work: false,
218            ..Metadata::default()
219        }));
220        let (maintenance, _) =
221            test_drive_with_models(router, maintenance, models.clone(), echo()).await?;
222
223        assert_eq!(selected_parent, "parent");
224        assert_eq!(first, "worker");
225        assert_eq!(same_child, "worker");
226        assert_eq!(sibling, "reviewer");
227        assert_eq!(defaulted, "worker");
228        assert_eq!(maintenance, "parent");
229        assert_eq!(classifier.calls.load(Ordering::Relaxed), 3);
230
231        let fixed = Arc::new(SubagentRouter::new(
232            Arc::new(Passthrough),
233            SubagentRouterConfig::fixed_target(),
234        )?);
235        let (fixed, _) = test_drive_with_models(fixed, child("fixed"), models, echo()).await?;
236        assert_eq!(fixed, "worker");
237        Ok(())
238    }
239}