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
50 .llm_request
51 .messages
52 .iter()
53 .rev()
54 .filter(|message| message.role == Role::User)
55 .find_map(|message| {
56 message.content.iter().rev().find_map(|block| match block {
57 ContentBlock::Text { text } if !text.trim().is_empty() => Some(text.clone()),
58 _ => None,
59 })
60 })?;
61
62 Some(Request {
63 llm_request: LlmRequest {
64 model: request.llm_request.model.clone(),
65 messages: vec![Message::text(Role::User, prompt)],
66 ..LlmRequest::default()
67 },
68 raw_request: None,
69 metadata: request.metadata.clone(),
70 })
71}
72
73#[async_trait]
74impl<S> Classifier<S> for SubagentGate<S>
75where
76 S: Send + 'static,
77{
78 async fn score(
79 &self,
80 state: &mut S,
81 request: &mut Request,
82 driver: &Driver,
83 ) -> Result<(Classification, Option<Response>)> {
84 if !request
85 .metadata
86 .as_ref()
87 .is_some_and(Metadata::is_subagent_work)
88 {
89 return Ok((Classification::Scores(Vec::new()), None));
90 }
91 let Some(mut classifier_request) = delegated_prompt_request(request) else {
92 return Ok((Classification::Scores(Vec::new()), None));
93 };
94 self.inner
95 .score(state, &mut classifier_request, driver)
96 .await
97 }
98}
99
100pub struct SubagentOverride {
102 worker: ModelId,
104}
105
106impl SubagentOverride {
107 pub fn new(worker: impl Into<ModelId>) -> Self {
112 Self {
113 worker: worker.into(),
114 }
115 }
116}
117
118#[async_trait]
119impl<S> Classifier<S> for SubagentOverride
120where
121 S: Send + 'static,
122{
123 async fn score(
124 &self,
125 _state: &mut S,
126 request: &mut Request,
127 driver: &Driver,
128 ) -> Result<(Classification, Option<Response>)> {
129 let is_delegated_work = request
132 .metadata
133 .as_ref()
134 .is_some_and(Metadata::is_subagent_work);
135 if is_delegated_work {
136 driver.set_evidence(serde_json::json!({"source": "subagent"}));
137 }
138 Ok((
139 Classification::Scores(if is_delegated_work {
140 vec![Score {
141 confidence: 1.0,
142 target: self.worker.clone(),
143 category: None,
144 }]
145 } else {
146 Vec::new()
147 }),
148 None,
149 ))
150 }
151}
152
153#[cfg(test)]
154mod tests {
155 use super::*;
156 use crate::core::testing::empty_driver;
157 use switchyard_protocol::{slice_to_header_map, text_request};
158
159 fn request(headers: &[(&str, &str)]) -> Request {
160 let metadata =
161 (!headers.is_empty()).then(|| Metadata::from_headers(&slice_to_header_map(headers)));
162 Request {
163 llm_request: text_request(Some(ModelId::from("auto").to_string()), "hi"),
164 raw_request: None,
165 metadata,
166 }
167 }
168
169 async fn selected(headers: &[(&str, &str)]) -> Result<Option<ModelId>> {
171 let mut state = ();
172 let classification = SubagentOverride::new("worker")
173 .score(&mut state, &mut request(headers), &empty_driver())
174 .await?;
175 Ok(classification.0.argmax(false)?.map(|score| score.target))
176 }
177
178 #[tokio::test]
179 async fn requests_without_metadata_abstain() -> Result<()> {
180 assert_eq!(selected(&[]).await?, None);
181 Ok(())
182 }
183
184 #[tokio::test]
185 async fn subagent_work_scores_the_worker() -> Result<()> {
186 let claude = &[
188 ("x-claude-code-session-id", "root"),
189 ("x-claude-code-agent-id", "child-1"),
190 ];
191 assert_eq!(selected(claude).await?, Some(ModelId::from("worker")));
192
193 assert_eq!(
195 selected(&[("x-openai-subagent", "review")]).await?,
196 Some(ModelId::from("worker"))
197 );
198 assert_eq!(
199 selected(&[("x-openai-subagent", "collab_spawn")]).await?,
200 Some(ModelId::from("worker"))
201 );
202 Ok(())
203 }
204
205 #[tokio::test]
206 async fn harness_maintenance_turns_abstain() -> Result<()> {
207 assert_eq!(selected(&[("x-openai-subagent", "compact")]).await?, None);
208 assert_eq!(
209 selected(&[("x-switchyard-is-subagent", "false")]).await?,
210 None
211 );
212 Ok(())
213 }
214
215 #[tokio::test]
216 async fn delegated_work_is_scored_definitively() -> Result<()> {
217 let mut state = ();
220 let classification = SubagentOverride::new("worker")
221 .score(
222 &mut state,
223 &mut request(&[("x-openai-subagent", "review")]),
224 &empty_driver(),
225 )
226 .await?;
227 match classification.0 {
228 Classification::Scores(scores) => {
229 assert_eq!(scores.len(), 1);
230 assert_eq!(scores[0].confidence, 1.0);
231 }
232 Classification::Ambiguous(_) => panic!("override must score definitively"),
233 }
234 Ok(())
235 }
236}