From 981aa662bdf19dce5d74000a09015a257ce2e00c Mon Sep 17 00:00:00 2001 From: ayushag Date: Wed, 19 Aug 2026 23:04:55 -0700 Subject: [PATCH 1/4] feat: classifier decision gate for sub-agent routing Signed-off-by: ayushag --- crates/libsy/src/algorithms/passthrough.rs | 349 ++++++++++++++++-- .../src/algorithms/subagent_affinity_tests.rs | 20 + crates/libsy/src/algorithms/util/affinity.rs | 17 +- crates/libsy/src/algorithms/util/subagent.rs | 174 ++++++++- crates/libsy/src/lib.rs | 4 +- crates/switchyard-server/src/config.rs | 160 +++++++- 6 files changed, 685 insertions(+), 39 deletions(-) diff --git a/crates/libsy/src/algorithms/passthrough.rs b/crates/libsy/src/algorithms/passthrough.rs index 0ea325bf..26721e4e 100644 --- a/crates/libsy/src/algorithms/passthrough.rs +++ b/crates/libsy/src/algorithms/passthrough.rs @@ -1,26 +1,103 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! Single-target routing for direct model calls and integration diagnostics. +//! Direct parent routing with an optional delegated-work classifier cascade. use std::sync::Arc; use switchyard_protocol::{ModelId, Request}; -use crate::core::algorithm::{Algorithm, Driver}; -use crate::{Result, RoutingOutcome}; +use super::fall_through::{DefaultTarget, FallThrough}; +use super::util::affinity::AffinityRouter; +use super::util::subagent::{SubagentGate, SubagentOverride}; +use super::util::turn_pin::ClassifyTrigger; +use crate::core::algorithm::{self, Algorithm, Driver}; +use crate::core::classifier::Classifier; +use crate::core::state::State; +use crate::{LibsyError, Result, RoutingOutcome}; -/// Routing algorithm that always selects one configured target. +/// Routes parent traffic directly and optionally classifies delegated sub-agent work. pub struct Passthrough { - target: ModelId, + parent_target: ModelId, + route: FallThrough, +} + +/// Runtime components for classifying and retaining delegated sub-agent work. +pub struct PassthroughSubagentConfig { + /// Targets the delegated-work classifier may select. + pub targets: Vec, + /// Classifier invoked for the first request from each identified child. + pub classifier: Arc>, + /// Child target used when `classifier` abstains. + pub default_target: ModelId, + /// Controls whether each child is classified once or on every request. + pub classify_trigger: ClassifyTrigger, + /// Unsupported for child routing because child identity must come from harness metadata. + pub message_hash_fallback: bool, +} + +/// Complete construction settings for [`Passthrough`]. +pub struct PassthroughConfig { + /// Target used for parent and harness-maintenance traffic. + pub parent_target: ModelId, + /// Optional delegated-work decision gate. + pub subagent: Option, } impl Passthrough { - /// Creates an algorithm that always selects `target`. - pub fn new(target: impl Into) -> Self { - Passthrough { - target: target.into(), - } + /// Creates direct parent routing, optionally with a decision gate for sub-agents. + /// + /// When configured, the sub-agent classifier runs once per identified child. Its first + /// decision is retained by `session + agent`; an abstaining classifier uses the child + /// default. Root and harness-maintenance traffic continue to the parent target. + /// + /// # Errors + /// + /// Returns an error when the configured child default is not a child target. + pub fn new(config: PassthroughConfig) -> Result { + let parent_target = config.parent_target; + let route = match config.subagent { + None => FallThrough::new_with_state(vec![parent_target.clone()]) + .with_name("passthrough") + .with_classifier(Arc::new(DefaultTarget::new(parent_target.clone()))), + Some(subagent) => { + algorithm::ensure_model_is_target(&subagent.targets, &subagent.default_target)?; + if subagent.message_hash_fallback { + return Err(LibsyError::AlgorithmError { + message: "sub-agent routing cannot use message_hash_fallback".to_string(), + }); + } + let mut targets = subagent.targets; + if !targets.contains(&parent_target) { + targets.push(parent_target.clone()); + } + let mut route = FallThrough::new_with_state(targets).with_name("passthrough"); + match subagent.classify_trigger { + ClassifyTrigger::EveryRequest => {} + ClassifyTrigger::NewSession => { + let affinity = Arc::new(AffinityRouter::for_subagents()); + route = route + .with_processor(affinity.clone()) + .with_classifier(affinity); + } + ClassifyTrigger::UserTurn => { + return Err(LibsyError::AlgorithmError { + message: "sub-agent routing cannot use classify_trigger = user_turn" + .to_string(), + }); + } + } + route + .with_classifier(Arc::new(SubagentGate::new(subagent.classifier))) + .with_classifier(Arc::new(SubagentOverride::new(subagent.default_target))) + .with_classifier(Arc::new(DefaultTarget::new(parent_target.clone()))) + } + }; + + Ok(Self { + parent_target, + route, + }) } } @@ -30,24 +107,114 @@ impl Algorithm for Passthrough { "passthrough" } - async fn route(self: Arc, _driver: Driver, request: Request) -> Result { - tracing::info!(target = %self.target, "passthrough selected target"); - Ok(RoutingOutcome::route_to( - self.target.clone(), - Vec::new(), - request, - )) + async fn route(self: Arc, driver: Driver, request: Request) -> Result { + let mut outcome = self.route.execute(driver, request).await?; + // Parent traffic preserves passthrough's no-fallback contract. Child traffic may + // fall back only within the child target set, never into the parent route. + if outcome.selected_model_id == self.parent_target { + outcome.fallback_models.clear(); + } else { + outcome + .fallback_models + .retain(|target| *target != self.parent_target); + } + Ok(outcome) } } #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; - use super::Passthrough; + use async_trait::async_trait; + use parking_lot::Mutex; + use serde_json::json; + + use super::{Passthrough, PassthroughConfig, PassthroughSubagentConfig}; use crate::core::algorithm::Algorithm; - use crate::core::testing::{echo, test_drive}; - use switchyard_protocol::{Request, completion_text, text_request}; + use crate::core::classifier::{Classification, Classifier, Score}; + use crate::core::testing::{echo, reply, test_drive}; + use crate::{ + ClassifyTrigger, CustomClassifierConfig, CustomClassifierPolicy, Driver, + LlmClassifierConfig, LlmTaskClassifier, State, + }; + use switchyard_protocol::{ + ContentBlock, InstructionBlock, Message, Metadata, ModelId, Request, Response, Role, + completion_text, text_request, + }; + + struct AlternatingClassifier { + calls: AtomicUsize, + } + + #[async_trait] + impl Classifier for AlternatingClassifier { + async fn score( + &self, + _state: &mut State, + _request: &mut Request, + _driver: Option<&Driver>, + ) -> crate::Result<(Classification, Option)> { + let target = if self.calls.fetch_add(1, Ordering::Relaxed) == 0 { + "worker" + } else { + "reviewer" + }; + Ok(( + Classification::Scores(vec![Score { + confidence: 1.0, + target: ModelId::from(target), + }]), + None, + )) + } + } + + struct AbstainingClassifier; + + #[async_trait] + impl Classifier for AbstainingClassifier { + async fn score( + &self, + _state: &mut State, + _request: &mut Request, + _driver: Option<&Driver>, + ) -> crate::Result<(Classification, Option)> { + Ok((Classification::Scores(Vec::new()), None)) + } + } + + fn request(metadata: Option) -> Request { + Request { + llm_request: text_request(Some("auto".to_string()), "hi"), + raw_request: None, + metadata, + } + } + + fn child(agent_id: &str) -> Request { + request(Some(Metadata { + session_id: Some("session-1".to_string()), + agent_id: Some(agent_id.to_string()), + is_subagent: true, + is_delegated_work: true, + ..Metadata::default() + })) + } + + fn configured(classifier: Arc>) -> crate::Result> { + Ok(Arc::new(Passthrough::new(PassthroughConfig { + parent_target: ModelId::from("parent"), + subagent: Some(PassthroughSubagentConfig { + targets: vec![ModelId::from("worker"), ModelId::from("reviewer")], + classifier, + default_target: ModelId::from("worker"), + classify_trigger: ClassifyTrigger::NewSession, + message_hash_fallback: false, + }), + })?)) + } #[tokio::test] async fn test_passthrough() -> crate::Result<()> { @@ -57,7 +224,10 @@ mod tests { raw_request: None, metadata: None, }; - let algorithm: Arc = Arc::new(Passthrough::new(MODEL_ID)); + let algorithm: Arc = Arc::new(Passthrough::new(PassthroughConfig { + parent_target: ModelId::from(MODEL_ID), + subagent: None, + })?); let (selected_model, response) = test_drive(algorithm, request, echo()).await?; assert_eq!( @@ -71,4 +241,141 @@ mod tests { assert_eq!(selected_model, MODEL_ID); Ok(()) } + + #[tokio::test] + async fn parent_traffic_keeps_the_passthrough_target() -> crate::Result<()> { + let router = configured(Arc::new(AlternatingClassifier { + calls: AtomicUsize::new(0), + }))?; + + let (selected, _) = test_drive(router, request(None), echo()).await?; + + assert_eq!(selected, "parent"); + Ok(()) + } + + #[tokio::test] + async fn classifier_decision_is_retained_per_child() -> crate::Result<()> { + let classifier = Arc::new(AlternatingClassifier { + calls: AtomicUsize::new(0), + }); + let router = configured(classifier.clone())?; + + let (first, _) = test_drive(router.clone(), child("child-1"), echo()).await?; + let (same_child, _) = test_drive(router.clone(), child("child-1"), echo()).await?; + let (sibling, _) = test_drive(router, child("child-2"), echo()).await?; + + assert_eq!(first, "worker"); + assert_eq!(same_child, "worker"); + assert_eq!(sibling, "reviewer"); + assert_eq!(classifier.calls.load(Ordering::Relaxed), 2); + Ok(()) + } + + #[tokio::test] + async fn an_abstaining_child_classifier_uses_the_child_default() -> crate::Result<()> { + let router = configured(Arc::new(AbstainingClassifier))?; + + let (selected, _) = test_drive(router, child("child-1"), echo()).await?; + + assert_eq!(selected, "worker"); + Ok(()) + } + + #[tokio::test] + async fn maintenance_traffic_ignores_an_existing_child_pin() -> crate::Result<()> { + let router = configured(Arc::new(AlternatingClassifier { + calls: AtomicUsize::new(0), + }))?; + test_drive(router.clone(), child("child-1"), echo()).await?; + let maintenance = request(Some(Metadata { + session_id: Some("session-1".to_string()), + agent_id: Some("child-1".to_string()), + is_subagent: true, + is_delegated_work: false, + ..Metadata::default() + })); + + let (selected, _) = test_drive(router, maintenance, echo()).await?; + + assert_eq!(selected, "parent"); + Ok(()) + } + + #[tokio::test] + async fn custom_classifier_receives_only_the_delegated_prompt() -> crate::Result<()> { + let classifier = LlmTaskClassifier::new(LlmClassifierConfig::Custom { + judge_target: ModelId::from("judge"), + targets: vec![ + ("worker".to_string(), ModelId::from("worker")), + ("reviewer".to_string(), ModelId::from("reviewer")), + ], + default_target: "worker".to_string(), + config: CustomClassifierConfig::new( + "classify the delegated task", + json!({ + "type": "object", + "properties": { + "target": {"type": "string", "enum": ["worker", "reviewer"]} + }, + "required": ["target"], + "additionalProperties": false + }), + CustomClassifierPolicy::target_selector("/target"), + ), + })?; + let router = configured(Arc::new(classifier))?; + let mut request = child("child-1"); + request.llm_request.instructions = vec![InstructionBlock { + role: Role::System, + content: Message::text(Role::System, "child system instructions").content, + }]; + request.llm_request.messages = vec![ + Message::text(Role::User, "harness context"), + Message { + role: Role::User, + content: vec![ + ContentBlock::Text { + text: "tool context".to_string(), + }, + ContentBlock::Text { + text: "review this parser".to_string(), + }, + ], + }, + ]; + let calls = Arc::new(Mutex::new(Vec::new())); + let served_calls = calls.clone(); + + let (selected, _) = test_drive(router, request, move |target, request| { + let calls = served_calls.clone(); + async move { + let completion = if target == "judge" { + r#"{"target":"reviewer"}"# + } else { + "child answer" + }; + calls.lock().push((target, request)); + Ok(reply(completion)) + } + }) + .await?; + + assert_eq!(selected, "reviewer"); + let calls = calls.lock(); + assert_eq!(calls.len(), 2); + assert_eq!(calls[0].0, "judge"); + assert_eq!( + calls[0].1.llm_request.instructions[0].content, + Message::text(Role::System, "classify the delegated task").content + ); + assert_eq!( + calls[0].1.llm_request.messages, + vec![Message::text(Role::User, "review this parser")] + ); + assert_eq!(calls[1].0, "reviewer"); + assert_eq!(calls[1].1.llm_request.instructions.len(), 1); + assert_eq!(calls[1].1.llm_request.messages.len(), 2); + Ok(()) + } } diff --git a/crates/libsy/src/algorithms/subagent_affinity_tests.rs b/crates/libsy/src/algorithms/subagent_affinity_tests.rs index 4a390351..92e7150c 100644 --- a/crates/libsy/src/algorithms/subagent_affinity_tests.rs +++ b/crates/libsy/src/algorithms/subagent_affinity_tests.rs @@ -137,6 +137,26 @@ async fn harness_maintenance_turns_are_not_forced_to_the_worker() -> Result<()> Ok(()) } +#[tokio::test] +async fn an_existing_pin_does_not_capture_harness_maintenance() -> Result<()> { + let router = router(); + + assert_eq!(turn(&router, &child("child-1")).await?, "worker"); + let served = turn( + &router, + &[ + ("x-switchyard-session-id", "session-1"), + ("x-switchyard-agent-id", "child-1"), + ("x-switchyard-is-subagent", "true"), + ("x-openai-subagent", "compact"), + ], + ) + .await?; + + assert_eq!(served, "orchestrator"); + Ok(()) +} + /// Builds a cascade whose override scores `worker`, sharing `affinity` across instances. fn router_overriding_to(affinity: Arc, worker: &str) -> Arc { Arc::new( diff --git a/crates/libsy/src/algorithms/util/affinity.rs b/crates/libsy/src/algorithms/util/affinity.rs index 5efe21d7..4f13e96f 100644 --- a/crates/libsy/src/algorithms/util/affinity.rs +++ b/crates/libsy/src/algorithms/util/affinity.rs @@ -14,8 +14,8 @@ //! //! Identity is derived from correlation metadata: by default, a request is keyed by its //! session, and a sub-agent is keyed more finely by `session + agent` — a subset of its -//! session. [`AffinityRouter::for_subagents`] narrows affinity to explicitly identified -//! child agents, leaving root traffic to later classifiers on every turn. +//! session. [`AffinityRouter::for_subagents`] narrows affinity to delegated child-agent +//! work, leaving root and harness-maintenance traffic to later classifiers on every turn. use std::collections::{HashMap, HashSet, hash_map::DefaultHasher}; use std::hash::{Hash, Hasher}; @@ -47,7 +47,7 @@ const MAX_ASSIGNMENTS: usize = 4096; pub struct AffinityRouter { /// When set, only these models are retained; a decision for any other model is not latched. latch_only: Option>, - /// Whether root-session requests should abstain instead of being retained. + /// Whether requests other than delegated sub-agent work should abstain. subagents_only: bool, /// In absence of headers, use the message hash based fallback key to do task based routing message_hash_fallback: bool, @@ -66,9 +66,10 @@ impl AffinityRouter { Self::default() } - /// Creates a router that retains assignments only for explicitly identified child agents. + /// Creates a router that retains assignments only for delegated child-agent work. /// - /// Root-agent requests always abstain, so a later classifier selects them on every turn. + /// Root and harness-maintenance requests always abstain, so a later classifier selects + /// them on every turn. pub fn for_subagents() -> Self { Self { subagents_only: true, @@ -100,8 +101,9 @@ impl AffinityRouter { fn affinity_key(&self, request: &Request) -> Option { let metadata = request.metadata.as_ref(); let is_subagent = metadata.is_some_and(|metadata| metadata.is_subagent); - // This mode handles only subagent requests; root requests intentionally fall through. - if self.subagents_only && !is_subagent { + let is_subagent_work = metadata.is_some_and(|metadata| metadata.is_subagent_work()); + // This mode handles only delegated work; root and maintenance requests fall through. + if self.subagents_only && !is_subagent_work { return None; } @@ -274,6 +276,7 @@ mod tests { agent_id: Some(agent_id.to_string()), task_id: Some(task_id.to_string()), is_subagent: true, + is_delegated_work: true, ..Metadata::default() } } diff --git a/crates/libsy/src/algorithms/util/subagent.rs b/crates/libsy/src/algorithms/util/subagent.rs index c2077ad1..d52e094c 100644 --- a/crates/libsy/src/algorithms/util/subagent.rs +++ b/crates/libsy/src/algorithms/util/subagent.rs @@ -15,13 +15,90 @@ //! target delegated work belongs on, affinity decides *how long* a decision lives, and //! neither needs to know about the other. +use std::sync::Arc; + use async_trait::async_trait; use crate::Result; use crate::core::algorithm::Driver; use crate::core::classifier::{Classification, Classifier, Score}; -use switchyard_protocol::ModelId; -use switchyard_protocol::{Metadata, Request, Response}; +use switchyard_protocol::{ + ContentBlock, LlmRequest, Message, Metadata, ModelId, Request, Response, Role, +}; + +/// Classifies delegated work from its parent-supplied prompt and abstains otherwise. +/// +/// The inner classifier receives a prompt-only request clone. The original request remains +/// unchanged for the selected child model. +pub struct SubagentGate { + inner: Arc>, +} + +impl SubagentGate { + /// Wraps `inner` with delegated-work detection. + pub fn new(inner: Arc>) -> Self { + Self { inner } + } +} + +/// Builds the prompt-only request shown to a delegated-work classifier. +fn delegated_prompt_request(request: &Request) -> Option { + // Coding harnesses append the parent's task after their injected user context and reminders. + let prompt = request + .llm_request + .messages + .iter() + .rev() + .find(|message| message.role == Role::User)? + .content + .iter() + .rev() + .find_map(|block| match block { + ContentBlock::Text { text } if !text.trim().is_empty() => Some(text.clone()), + _ => None, + })?; + + Some(Request { + llm_request: LlmRequest { + model: request.llm_request.model.clone(), + messages: vec![Message::text(Role::User, prompt)], + ..LlmRequest::default() + }, + raw_request: None, + metadata: request.metadata.clone(), + }) +} + +#[async_trait] +impl Classifier for SubagentGate +where + S: Send + 'static, +{ + fn routing_tier(&self, selected_model_id: &ModelId) -> Option<&'static str> { + self.inner.routing_tier(selected_model_id) + } + + async fn score( + &self, + state: &mut S, + request: &mut Request, + driver: Option<&Driver>, + ) -> Result<(Classification, Option)> { + if !request + .metadata + .as_ref() + .is_some_and(Metadata::is_subagent_work) + { + return Ok((Classification::Scores(Vec::new()), None)); + } + let Some(mut classifier_request) = delegated_prompt_request(request) else { + return Ok((Classification::Scores(Vec::new()), None)); + }; + self.inner + .score(state, &mut classifier_request, driver) + .await + } +} /// Scores a fixed worker target for delegated sub-agent work; abstains otherwise. pub struct SubagentOverride { @@ -74,9 +151,36 @@ where #[cfg(test)] mod tests { + use parking_lot::Mutex; + use switchyard_protocol::InstructionBlock; + use super::*; use switchyard_protocol::{slice_to_header_map, text_request}; + #[derive(Default)] + struct CapturingClassifier { + requests: Mutex>, + } + + #[async_trait] + impl Classifier<()> for CapturingClassifier { + async fn score( + &self, + _state: &mut (), + request: &mut Request, + _driver: Option<&Driver>, + ) -> Result<(Classification, Option)> { + self.requests.lock().push(request.clone()); + Ok(( + Classification::Scores(vec![Score { + confidence: 1.0, + target: ModelId::from("worker"), + }]), + None, + )) + } + } + fn request(headers: &[(&str, &str)]) -> Request { let metadata = (!headers.is_empty()).then(|| Metadata::from_headers(&slice_to_header_map(headers))); @@ -154,4 +258,70 @@ mod tests { } Ok(()) } + + #[tokio::test] + async fn gate_shows_the_classifier_only_the_parent_prompt() -> Result<()> { + let classifier = Arc::new(CapturingClassifier::default()); + let gate = SubagentGate::new(classifier.clone()); + let mut request = request(&[("x-openai-subagent", "collab_spawn")]); + request.llm_request.instructions = vec![InstructionBlock { + role: Role::System, + content: Message::text(Role::System, "child system instructions").content, + }]; + request.llm_request.messages = vec![ + Message::text(Role::User, "harness context"), + Message { + role: Role::User, + content: vec![ + ContentBlock::Text { + text: "tool context".to_string(), + }, + ContentBlock::Text { + text: "implement the delegated task".to_string(), + }, + ], + }, + ]; + let original = request.clone(); + + let mut state = (); + gate.score(&mut state, &mut request, None).await?; + + let captured = classifier.requests.lock(); + assert_eq!(captured.len(), 1); + assert!(captured[0].llm_request.instructions.is_empty()); + assert_eq!( + captured[0].llm_request.messages, + vec![Message::text(Role::User, "implement the delegated task")] + ); + assert_eq!(request.llm_request, original.llm_request); + assert_eq!(request.raw_request, original.raw_request); + assert_eq!( + request + .metadata + .as_ref() + .and_then(|metadata| metadata.agent_id.as_deref()), + original + .metadata + .as_ref() + .and_then(|metadata| metadata.agent_id.as_deref()) + ); + Ok(()) + } + + #[tokio::test] + async fn gate_abstains_when_delegated_work_has_no_text_prompt() -> Result<()> { + let classifier = Arc::new(CapturingClassifier::default()); + let gate = SubagentGate::new(classifier.clone()); + let mut request = request(&[("x-openai-subagent", "collab_spawn")]); + request.llm_request.messages = vec![Message::text(Role::Assistant, "no user prompt")]; + + let mut state = (); + let (classification, response) = gate.score(&mut state, &mut request, None).await?; + + assert!(classification.argmax(false)?.is_none()); + assert!(response.is_none()); + assert!(classifier.requests.lock().is_empty()); + Ok(()) + } } diff --git a/crates/libsy/src/lib.rs b/crates/libsy/src/lib.rs index 9abf46da..231072ba 100644 --- a/crates/libsy/src/lib.rs +++ b/crates/libsy/src/lib.rs @@ -20,7 +20,7 @@ pub use algorithms::llm_class::{ TaskClassifierConfig, }; pub use algorithms::noop::Noop; -pub use algorithms::passthrough::Passthrough; +pub use algorithms::passthrough::{Passthrough, PassthroughConfig, PassthroughSubagentConfig}; pub use algorithms::rand::{Random, RandomClassifier}; pub use algorithms::stage::{LlmFallback, StageRouter, StageRouterConfig}; pub use algorithms::util::affinity::AffinityRouter; @@ -29,7 +29,7 @@ pub use algorithms::util::classifier_contract::{ }; pub use algorithms::util::escalation::EscalationJudgeConfig; pub use algorithms::util::prompts::{SystemPromptProcessor, TargetPrompts, append_note}; -pub use algorithms::util::subagent::SubagentOverride; +pub use algorithms::util::subagent::{SubagentGate, SubagentOverride}; pub use algorithms::util::tool_signals::{DEFAULT_RECENT_WINDOW, ToolSignals}; pub use algorithms::util::turn_pin::ClassifyTrigger; diff --git a/crates/switchyard-server/src/config.rs b/crates/switchyard-server/src/config.rs index 22b20f28..27059a77 100644 --- a/crates/switchyard-server/src/config.rs +++ b/crates/switchyard-server/src/config.rs @@ -12,8 +12,8 @@ use libsy::{ AdvisorGate, AdvisorGateConfig, Algorithm, ClassifierContractConfig, ClassifierResponseFormat, ClassifyTrigger, CustomClassifierConfig, CustomClassifierPolicy, EscalationJudgeConfig, GateTrigger, HandoffNoteConfig, LlmClassifierConfig, LlmFallback, LlmTaskClassifier, Noop, - Passthrough, PickerMode, Random, StageRouter, StageRouterConfig, TargetPrompts, - TaskClassifierConfig, + Passthrough, PassthroughConfig, PassthroughSubagentConfig, PickerMode, Random, StageRouter, + StageRouterConfig, TargetPrompts, TaskClassifierConfig, }; use serde::Deserialize; use serde_json::Value; @@ -363,16 +363,22 @@ struct EscalationClassifierRouteConfig { judge: EscalationJudgeConfig, } -#[derive(Debug)] +#[derive(Clone, Debug, Deserialize)] +#[serde(deny_unknown_fields)] struct CustomClassifierRouteConfig { + classifier_target: String, targets: Vec, default_target: String, prompt: String, response_schema: String, policy: ClassifierPolicyConfig, + #[serde(default)] classify_trigger: ClassifyTrigger, + #[serde(default)] message_hash_fallback: bool, + #[serde(default)] recent_turn_window: Option, + #[serde(default = "default_classifier_max_output_tokens")] max_output_tokens: u64, } @@ -409,6 +415,8 @@ enum RouteConfig { #[serde(default)] reasoning: Option, target: String, + #[serde(default)] + subagent_classifier: Option, }, LlmClassifier { id: ModelId, @@ -585,7 +593,17 @@ impl RouteConfig { match self { Self::Noop { .. } => Vec::new(), Self::Random { targets, .. } => targets.iter().map(String::as_str).collect(), - Self::Passthrough { target, .. } => vec![target], + Self::Passthrough { + target, + subagent_classifier, + .. + } => { + let mut names = vec![target.as_str()]; + if let Some(classifier) = subagent_classifier { + names.extend(classifier.targets.iter().map(String::as_str)); + } + names + } Self::LlmClassifier { mode, strong_target, @@ -633,6 +651,10 @@ impl RouteConfig { Self::LlmClassifier { classifier_target, .. } => names.push(classifier_target), + Self::Passthrough { + subagent_classifier: Some(classifier), + .. + } => names.push(&classifier.classifier_target), Self::StageRouter { classifier: Some(classifier), .. @@ -691,6 +713,7 @@ impl RouteConfig { fn classifier_mode(&self, route_name: &str) -> ServerResult { let Self::LlmClassifier { + classifier_target, mode, strong_target, weak_target, @@ -820,6 +843,7 @@ impl RouteConfig { } Ok(LlmClassifierModeConfig::Custom( CustomClassifierRouteConfig { + classifier_target: classifier_target.clone(), targets: required_classifier_field(route_name, "targets", targets)?, default_target: required_classifier_field( route_name, @@ -963,9 +987,81 @@ fn build_algorithm( .map_err(|error| ServerError::new(format!("random route {route_name}: {error}")))?; Ok(Arc::new(algorithm)) } - RouteConfig::Passthrough { target, .. } => { - let target = resolve_target_model_id(route_name, target, targets)?; - Ok(Arc::new(Passthrough::new(target))) + RouteConfig::Passthrough { + target, + subagent_classifier, + .. + } => { + let parent_target = resolve_target_model_id(route_name, target, targets)?; + let subagent = if let Some(config) = subagent_classifier { + let judge_target = + resolve_target_model_id(route_name, &config.classifier_target, targets)?; + let resolved_targets = config + .targets + .iter() + .map(|name| { + resolve_target_model_id(route_name, name, targets) + .map(|target| (name.clone(), target)) + }) + .collect::>>()?; + let subagent_default_target = resolved_targets + .iter() + .find(|(name, _)| *name == config.default_target) + .map(|(_, target)| target.clone()) + .ok_or_else(|| { + ServerError::new(format!( + "passthrough route {route_name}: subagent_classifier default_target {:?} must be one of its configured targets", + config.default_target + )) + })?; + let response_schema = + serde_json::from_str(&config.response_schema).map_err(|error| { + ServerError::new(format!( + "passthrough route {route_name}: subagent_classifier response_schema is invalid JSON: {error}" + )) + })?; + let mut classifier_config = CustomClassifierConfig::new( + config.prompt.clone(), + response_schema, + config.policy.clone().into_libsy(), + ); + classifier_config.recent_turn_window = config.recent_turn_window; + classifier_config.max_output_tokens = config.max_output_tokens; + let subagent_targets = resolved_targets + .iter() + .map(|(_, target)| target.clone()) + .collect(); + let classifier = Arc::new( + LlmTaskClassifier::new(LlmClassifierConfig::Custom { + judge_target, + targets: resolved_targets, + default_target: config.default_target.clone(), + config: classifier_config, + }) + .map_err(|error| { + ServerError::new(format!( + "passthrough route {route_name}: subagent_classifier: {error}" + )) + })?, + ); + Some(PassthroughSubagentConfig { + targets: subagent_targets, + classifier, + default_target: subagent_default_target, + classify_trigger: config.classify_trigger, + message_hash_fallback: config.message_hash_fallback, + }) + } else { + None + }; + let algorithm = Passthrough::new(PassthroughConfig { + parent_target, + subagent, + }) + .map_err(|error| { + ServerError::new(format!("passthrough route {route_name}: {error}")) + })?; + Ok(Arc::new(algorithm)) } RouteConfig::LlmClassifier { classifier_target, .. @@ -1299,6 +1395,56 @@ target = "weak" Ok(()) } + #[test] + fn passthrough_accepts_a_custom_subagent_classifier() -> ServerResult<()> { + let configured = VALID_CONFIG.replace( + "[routes.passthrough]\nid = \"switchyard/passthrough\"\ntype = \"passthrough\"\ntarget = \"weak\"", + r#"[routes.passthrough] +id = "switchyard/passthrough" +type = "passthrough" +target = "weak" + +[routes.passthrough.subagent_classifier] +classifier_target = "classifier" +targets = ["strong", "weak"] +default_target = "weak" +prompt = "Select a target for this delegated task." +response_schema = '{"type":"object","properties":{"target":{"type":"string","enum":["strong","weak"]}},"required":["target"],"additionalProperties":false}' +policy = { type = "target_selector", selector = "/target" } +classify_trigger = "new_session""#, + ); + + server_state_from_toml(&configured)?; + Ok(()) + } + + #[test] + fn passthrough_subagent_classifier_rejects_message_hash_fallback() { + let configured = VALID_CONFIG.replace( + "[routes.passthrough]\nid = \"switchyard/passthrough\"\ntype = \"passthrough\"\ntarget = \"weak\"", + r#"[routes.passthrough] +id = "switchyard/passthrough" +type = "passthrough" +target = "weak" + +[routes.passthrough.subagent_classifier] +classifier_target = "classifier" +targets = ["strong", "weak"] +default_target = "weak" +prompt = "Select a target for this delegated task." +response_schema = '{"type":"object","properties":{"target":{"type":"string","enum":["strong","weak"]}},"required":["target"],"additionalProperties":false}' +policy = { type = "target_selector", selector = "/target" } +classify_trigger = "new_session" +message_hash_fallback = true"#, + ); + + assert!( + error_message(&configured).contains("cannot use message_hash_fallback"), + "{}", + error_message(&configured) + ); + } + #[test] fn an_escalation_table_switches_the_classifier_route_to_escalation() -> ServerResult<()> { // Present: the classifier target judges the weak tier's reply each turn instead of From 91ef52076e7c31d0e2783030350aeb1b3045980a Mon Sep 17 00:00:00 2001 From: ayushag Date: Thu, 20 Aug 2026 00:11:43 -0700 Subject: [PATCH 2/4] chore: metadata fixes with latest codex version + docs Signed-off-by: ayushag --- .../src/algorithms/subagent_affinity_tests.rs | 20 ----- crates/libsy/src/algorithms/util/subagent.rs | 54 +----------- crates/protocol/src/metadata.rs | 32 ++++++- docs/reference/toml_schema.md | 4 +- docs/routing_algorithms/overview.md | 1 + docs/routing_algorithms/subagent_routing.md | 85 +++++++++++++++++++ mkdocs.yml | 1 + 7 files changed, 119 insertions(+), 78 deletions(-) create mode 100644 docs/routing_algorithms/subagent_routing.md diff --git a/crates/libsy/src/algorithms/subagent_affinity_tests.rs b/crates/libsy/src/algorithms/subagent_affinity_tests.rs index 92e7150c..4a390351 100644 --- a/crates/libsy/src/algorithms/subagent_affinity_tests.rs +++ b/crates/libsy/src/algorithms/subagent_affinity_tests.rs @@ -137,26 +137,6 @@ async fn harness_maintenance_turns_are_not_forced_to_the_worker() -> Result<()> Ok(()) } -#[tokio::test] -async fn an_existing_pin_does_not_capture_harness_maintenance() -> Result<()> { - let router = router(); - - assert_eq!(turn(&router, &child("child-1")).await?, "worker"); - let served = turn( - &router, - &[ - ("x-switchyard-session-id", "session-1"), - ("x-switchyard-agent-id", "child-1"), - ("x-switchyard-is-subagent", "true"), - ("x-openai-subagent", "compact"), - ], - ) - .await?; - - assert_eq!(served, "orchestrator"); - Ok(()) -} - /// Builds a cascade whose override scores `worker`, sharing `affinity` across instances. fn router_overriding_to(affinity: Arc, worker: &str) -> Arc { Arc::new( diff --git a/crates/libsy/src/algorithms/util/subagent.rs b/crates/libsy/src/algorithms/util/subagent.rs index d52e094c..06387de5 100644 --- a/crates/libsy/src/algorithms/util/subagent.rs +++ b/crates/libsy/src/algorithms/util/subagent.rs @@ -151,10 +151,8 @@ where #[cfg(test)] mod tests { - use parking_lot::Mutex; - use switchyard_protocol::InstructionBlock; - use super::*; + use parking_lot::Mutex; use switchyard_protocol::{slice_to_header_map, text_request}; #[derive(Default)] @@ -259,56 +257,6 @@ mod tests { Ok(()) } - #[tokio::test] - async fn gate_shows_the_classifier_only_the_parent_prompt() -> Result<()> { - let classifier = Arc::new(CapturingClassifier::default()); - let gate = SubagentGate::new(classifier.clone()); - let mut request = request(&[("x-openai-subagent", "collab_spawn")]); - request.llm_request.instructions = vec![InstructionBlock { - role: Role::System, - content: Message::text(Role::System, "child system instructions").content, - }]; - request.llm_request.messages = vec![ - Message::text(Role::User, "harness context"), - Message { - role: Role::User, - content: vec![ - ContentBlock::Text { - text: "tool context".to_string(), - }, - ContentBlock::Text { - text: "implement the delegated task".to_string(), - }, - ], - }, - ]; - let original = request.clone(); - - let mut state = (); - gate.score(&mut state, &mut request, None).await?; - - let captured = classifier.requests.lock(); - assert_eq!(captured.len(), 1); - assert!(captured[0].llm_request.instructions.is_empty()); - assert_eq!( - captured[0].llm_request.messages, - vec![Message::text(Role::User, "implement the delegated task")] - ); - assert_eq!(request.llm_request, original.llm_request); - assert_eq!(request.raw_request, original.raw_request); - assert_eq!( - request - .metadata - .as_ref() - .and_then(|metadata| metadata.agent_id.as_deref()), - original - .metadata - .as_ref() - .and_then(|metadata| metadata.agent_id.as_deref()) - ); - Ok(()) - } - #[tokio::test] async fn gate_abstains_when_delegated_work_has_no_text_prompt() -> Result<()> { let classifier = Arc::new(CapturingClassifier::default()); diff --git a/crates/protocol/src/metadata.rs b/crates/protocol/src/metadata.rs index 59f13122..9c62add8 100644 --- a/crates/protocol/src/metadata.rs +++ b/crates/protocol/src/metadata.rs @@ -16,6 +16,7 @@ const CODEX_SESSION_ID_PATH: &str = "x-codex-turn-metadata.session_id"; const CODEX_THREAD_ID_PATH: &str = "x-codex-turn-metadata.thread_id"; const CODEX_PARENT_THREAD_ID_PATH: &str = "x-codex-turn-metadata.parent_thread_id"; const CODEX_TURN_ID_PATH: &str = "x-codex-turn-metadata.turn_id"; +const CODEX_THREAD_SOURCE_PATH: &str = "x-codex-turn-metadata.thread_source"; const CODEX_SUBAGENT_KIND_PATH: &str = "x-codex-turn-metadata.subagent_kind"; const CODEX_AGENT_ROLE_PATH: &str = "x-codex-turn-metadata.agent_role"; const CODEX_TASK_ID_PATH: &str = "x-codex-turn-metadata.task_id"; @@ -232,9 +233,10 @@ impl Metadata { /// Returns `(parent_agent_id, is_subagent, is_delegated_work)` from the headers. /// /// Recognized sub-agent signals include `x-claude-code-agent-id`, -/// `x-openai-subagent`, `x-codex-turn-metadata.subagent_kind`, and explicit -/// `x-switchyard-is-subagent`. Other host correlation and parent-session headers -/// may populate metadata but do not drive sub-agent classification. +/// `x-openai-subagent`, Codex's `subagent_kind`, Codex's current child lineage +/// (`thread_source = subagent` plus `parent_thread_id`), and explicit +/// `x-switchyard-is-subagent`. Other host correlation and parent-session headers may +/// populate metadata but do not drive sub-agent classification. /// /// `is_delegated_work` is computed from raw harness signals, not from `agent_kind`, /// which may be set by an unrelated operator label (`x-switchyard-agent-kind`). @@ -253,7 +255,13 @@ fn parse_sub_agent(headers: &http::HeaderMap) -> (Option, bool, bool) { let parent = sy_header(headers, SWITCHYARD_PARENT_AGENT_ID_HEADER) .or_else(|| claude_parent.map(str::to_string)); - let is_subagent = explicit.unwrap_or(claude_subagent || harness_kind.is_some()); + // Current Codex releases identify spawned children through lineage rather than + // `subagent_kind`. Require both fields so a parent id used only for correlation + // cannot accidentally route an ordinary turn as delegated work. + let codex_child = parent.is_some() + && resolve_path(headers, CODEX_THREAD_SOURCE_PATH).as_deref() == Some("subagent"); + + let is_subagent = explicit.unwrap_or(claude_subagent || codex_child || harness_kind.is_some()); let is_delegated_work = match explicit { Some(false) => false, @@ -263,6 +271,7 @@ fn parse_sub_agent(headers: &http::HeaderMap) -> (Option, bool, bool) { .unwrap_or(true), None => { claude_subagent + || codex_child || harness_kind .as_deref() .is_some_and(|k| SUBAGENT_WORK_KINDS.contains(&k)) @@ -412,6 +421,21 @@ mod tests { assert_eq!(correlated.parent_agent_id.as_deref(), Some("root-thread")); assert!(!correlated.is_subagent); assert!(!correlated.is_subagent_work()); + + // Current Codex children add `thread_source = subagent` to the same + // lineage. Together the two fields are a delegated-work signal. + let child_body = serde_json::json!({ + "session_id": "root-session", + "thread_id": "child-thread", + "parent_thread_id": "root-thread", + "thread_source": "subagent", + "turn_id": "turn-4", + }) + .to_string(); + let child = metadata(&[(CODEX_TURN_METADATA_HEADER, child_body.as_str())]); + assert_eq!(child.parent_agent_id.as_deref(), Some("root-thread")); + assert!(child.is_subagent); + assert!(child.is_subagent_work()); } #[test] diff --git a/docs/reference/toml_schema.md b/docs/reference/toml_schema.md index 27079e37..b6e8cb4a 100644 --- a/docs/reference/toml_schema.md +++ b/docs/reference/toml_schema.md @@ -116,11 +116,13 @@ type = "noop" ### `passthrough` -Sends every request to one target. +Sends every request to one target. It can also classify delegated sub-agent work; +see [Sub-Agent-Aware Routing](../routing_algorithms/subagent_routing.md). | Key | Required | Meaning | |---|:---:|---| | `target` | Yes | Target every request is sent to. | +| `subagent_classifier` | No | Custom classifier used only for delegated sub-agent work. | ### `random` diff --git a/docs/routing_algorithms/overview.md b/docs/routing_algorithms/overview.md index dad214d8..f7b97198 100644 --- a/docs/routing_algorithms/overview.md +++ b/docs/routing_algorithms/overview.md @@ -12,6 +12,7 @@ configuration and tuning. For the vocabulary these pages use, see | Strategy | Use it when | Route `type` | |---|---|---| +| [Sub-Agent-Aware Routing](subagent_routing.md) | Parent traffic should use one target while delegated sub-agents are classified across other targets. | `passthrough` with `subagent_classifier` | | [Random Routing](random_routing.md) | You need a fixed traffic split for A/B tests, baselines, or cost experiments. | `random` | | [LLM Classifier Routing](llm_classifier_routing.md) | Request content should decide whether a turn needs the weak or strong tier. | `llm_classifier` | | [Stage-Router Routing](stage_router_routing.md) | Tool-result and agent-progress signals should route most turns without an extra classifier call. | `stage_router` | diff --git a/docs/routing_algorithms/subagent_routing.md b/docs/routing_algorithms/subagent_routing.md new file mode 100644 index 00000000..d4e6add7 --- /dev/null +++ b/docs/routing_algorithms/subagent_routing.md @@ -0,0 +1,85 @@ +# Sub-Agent-Aware Routing + +Sub-agent-aware routing keeps parent-agent traffic on one target while routing +delegated sub-agent work across separate targets. Current support is available +through the `passthrough` route's optional `subagent_classifier` table. + +```toml +schema_version = 1 + +[llm_clients.openrouter] +format = "openai_chat" +base_url = "https://openrouter.ai/api/v1" +api_key_env = "OPENROUTER_API_KEY" + +[targets.parent] +id = "anthropic/claude-sonnet-5" +llm_client = "openrouter" + +[targets.classifier] +id = "openai/gpt-4.1-mini" +llm_client = "openrouter" + +[targets.worker] +id = "openai/gpt-5.4-mini" +llm_client = "openrouter" + +[targets.reviewer] +id = "anthropic/claude-opus-5" +llm_client = "openrouter" + +[routes.agent] +id = "agent" +type = "passthrough" +target = "parent" +context_window = 400000 +tool_calling = true +reasoning = true + +[routes.agent.subagent_classifier] +classifier_target = "classifier" +targets = ["worker", "reviewer"] +default_target = "worker" +classify_trigger = "new_session" +max_output_tokens = 64 +prompt = """ +Select exactly one target for the delegated task. + +- Select "reviewer" for code review, critique, auditing, or correctness analysis. +- Select "worker" for implementation, research, explanation, and other delegated work. + +Return only JSON matching the response schema. +""" +response_schema = ''' +{ + "type": "object", + "properties": { + "target": {"type": "string", "enum": ["worker", "reviewer"]} + }, + "required": ["target"], + "additionalProperties": false +} +''' +policy = { type = "target_selector", selector = "/target" } +``` + +Set `OPENROUTER_API_KEY`, save the configuration as `routes.toml`, and validate it +before starting the server: + +```bash +export OPENROUTER_API_KEY="sk-or-v1-..." # pragma: allowlist secret +switchyard-server --config routes.toml --dry-run +switchyard-server --config routes.toml --port 4000 +``` + +The parent always uses `parent`. For a delegated request, the classifier sees +the prompt supplied by the parent and selects one configured target. With +`classify_trigger = "new_session"`, Switchyard reuses that decision for later +requests from the same `session + agent` identity. Use `every_request` to +classify each delegated request. `user_turn` is not supported for sub-agent +routing. Harness-maintenance requests continue to the parent target. + +Clients must still request the route ID (`agent` above). An explicit model name +that is not registered as a route is rejected before sub-agent classification. +`message_hash_fallback` is not supported for sub-agent routing because affinity +requires harness-provided child identity. diff --git a/mkdocs.yml b/mkdocs.yml index d934b61e..cbcea947 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -25,6 +25,7 @@ nav: - Architecture: architecture.md - Routing: - Overview: routing_algorithms/overview.md + - Sub-Agent-Aware Routing: routing_algorithms/subagent_routing.md - Random Routing: routing_algorithms/random_routing.md - LLM Classifier Routing: routing_algorithms/llm_classifier_routing.md - Stage-Router Routing: routing_algorithms/stage_router_routing.md From 1540d414ec781117850708487b0a02ffe016df7f Mon Sep 17 00:00:00 2001 From: ayushag Date: Thu, 20 Aug 2026 00:56:39 -0700 Subject: [PATCH 3/4] fix: routing record not populated issue in terminal events Signed-off-by: ayushag --- crates/switchyard-server/src/usage_metrics.rs | 96 +++++++++++++++++-- 1 file changed, 90 insertions(+), 6 deletions(-) diff --git a/crates/switchyard-server/src/usage_metrics.rs b/crates/switchyard-server/src/usage_metrics.rs index 1a3e10cb..754df811 100644 --- a/crates/switchyard-server/src/usage_metrics.rs +++ b/crates/switchyard-server/src/usage_metrics.rs @@ -39,6 +39,8 @@ pub(crate) fn observe( LlmResponse::Stream(mut stream) => { let wrapped = async_stream::stream! { let mut latest_usage = None; + let mut terminal_seen = false; + let mut recorded = false; while let Some(item) = stream.next().await { let failed = match &item { Err(_) => true, @@ -52,23 +54,42 @@ pub(crate) fn observe( }; if let Ok(event) = &item { for chunk in event.normalized() { - if let LlmResponseChunk::Usage(usage) = chunk { - latest_usage = Some(usage.clone()); + match chunk { + LlmResponseChunk::Usage(usage) => { + latest_usage = Some(usage.clone()); + } + LlmResponseChunk::MessageStop { .. } => { + terminal_seen = true; + } + _ => {} } } } if failed { record_stream_error(&stats, &model); } + // Responses clients may stop polling immediately after the terminal event. + // Commit first when that event already carries the final usage. + if !failed && !recorded && terminal_seen + && let Some(usage) = latest_usage.as_ref() + { + record_terminal(&stats, usage, &model, started, cache_eligible); + if let Some((log, context)) = routing_log.as_ref() { + log.append(context.clone(), &model, None, usage); + } + recorded = true; + } yield item; if failed { return; } } - let usage = latest_usage.unwrap_or_default(); - record_terminal(&stats, &usage, &model, started, cache_eligible); - if let Some((log, context)) = routing_log { - log.append(context, &model, None, &usage); + if !recorded { + let usage = latest_usage.unwrap_or_default(); + record_terminal(&stats, &usage, &model, started, cache_eligible); + if let Some((log, context)) = routing_log { + log.append(context, &model, None, &usage); + } } }; LlmResponse::Stream(Box::pin(wrapped)) @@ -160,3 +181,66 @@ fn record_latency(model: &str, latency: Duration) { .build() .record(latency.as_secs_f64() * 1000.0, &attributes(model)); } + +#[cfg(test)] +mod tests { + use futures_util::{StreamExt, stream}; + use switchyard_protocol::{LlmResponseChunk, LlmResponseStreamEvent, Metadata, Response}; + + use super::*; + + /// An OpenAI Responses client may stop polling immediately after receiving the + /// terminal `response.completed` event. Switchyard must record usage and routing + /// data before returning that event because the stream wrapper will not resume + /// after the client drops it. + #[tokio::test] + async fn terminal_event_is_recorded_before_the_client_drops_the_stream() { + let dir = tempfile::tempdir().expect("temp dir"); + let log = SharedRoutingLog::new(dir.path().join("routing.jsonl")).expect("routing log"); + let context = RoutingLogContext::from_metadata(&Metadata { + session_id: Some("streaming-session".to_string()), + ..Metadata::default() + }); + let usage = Usage { + input_tokens: Some(10), + output_tokens: Some(3), + ..Usage::default() + }; + let source = stream::iter([Ok(LlmResponseStreamEvent::new(vec![ + LlmResponseChunk::Usage(usage), + LlmResponseChunk::MessageStop { reason: None }, + ]))]); + let response = Response { + llm_response: LlmResponse::Stream(Box::pin(source)), + metadata: None, + }; + let stats = StatsAccumulator::default(); + let observed = observe( + response, + "model/worker", + Instant::now(), + stats.clone(), + 0.0, + Some((log.clone(), context)), + ); + + let LlmResponse::Stream(mut observed) = observed.llm_response else { + panic!("expected stream"); + }; + assert!(observed.next().await.is_some()); + drop(observed); + + let routing = log + .snapshot_session("streaming-session") + .expect("read routing log") + .expect("terminal event was recorded"); + let routing = serde_json::to_value(routing).expect("serialize routing stats"); + assert_eq!(routing["models"]["model/worker"]["calls"], 1); + assert_eq!(routing["models"]["model/worker"]["prompt_tokens"], 10); + assert_eq!(routing["models"]["model/worker"]["completion_tokens"], 3); + + let process = stats.snapshot(); + assert_eq!(process.models["model/worker"].prompt_tokens, 10); + assert_eq!(process.models["model/worker"].completion_tokens, 3); + } +} From e1bfdbaa7bec9a8d1397c1a983a86a1f3a39fa75 Mon Sep 17 00:00:00 2001 From: ayushag Date: Thu, 20 Aug 2026 08:49:13 -0700 Subject: [PATCH 4/4] chore: merge tests Signed-off-by: ayushag --- crates/libsy/src/algorithms/passthrough.rs | 95 ++++++---------------- crates/switchyard-server/src/config.rs | 69 ++++++---------- 2 files changed, 53 insertions(+), 111 deletions(-) diff --git a/crates/libsy/src/algorithms/passthrough.rs b/crates/libsy/src/algorithms/passthrough.rs index 26721e4e..7358a6b1 100644 --- a/crates/libsy/src/algorithms/passthrough.rs +++ b/crates/libsy/src/algorithms/passthrough.rs @@ -144,44 +144,30 @@ mod tests { completion_text, text_request, }; - struct AlternatingClassifier { + struct ScriptedClassifier { calls: AtomicUsize, } #[async_trait] - impl Classifier for AlternatingClassifier { + impl Classifier for ScriptedClassifier { async fn score( &self, _state: &mut State, _request: &mut Request, _driver: Option<&Driver>, ) -> crate::Result<(Classification, Option)> { - let target = if self.calls.fetch_add(1, Ordering::Relaxed) == 0 { - "worker" - } else { - "reviewer" - }; - Ok(( - Classification::Scores(vec![Score { + let scores = match self.calls.fetch_add(1, Ordering::Relaxed) { + 0 => vec![Score { confidence: 1.0, - target: ModelId::from(target), - }]), - None, - )) - } - } - - struct AbstainingClassifier; - - #[async_trait] - impl Classifier for AbstainingClassifier { - async fn score( - &self, - _state: &mut State, - _request: &mut Request, - _driver: Option<&Driver>, - ) -> crate::Result<(Classification, Option)> { - Ok((Classification::Scores(Vec::new()), None)) + target: ModelId::from("worker"), + }], + 1 => vec![Score { + confidence: 1.0, + target: ModelId::from("reviewer"), + }], + _ => Vec::new(), + }; + Ok((Classification::Scores(scores), None)) } } @@ -243,51 +229,17 @@ mod tests { } #[tokio::test] - async fn parent_traffic_keeps_the_passthrough_target() -> crate::Result<()> { - let router = configured(Arc::new(AlternatingClassifier { - calls: AtomicUsize::new(0), - }))?; - - let (selected, _) = test_drive(router, request(None), echo()).await?; - - assert_eq!(selected, "parent"); - Ok(()) - } - - #[tokio::test] - async fn classifier_decision_is_retained_per_child() -> crate::Result<()> { - let classifier = Arc::new(AlternatingClassifier { + async fn routes_parent_and_children_with_affinity_and_default() -> crate::Result<()> { + let classifier = Arc::new(ScriptedClassifier { calls: AtomicUsize::new(0), }); let router = configured(classifier.clone())?; + let (parent, _) = test_drive(router.clone(), request(None), echo()).await?; let (first, _) = test_drive(router.clone(), child("child-1"), echo()).await?; let (same_child, _) = test_drive(router.clone(), child("child-1"), echo()).await?; - let (sibling, _) = test_drive(router, child("child-2"), echo()).await?; - - assert_eq!(first, "worker"); - assert_eq!(same_child, "worker"); - assert_eq!(sibling, "reviewer"); - assert_eq!(classifier.calls.load(Ordering::Relaxed), 2); - Ok(()) - } - - #[tokio::test] - async fn an_abstaining_child_classifier_uses_the_child_default() -> crate::Result<()> { - let router = configured(Arc::new(AbstainingClassifier))?; - - let (selected, _) = test_drive(router, child("child-1"), echo()).await?; - - assert_eq!(selected, "worker"); - Ok(()) - } - - #[tokio::test] - async fn maintenance_traffic_ignores_an_existing_child_pin() -> crate::Result<()> { - let router = configured(Arc::new(AlternatingClassifier { - calls: AtomicUsize::new(0), - }))?; - test_drive(router.clone(), child("child-1"), echo()).await?; + let (sibling, _) = test_drive(router.clone(), child("child-2"), echo()).await?; + let (defaulted, _) = test_drive(router.clone(), child("child-3"), echo()).await?; let maintenance = request(Some(Metadata { session_id: Some("session-1".to_string()), agent_id: Some("child-1".to_string()), @@ -295,10 +247,15 @@ mod tests { is_delegated_work: false, ..Metadata::default() })); + let (maintenance, _) = test_drive(router, maintenance, echo()).await?; - let (selected, _) = test_drive(router, maintenance, echo()).await?; - - assert_eq!(selected, "parent"); + assert_eq!(parent, "parent"); + assert_eq!(first, "worker"); + assert_eq!(same_child, "worker"); + assert_eq!(sibling, "reviewer"); + assert_eq!(defaulted, "worker"); + assert_eq!(maintenance, "parent"); + assert_eq!(classifier.calls.load(Ordering::Relaxed), 3); Ok(()) } diff --git a/crates/switchyard-server/src/config.rs b/crates/switchyard-server/src/config.rs index 27059a77..ebecfa21 100644 --- a/crates/switchyard-server/src/config.rs +++ b/crates/switchyard-server/src/config.rs @@ -1379,24 +1379,7 @@ target = "weak" } } - #[test] - fn builds_all_supported_algorithm_types() -> ServerResult<()> { - let state = server_state_from_toml(VALID_CONFIG)?; - // The model id array is sorted alphabetically - assert_eq!( - state.models().collect::>(), - [ - "switchyard/classifier", - "switchyard/noop", - "switchyard/passthrough", - "switchyard/random", - ] - ); - Ok(()) - } - - #[test] - fn passthrough_accepts_a_custom_subagent_classifier() -> ServerResult<()> { + fn passthrough_with_subagent_classifier(extra: &str) -> String { let configured = VALID_CONFIG.replace( "[routes.passthrough]\nid = \"switchyard/passthrough\"\ntype = \"passthrough\"\ntarget = \"weak\"", r#"[routes.passthrough] @@ -1413,36 +1396,34 @@ response_schema = '{"type":"object","properties":{"target":{"type":"string","enu policy = { type = "target_selector", selector = "/target" } classify_trigger = "new_session""#, ); + configured.replace( + "classify_trigger = \"new_session\"", + &format!("classify_trigger = \"new_session\"{extra}"), + ) + } - server_state_from_toml(&configured)?; + #[test] + fn builds_all_supported_algorithm_types() -> ServerResult<()> { + let state = server_state_from_toml(VALID_CONFIG)?; + // The model id array is sorted alphabetically + assert_eq!( + state.models().collect::>(), + [ + "switchyard/classifier", + "switchyard/noop", + "switchyard/passthrough", + "switchyard/random", + ] + ); Ok(()) } #[test] - fn passthrough_subagent_classifier_rejects_message_hash_fallback() { - let configured = VALID_CONFIG.replace( - "[routes.passthrough]\nid = \"switchyard/passthrough\"\ntype = \"passthrough\"\ntarget = \"weak\"", - r#"[routes.passthrough] -id = "switchyard/passthrough" -type = "passthrough" -target = "weak" - -[routes.passthrough.subagent_classifier] -classifier_target = "classifier" -targets = ["strong", "weak"] -default_target = "weak" -prompt = "Select a target for this delegated task." -response_schema = '{"type":"object","properties":{"target":{"type":"string","enum":["strong","weak"]}},"required":["target"],"additionalProperties":false}' -policy = { type = "target_selector", selector = "/target" } -classify_trigger = "new_session" -message_hash_fallback = true"#, - ); + fn passthrough_accepts_a_custom_subagent_classifier() -> ServerResult<()> { + let configured = passthrough_with_subagent_classifier(""); - assert!( - error_message(&configured).contains("cannot use message_hash_fallback"), - "{}", - error_message(&configured) - ); + server_state_from_toml(&configured)?; + Ok(()) } #[test] @@ -1649,6 +1630,10 @@ classifier_magic = true ), "message_hash_fallback requires classify_trigger = new_session", ), + ( + passthrough_with_subagent_classifier("\nmessage_hash_fallback = true"), + "cannot use message_hash_fallback", + ), ( VALID_CONFIG.replace( "base_threshold = 0.5",