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