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 crate::core::state::State;
26use switchyard_protocol::{
27    ContentBlock, LlmRequest, Message, Metadata, ModelId, Request, Response, Role,
28};
29
30/// Classifies delegated work from its parent-supplied prompt and abstains otherwise.
31///
32/// The inner classifier receives a prompt-only request clone. The original request remains
33/// unchanged for the selected child model.
34pub struct SubagentGate {
35    inner: Arc<dyn Classifier>,
36}
37
38impl SubagentGate {
39    /// Wraps `inner` with delegated-work detection.
40    pub fn new(inner: Arc<dyn Classifier>) -> Self {
41        Self { inner }
42    }
43}
44
45/// Builds a classifier request containing only the delegated task text.
46fn delegated_prompt_request(request: &Request) -> Option<Request> {
47    // Agent clients put the delegated task after added context and reminders.
48    // A user message containing only tool-result blocks has no top-level text.
49    // Find the latest nonblank user text without using text inside tool results.
50    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
98/// Scores a fixed worker target for delegated sub-agent work; abstains otherwise.
99pub struct SubagentOverride {
100    /// Name of the worker target, resolved by the cascade against its target set.
101    worker: ModelId,
102}
103
104impl SubagentOverride {
105    /// Creates an override scoring `worker` for delegated sub-agent work.
106    ///
107    /// `worker` must name a target in the cascade's set, or routing a sub-agent request
108    /// fails with [`LibsyError::TargetNotFound`](crate::LibsyError::TargetNotFound).
109    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        // Delegated *work* only. A harness maintenance turn (e.g. Codex `compact`) carries
125        // sub-agent lineage but is not delegated work, so it abstains and routes normally.
126        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    /// Scores `headers` through the override, returning the winning target if it scored.
165    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        // Claude Code child-agent lineage.
182        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        // Codex delegated-work kinds.
189        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        // Confidence 1.0 under `Scores` (never `Ambiguous`), so the cascade stops here
213        // rather than consulting later classifiers.
214        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}