switchyard_libsy/algorithms/util/
subagent.rs1use std::sync::Arc;
19
20use async_trait::async_trait;
21
22use crate::Result;
23use crate::core::algorithm::Driver;
24use crate::core::classifier::{Classification, Classifier, Score};
25use switchyard_protocol::{
26 ContentBlock, LlmRequest, Message, Metadata, ModelId, Request, Response, Role,
27};
28
29pub struct SubagentGate<S> {
34 inner: Arc<dyn Classifier<S>>,
35}
36
37impl<S> SubagentGate<S> {
38 pub fn new(inner: Arc<dyn Classifier<S>>) -> Self {
40 Self { inner }
41 }
42}
43
44fn delegated_prompt_request(request: &Request) -> Option<Request> {
46 let prompt = request
48 .llm_request
49 .messages
50 .iter()
51 .rev()
52 .find(|message| message.role == Role::User)?
53 .content
54 .iter()
55 .rev()
56 .find_map(|block| match block {
57 ContentBlock::Text { text } if !text.trim().is_empty() => Some(text.clone()),
58 _ => None,
59 })?;
60
61 Some(Request {
62 llm_request: LlmRequest {
63 model: request.llm_request.model.clone(),
64 messages: vec![Message::text(Role::User, prompt)],
65 ..LlmRequest::default()
66 },
67 raw_request: None,
68 metadata: request.metadata.clone(),
69 })
70}
71
72#[async_trait]
73impl<S> Classifier<S> for SubagentGate<S>
74where
75 S: Send + 'static,
76{
77 async fn score(
78 &self,
79 state: &mut S,
80 request: &mut Request,
81 driver: &Driver,
82 ) -> Result<(Classification, Option<Response>)> {
83 if !request
84 .metadata
85 .as_ref()
86 .is_some_and(Metadata::is_subagent_work)
87 {
88 return Ok((Classification::Scores(Vec::new()), None));
89 }
90 let Some(mut classifier_request) = delegated_prompt_request(request) else {
91 return Ok((Classification::Scores(Vec::new()), None));
92 };
93 self.inner
94 .score(state, &mut classifier_request, driver)
95 .await
96 }
97}
98
99pub struct SubagentOverride {
101 worker: ModelId,
103}
104
105impl SubagentOverride {
106 pub fn new(worker: impl Into<ModelId>) -> Self {
111 Self {
112 worker: worker.into(),
113 }
114 }
115}
116
117#[async_trait]
118impl<S> Classifier<S> for SubagentOverride
119where
120 S: Send + 'static,
121{
122 async fn score(
123 &self,
124 _state: &mut S,
125 request: &mut Request,
126 driver: &Driver,
127 ) -> Result<(Classification, Option<Response>)> {
128 let is_delegated_work = request
131 .metadata
132 .as_ref()
133 .is_some_and(Metadata::is_subagent_work);
134 if is_delegated_work {
135 driver.set_evidence(serde_json::json!({"source": "subagent"}));
136 }
137 Ok((
138 Classification::Scores(if is_delegated_work {
139 vec![Score {
140 confidence: 1.0,
141 target: self.worker.clone(),
142 category: None,
143 }]
144 } else {
145 Vec::new()
146 }),
147 None,
148 ))
149 }
150}
151
152#[cfg(test)]
153mod tests {
154 use super::*;
155 use crate::core::testing::empty_driver;
156 use parking_lot::Mutex;
157 use switchyard_protocol::{slice_to_header_map, text_request};
158
159 #[derive(Default)]
160 struct CapturingClassifier {
161 requests: Mutex<Vec<Request>>,
162 }
163
164 #[async_trait]
165 impl Classifier<()> for CapturingClassifier {
166 async fn score(
167 &self,
168 _state: &mut (),
169 request: &mut Request,
170 _driver: &Driver,
171 ) -> Result<(Classification, Option<Response>)> {
172 self.requests.lock().push(request.clone());
173 Ok((
174 Classification::Scores(vec![Score {
175 confidence: 1.0,
176 target: ModelId::from("worker"),
177 category: None,
178 }]),
179 None,
180 ))
181 }
182 }
183
184 fn request(headers: &[(&str, &str)]) -> Request {
185 let metadata =
186 (!headers.is_empty()).then(|| Metadata::from_headers(&slice_to_header_map(headers)));
187 Request {
188 llm_request: text_request(Some(ModelId::from("auto").to_string()), "hi"),
189 raw_request: None,
190 metadata,
191 }
192 }
193
194 async fn selected(headers: &[(&str, &str)]) -> Result<Option<ModelId>> {
196 let mut state = ();
197 let classification = SubagentOverride::new("worker")
198 .score(&mut state, &mut request(headers), &empty_driver())
199 .await?;
200 Ok(classification.0.argmax(false)?.map(|score| score.target))
201 }
202
203 #[tokio::test]
204 async fn requests_without_metadata_abstain() -> Result<()> {
205 assert_eq!(selected(&[]).await?, None);
206 Ok(())
207 }
208
209 #[tokio::test]
210 async fn subagent_work_scores_the_worker() -> Result<()> {
211 let claude = &[
213 ("x-claude-code-session-id", "root"),
214 ("x-claude-code-agent-id", "child-1"),
215 ];
216 assert_eq!(selected(claude).await?, Some(ModelId::from("worker")));
217
218 assert_eq!(
220 selected(&[("x-openai-subagent", "review")]).await?,
221 Some(ModelId::from("worker"))
222 );
223 assert_eq!(
224 selected(&[("x-openai-subagent", "collab_spawn")]).await?,
225 Some(ModelId::from("worker"))
226 );
227 Ok(())
228 }
229
230 #[tokio::test]
231 async fn harness_maintenance_turns_abstain() -> Result<()> {
232 assert_eq!(selected(&[("x-openai-subagent", "compact")]).await?, None);
233 assert_eq!(
234 selected(&[("x-switchyard-is-subagent", "false")]).await?,
235 None
236 );
237 Ok(())
238 }
239
240 #[tokio::test]
241 async fn delegated_work_is_scored_definitively() -> Result<()> {
242 let mut state = ();
245 let classification = SubagentOverride::new("worker")
246 .score(
247 &mut state,
248 &mut request(&[("x-openai-subagent", "review")]),
249 &empty_driver(),
250 )
251 .await?;
252 match classification.0 {
253 Classification::Scores(scores) => {
254 assert_eq!(scores.len(), 1);
255 assert_eq!(scores[0].confidence, 1.0);
256 }
257 Classification::Ambiguous(_) => panic!("override must score definitively"),
258 }
259 Ok(())
260 }
261
262 #[tokio::test]
263 async fn gate_abstains_when_delegated_work_has_no_text_prompt() -> Result<()> {
264 let classifier = Arc::new(CapturingClassifier::default());
265 let gate = SubagentGate::new(classifier.clone());
266 let mut request = request(&[("x-openai-subagent", "collab_spawn")]);
267 request.llm_request.messages = vec![Message::text(Role::Assistant, "no user prompt")];
268
269 let mut state = ();
270 let (classification, response) = gate
271 .score(&mut state, &mut request, &empty_driver())
272 .await?;
273
274 assert!(classification.argmax(false)?.is_none());
275 assert!(response.is_none());
276 assert!(classifier.requests.lock().is_empty());
277 Ok(())
278 }
279}