Skip to main content

switchyard_libsy/algorithms/util/
subagent.rs

1// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Sub-agent override as a single SDK component.
5//!
6//! [`SubagentOverride`] scores one fixed worker target for requests carrying delegated
7//! sub-agent work ([`Metadata::is_subagent_work`]) and abstains for everything else, so a
8//! cascade falls through to its later classifiers on ordinary traffic.
9//!
10//! It is stateless and holds only the worker's *name*: the fall-through cascade resolves it
11//! against its
12//! target set. Keeping the policy independent of any memory of past decisions is what lets
13//! it compose with a stateful classifier such as
14//! [`AffinityRouter`](crate::algorithms::AffinityRouter) — the override decides *which*
15//! target delegated work belongs on, affinity decides *how long* a decision lives, and
16//! neither needs to know about the other.
17
18use 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
29/// Classifies delegated work from its parent-supplied prompt and abstains otherwise.
30///
31/// The inner classifier receives a prompt-only request clone. The original request remains
32/// unchanged for the selected child model.
33pub struct SubagentGate<S> {
34    inner: Arc<dyn Classifier<S>>,
35}
36
37impl<S> SubagentGate<S> {
38    /// Wraps `inner` with delegated-work detection.
39    pub fn new(inner: Arc<dyn Classifier<S>>) -> Self {
40        Self { inner }
41    }
42}
43
44/// Builds a classifier request containing only the delegated task text.
45fn delegated_prompt_request(request: &Request) -> Option<Request> {
46    // Agent clients put the delegated task after added context and reminders.
47    // A user message containing only tool-result blocks has no top-level text.
48    // Find the latest nonblank user text without using text inside tool results.
49    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
100/// Scores a fixed worker target for delegated sub-agent work; abstains otherwise.
101pub struct SubagentOverride {
102    /// Name of the worker target, resolved by the cascade against its target set.
103    worker: ModelId,
104}
105
106impl SubagentOverride {
107    /// Creates an override scoring `worker` for delegated sub-agent work.
108    ///
109    /// `worker` must name a target in the cascade's set, or routing a sub-agent request
110    /// fails with [`LibsyError::TargetNotFound`](crate::LibsyError::TargetNotFound).
111    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        // Delegated *work* only. A harness maintenance turn (e.g. Codex `compact`) carries
130        // sub-agent lineage but is not delegated work, so it abstains and routes normally.
131        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    /// Scores `headers` through the override, returning the winning target if it scored.
170    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        // Claude Code child-agent lineage.
187        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        // Codex delegated-work kinds.
194        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        // Confidence 1.0 under `Scores` (never `Ambiguous`), so the cascade stops here
218        // rather than consulting later classifiers.
219        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}