switchyard_libsy/algorithms/util/
subagent.rs1use async_trait::async_trait;
19
20use crate::Result;
21use crate::core::algorithm::Driver;
22use crate::core::classifier::{Classification, Classifier, Score};
23use switchyard_protocol::{Metadata, Request, Response};
24
25pub struct SubagentOverride {
27 worker: String,
29}
30
31impl SubagentOverride {
32 pub fn new(worker: impl Into<String>) -> Self {
37 Self {
38 worker: worker.into(),
39 }
40 }
41}
42
43#[async_trait]
44impl<S> Classifier<S> for SubagentOverride
45where
46 S: Send + 'static,
47{
48 async fn score(
49 &self,
50 _state: &mut S,
51 request: &mut Request,
52 _driver: Option<&Driver>,
53 ) -> Result<(Classification, Option<Response>)> {
54 let is_delegated_work = request
57 .metadata
58 .as_ref()
59 .is_some_and(Metadata::is_subagent_work);
60 Ok((
61 Classification::Scores(if is_delegated_work {
62 vec![Score {
63 confidence: 1.0,
64 target: self.worker.clone(),
65 }]
66 } else {
67 Vec::new()
68 }),
69 None,
70 ))
71 }
72}
73
74#[cfg(test)]
75mod tests {
76 use super::*;
77 use std::collections::BTreeMap;
78 use switchyard_protocol::text_request;
79
80 fn request(headers: &[(&str, &str)]) -> Request {
81 let metadata = (!headers.is_empty()).then(|| {
82 Metadata::from_headers(
83 &headers
84 .iter()
85 .map(|(name, value)| ((*name).to_string(), (*value).to_string()))
86 .collect::<BTreeMap<_, _>>(),
87 )
88 });
89 Request {
90 llm_request: text_request(Some("auto".to_string()), "hi"),
91 raw_request: None,
92 metadata,
93 }
94 }
95
96 async fn selected(headers: &[(&str, &str)]) -> Result<Option<String>> {
98 let mut state = ();
99 let classification = SubagentOverride::new("worker")
100 .score(&mut state, &mut request(headers), None)
101 .await?;
102 Ok(classification.0.argmax(false)?.map(|score| score.target))
103 }
104
105 #[tokio::test]
106 async fn requests_without_metadata_abstain() -> Result<()> {
107 assert_eq!(selected(&[]).await?, None);
108 Ok(())
109 }
110
111 #[tokio::test]
112 async fn subagent_work_scores_the_worker() -> Result<()> {
113 let claude = &[
115 ("x-claude-code-session-id", "root"),
116 ("x-claude-code-agent-id", "child-1"),
117 ];
118 assert_eq!(selected(claude).await?, Some("worker".to_string()));
119
120 assert_eq!(
122 selected(&[("x-openai-subagent", "review")]).await?,
123 Some("worker".to_string())
124 );
125 assert_eq!(
126 selected(&[("x-openai-subagent", "collab_spawn")]).await?,
127 Some("worker".to_string())
128 );
129 Ok(())
130 }
131
132 #[tokio::test]
133 async fn harness_maintenance_turns_abstain() -> Result<()> {
134 assert_eq!(selected(&[("x-openai-subagent", "compact")]).await?, None);
135 assert_eq!(
136 selected(&[("x-switchyard-is-subagent", "false")]).await?,
137 None
138 );
139 Ok(())
140 }
141
142 #[tokio::test]
143 async fn delegated_work_is_scored_definitively() -> Result<()> {
144 let mut state = ();
147 let classification = SubagentOverride::new("worker")
148 .score(
149 &mut state,
150 &mut request(&[("x-openai-subagent", "review")]),
151 None,
152 )
153 .await?;
154 match classification.0 {
155 Classification::Scores(scores) => {
156 assert_eq!(scores.len(), 1);
157 assert_eq!(scores[0].confidence, 1.0);
158 }
159 Classification::Ambiguous(_) => panic!("override must score definitively"),
160 }
161 Ok(())
162 }
163}