diff --git a/benchmark/run_cost.py b/benchmark/run_cost.py new file mode 100644 index 000000000..3b7f3250f --- /dev/null +++ b/benchmark/run_cost.py @@ -0,0 +1,105 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-model cost breakdown for a benchmark run's routing log. + +Aggregates ``routing_requests.jsonl`` by served model and prices each bucket with +:func:`switchyard.lib.cost_estimator.estimate_model_cost`, which splits input into +base / cache-read / cache-write tiers. That split matters: agentic runs are +cache-dominated (90%+ of prompt tokens are cache reads is normal), so a flat +input rate overstates cost several-fold. + +On a run that routes sub-agent work to a different tier, the served-model column +doubles as the parent/child attribution key — the routing log records no +``agent_id``, so which model answered is the only role signal available. + +Models absent from the price table contribute $0 and are listed separately, so an +unpriced model reads as a gap rather than as a free one. + +Usage: + uv run python benchmark/run_cost.py --run benchmark/tb_runs/ +""" + +from __future__ import annotations + +import argparse +import collections +import json +from pathlib import Path + +from switchyard.lib.cost_estimator import MODEL_PRICING, estimate_model_cost + +_TOKEN_FIELDS = ( + "prompt_tokens", + "completion_tokens", + "cached_tokens", + "cache_creation_tokens", +) + + +def _aggregate(routing_log: Path) -> dict[str, collections.Counter[str]]: + """Sum request counts and token fields per served model.""" + totals: dict[str, collections.Counter[str]] = collections.defaultdict(collections.Counter) + for line in routing_log.read_text().splitlines(): + if not line.strip(): + continue + record = json.loads(line) + bucket = totals[record.get("model") or ""] + bucket["reqs"] += 1 + for field in _TOKEN_FIELDS: + bucket[field] += record.get(field) or 0 + return totals + + +def _mean_agg_score(run: Path) -> tuple[float, int] | None: + """Mean rubric ``agg_score`` across the run's scored tasks, if any.""" + scores = [ + json.loads(path.read_text())["agg_score"] + for path in run.glob("jobs/*/task-*/verifier/evaluation_results.json") + ] + return (sum(scores) / len(scores), len(scores)) if scores else None + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--run", required=True, type=Path, help="Run directory under tb_runs/") + args = parser.parse_args() + + totals = _aggregate(args.run / "routing_requests.jsonl") + if not totals: + print("no routing records found") + return + + unpriced = [model for model in totals if model not in MODEL_PRICING] + total_cost = 0.0 + header = f"{'model':46s} {'reqs':>5s} {'in':>11s} {'out':>9s} {'cached':>11s} {'USD':>9s}" + print(header) + for model, bucket in sorted(totals.items(), key=lambda item: -item[1]["reqs"]): + cost = estimate_model_cost( + model, + bucket["prompt_tokens"], + bucket["completion_tokens"], + bucket["cached_tokens"], + bucket["cache_creation_tokens"], + )["total_cost"] + total_cost += cost + print( + f"{model:46s} {bucket['reqs']:5d} {bucket['prompt_tokens']:11,} " + f"{bucket['completion_tokens']:9,} {bucket['cached_tokens']:11,} {cost:9.3f}" + ) + + requests = sum(bucket["reqs"] for bucket in totals.values()) + print(f"\nTOTAL: ${total_cost:.2f} over {requests} requests") + if unpriced: + # Priced at zero by estimate_model_cost — surface it so the total is not + # mistaken for complete. + print(f"NOTE: no price entry for {', '.join(sorted(unpriced))} — excluded from the total") + + scored = _mean_agg_score(args.run) + if scored: + mean, count = scored + print(f"mean agg_score: {mean:.3f} over {count} tasks → ${total_cost / count:.2f}/task") + + +if __name__ == "__main__": + main() diff --git a/crates/libsy/src/algorithms.rs b/crates/libsy/src/algorithms.rs index b81e62217..25f692e25 100644 --- a/crates/libsy/src/algorithms.rs +++ b/crates/libsy/src/algorithms.rs @@ -10,12 +10,11 @@ pub mod fall_through; pub mod llm_class; pub mod noop; pub mod rand; -pub mod subagent_override; pub use fall_through::{FallThrough, FallThroughDecision}; pub use llm_class::{ClassifierDecision, ClassifierTier, LlmClassifier}; pub use noop::{Noop, NoopDecision}; pub use rand::{Random, RandomDecision}; -pub use subagent_override::{SubagentDecision, SubagentOverride}; +pub use util::{AffinityRouter, SubagentOverride}; -pub(crate) mod util; +pub mod util; diff --git a/crates/libsy/src/algorithms/fall_through.rs b/crates/libsy/src/algorithms/fall_through.rs index 3e08579f9..40c302a04 100644 --- a/crates/libsy/src/algorithms/fall_through.rs +++ b/crates/libsy/src/algorithms/fall_through.rs @@ -75,6 +75,19 @@ impl FallThrough { self.classifiers.push(classifier); self } + + /// Registers one dual-role component in *both* the processor chain and the classifier + /// cascade. + /// + /// A component that writes state as a [`Processor`] and reads it back as a + /// [`Classifier`] — such as [`AffinityRouter`](crate::algorithms::AffinityRouter) — + /// shares that state through the instance, so both roles must be the same `Arc`. + /// Registering the two separately is easy to half-wire: omit the processor and the + /// classifier silently never sees an assignment. This registers both at once. + pub fn with_component(self, component: Arc) -> Self { + self.with_processor(component.clone()) + .with_classifier(component) + } } #[async_trait] impl Algorithm for FallThrough { diff --git a/crates/libsy/src/algorithms/subagent_override.rs b/crates/libsy/src/algorithms/subagent_override.rs deleted file mode 100644 index b64c6b868..000000000 --- a/crates/libsy/src/algorithms/subagent_override.rs +++ /dev/null @@ -1,206 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 - -//! Sub-agent override combinator built on the [`Algorithm`] interfaces. -//! -//! Wraps any algorithm without changing its behavior for normal traffic. A -//! request whose [`Metadata`] marks delegated sub-agent work -//! ([`Metadata::is_subagent_work`]) is served by one fixed worker target — -//! keeping a sub-agent loop on an intentional, cache-compatible target — -//! while every other request delegates to the wrapped algorithm. The wrapped -//! algorithm never learns about harnesses or lineage headers, and a worker -//! failure surfaces as a normal target error rather than re-entering the -//! wrapped algorithm. - -use std::sync::Arc; - -use async_trait::async_trait; - -use crate::{Algorithm, Context, Decision, Driver, LlmTarget, Metadata, Request, Response, Result}; - -/// Decision produced by [`SubagentOverride`] when it routes to the worker target. -pub struct SubagentDecision { - /// The fixed worker target selected for the sub-agent request. - pub selected_model: String, - /// Human-readable explanation of the override. - pub reasoning: String, -} - -impl Decision for SubagentDecision { - fn selected_model(&self) -> &str { - &self.selected_model - } - - fn reasoning(&self) -> Option<&str> { - Some(&self.reasoning) - } - - fn as_any(&self) -> &dyn std::any::Any { - self - } -} - -/// Routes delegated sub-agent work to a fixed worker target; delegates the rest. -pub struct SubagentOverride { - inner: Arc, - worker: LlmTarget, -} - -impl SubagentOverride { - /// Wraps `inner`, sending recognized sub-agent work to `worker` instead. - /// - /// Wrap it in an [`Arc`] and drive it with [`Algorithm::run`] or - /// [`Algorithm::run_stream`]. - pub fn new(inner: Arc, worker: LlmTarget) -> Self { - Self { inner, worker } - } -} - -#[async_trait] -impl Algorithm for SubagentOverride { - fn name(&self) -> &str { - "subagent_override" - } - - async fn create_run_task( - self: Arc, - ctx: Context, - driver: Driver, - request: Request, - ) -> Result { - let is_subagent_work = request - .metadata - .as_ref() - .is_some_and(Metadata::is_subagent_work); - if !is_subagent_work { - return Arc::clone(&self.inner) - .create_run_task(ctx, driver, request) - .await; - } - - let selected = self.worker.semantic_name.clone(); - let decision: Arc = Arc::new(SubagentDecision { - reasoning: format!("sub-agent work routed to fixed worker target '{selected}'"), - selected_model: selected, - }); - driver.info(ctx.clone(), Arc::clone(&decision)).await?; - driver - .call_llm_target(ctx, &self.worker, request, decision) - .await - } -} - -#[cfg(test)] -mod tests { - use super::*; - use std::collections::BTreeMap; - - use switchyard_protocol::{completion_text, text_request, text_response}; - - use crate::algorithms::Random; - use crate::{LlmResponse, LlmTargetSet, RoutedLlmClient}; - - /// Echoes the selected target so tests can inspect which target was called. - struct EchoClient; - - #[async_trait] - impl RoutedLlmClient for EchoClient { - async fn call( - &self, - _ctx: Context, - _request: Request, - decision: Arc, - ) -> std::result::Result { - Ok(Response { - llm_response: LlmResponse::Agg(text_response(None, decision.selected_model())), - metadata: None, - }) - } - } - - fn target(name: &str) -> LlmTarget { - LlmTarget { - semantic_name: name.to_string(), - llm_client: Some(Arc::new(EchoClient)), - } - } - - fn request(headers: &[(&str, &str)]) -> Request { - let metadata = (!headers.is_empty()).then(|| { - Metadata::from_headers( - &headers - .iter() - .map(|(name, value)| ((*name).to_string(), (*value).to_string())) - .collect::>(), - ) - }); - Request { - llm_request: text_request(Some("auto".to_string()), "hi"), - raw_request: None, - metadata, - } - } - - /// Wraps single-target random routing so the inner selection is deterministic. - fn algorithm() -> Arc { - let inner: Arc = - Arc::new(Random::new(LlmTargetSet::new(vec![target("orchestrator")]))); - Arc::new(SubagentOverride::new(inner, target("worker"))) - } - - async fn selected_model(headers: &[(&str, &str)]) -> Result { - let (trace, response) = algorithm() - .run(Context::default(), request(headers)) - .await?; - let selected = response - .llm_response - .as_agg() - .map(completion_text) - .unwrap_or_default(); - assert_eq!( - trace.last().map(|d| d.selected_model().to_string()), - Some(selected.clone()) - ); - Ok(selected) - } - - #[tokio::test] - async fn requests_without_metadata_delegate_to_the_wrapped_algorithm() -> Result<()> { - assert_eq!(selected_model(&[]).await?, "orchestrator"); - Ok(()) - } - - #[tokio::test] - async fn subagent_work_is_routed_to_the_worker_target() -> Result<()> { - // Claude Code child-agent lineage. - let claude = &[ - ("x-claude-code-session-id", "root"), - ("x-claude-code-agent-id", "child-1"), - ]; - assert_eq!(selected_model(claude).await?, "worker"); - - // Codex delegated-work kinds. - assert_eq!( - selected_model(&[("x-openai-subagent", "review")]).await?, - "worker" - ); - assert_eq!( - selected_model(&[("x-openai-subagent", "collab_spawn")]).await?, - "worker" - ); - Ok(()) - } - - #[tokio::test] - async fn harness_maintenance_turns_stay_on_the_wrapped_algorithm() -> Result<()> { - assert_eq!( - selected_model(&[("x-openai-subagent", "compact")]).await?, - "orchestrator" - ); - assert_eq!( - selected_model(&[("x-switchyard-is-subagent", "false")]).await?, - "orchestrator" - ); - Ok(()) - } -} diff --git a/crates/libsy/src/algorithms/util.rs b/crates/libsy/src/algorithms/util.rs index 99c18a44f..601cffef3 100644 --- a/crates/libsy/src/algorithms/util.rs +++ b/crates/libsy/src/algorithms/util.rs @@ -1,7 +1,8 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -mod affinity; +pub mod affinity; +pub mod subagent; -#[allow(unused_imports)] -pub(crate) use affinity::AffinityRouter; +pub use affinity::AffinityRouter; +pub use subagent::SubagentOverride; diff --git a/crates/libsy/src/algorithms/util/affinity.rs b/crates/libsy/src/algorithms/util/affinity.rs index 4a12faf38..9f1c793d7 100644 --- a/crates/libsy/src/algorithms/util/affinity.rs +++ b/crates/libsy/src/algorithms/util/affinity.rs @@ -18,8 +18,6 @@ //! session. [`AffinityRouter::for_subagents`] narrows affinity to explicitly identified //! child agents, leaving root traffic to later classifiers on every turn. -#![allow(dead_code)] - use std::collections::{HashMap, HashSet}; use async_trait::async_trait; diff --git a/crates/libsy/src/algorithms/util/subagent.rs b/crates/libsy/src/algorithms/util/subagent.rs new file mode 100644 index 000000000..6b5d5b971 --- /dev/null +++ b/crates/libsy/src/algorithms/util/subagent.rs @@ -0,0 +1,155 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Sub-agent override as a single SDK component. +//! +//! [`SubagentOverride`] scores one fixed worker target for requests carrying delegated +//! sub-agent work ([`Metadata::is_subagent_work`]) and abstains for everything else, so a +//! cascade falls through to its later classifiers on ordinary traffic. +//! +//! It is stateless and holds only the worker's *name*: the +//! [`FallThrough`](crate::algorithms::FallThrough) cascade resolves that name against its +//! target set. Keeping the policy independent of any memory of past decisions is what lets +//! it compose with a stateful classifier such as +//! [`AffinityRouter`](crate::algorithms::AffinityRouter) — the override decides *which* +//! target delegated work belongs on, affinity decides *how long* a decision lives, and +//! neither needs to know about the other. + +use async_trait::async_trait; + +use crate::{Classification, Classifier, Driver, Metadata, Request, Result, Score, State}; + +/// Scores a fixed worker target for delegated sub-agent work; abstains otherwise. +pub struct SubagentOverride { + /// Name of the worker target, resolved by the cascade against its target set. + worker: String, +} + +impl SubagentOverride { + /// Creates an override scoring `worker` for delegated sub-agent work. + /// + /// `worker` must name a target in the cascade's set, or routing a sub-agent request + /// fails with [`LibsyError::TargetNotFound`](crate::LibsyError::TargetNotFound). + pub fn new(worker: impl Into) -> Self { + Self { + worker: worker.into(), + } + } +} + +#[async_trait] +impl Classifier for SubagentOverride { + async fn score( + &self, + _state: &mut State, + request: &Request, + _driver: Option<&Driver>, + ) -> Result { + // Delegated *work* only. A harness maintenance turn (e.g. Codex `compact`) carries + // sub-agent lineage but is not delegated work, so it abstains and routes normally. + let is_delegated_work = request + .metadata + .as_ref() + .is_some_and(Metadata::is_subagent_work); + Ok(Classification::Scores(if is_delegated_work { + vec![Score { + confidence: 1.0, + target: self.worker.clone(), + }] + } else { + Vec::new() + })) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::BTreeMap; + + use switchyard_protocol::text_request; + + fn request(headers: &[(&str, &str)]) -> Request { + let metadata = (!headers.is_empty()).then(|| { + Metadata::from_headers( + &headers + .iter() + .map(|(name, value)| ((*name).to_string(), (*value).to_string())) + .collect::>(), + ) + }); + Request { + llm_request: text_request(Some("auto".to_string()), "hi"), + raw_request: None, + metadata, + } + } + + /// Scores `headers` through the override, returning the winning target if it scored. + async fn selected(headers: &[(&str, &str)]) -> Result> { + let mut state = State::default(); + let classification = SubagentOverride::new("worker") + .score(&mut state, &request(headers), None) + .await?; + Ok(classification.argmax(false)?.map(|score| score.target)) + } + + #[tokio::test] + async fn requests_without_metadata_abstain() -> Result<()> { + assert_eq!(selected(&[]).await?, None); + Ok(()) + } + + #[tokio::test] + async fn subagent_work_scores_the_worker() -> Result<()> { + // Claude Code child-agent lineage. + let claude = &[ + ("x-claude-code-session-id", "root"), + ("x-claude-code-agent-id", "child-1"), + ]; + assert_eq!(selected(claude).await?, Some("worker".to_string())); + + // Codex delegated-work kinds. + assert_eq!( + selected(&[("x-openai-subagent", "review")]).await?, + Some("worker".to_string()) + ); + assert_eq!( + selected(&[("x-openai-subagent", "collab_spawn")]).await?, + Some("worker".to_string()) + ); + Ok(()) + } + + #[tokio::test] + async fn harness_maintenance_turns_abstain() -> Result<()> { + assert_eq!(selected(&[("x-openai-subagent", "compact")]).await?, None); + assert_eq!( + selected(&[("x-switchyard-is-subagent", "false")]).await?, + None + ); + Ok(()) + } + + #[tokio::test] + async fn delegated_work_is_scored_definitively() -> Result<()> { + // Confidence 1.0 under `Scores` (never `Ambiguous`), so the cascade stops here + // rather than consulting later classifiers. + let mut state = State::default(); + let classification = SubagentOverride::new("worker") + .score( + &mut state, + &request(&[("x-openai-subagent", "review")]), + None, + ) + .await?; + match classification { + Classification::Scores(scores) => { + assert_eq!(scores.len(), 1); + assert_eq!(scores[0].confidence, 1.0); + } + Classification::Ambiguous(_) => panic!("override must score definitively"), + } + Ok(()) + } +} diff --git a/crates/libsy/tests/subagent_affinity.rs b/crates/libsy/tests/subagent_affinity.rs new file mode 100644 index 000000000..7e54d6380 --- /dev/null +++ b/crates/libsy/tests/subagent_affinity.rs @@ -0,0 +1,269 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Integration tests for composing the sub-agent override with affinity routing. +//! +//! The two are independent classifiers in one [`FallThrough`] cascade: the override decides +//! *which* target delegated work belongs on, affinity decides *how long* that decision +//! lives. These tests drive them through the crate's public API, so they also pin that the +//! pieces needed to compose a cascade are actually exported. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use async_trait::async_trait; + +use switchyard_libsy::algorithms::{AffinityRouter, FallThrough, SubagentOverride}; +use switchyard_libsy::{ + Algorithm, Classification, Classifier, Context, Decision, Driver, LlmResponse, LlmTarget, + LlmTargetSet, Metadata, Request, Response, Result, RoutedLlmClient, Score, SharedState, State, +}; +use switchyard_protocol::{completion_text, text_request, text_response}; + +/// A client that echoes the routed target name back as the completion. +struct EchoClient; + +#[async_trait] +impl RoutedLlmClient for EchoClient { + async fn call( + &self, + _ctx: Context, + _request: Request, + decision: Arc, + ) -> std::result::Result { + Ok(Response { + llm_response: LlmResponse::Agg(text_response(None, decision.selected_model())), + metadata: None, + }) + } +} + +/// The cascade's terminal classifier: always picks the orchestrator. +struct AlwaysOrchestrator; + +#[async_trait] +impl Classifier for AlwaysOrchestrator { + async fn score( + &self, + _state: &mut State, + _request: &Request, + _driver: Option<&Driver>, + ) -> Result { + Ok(Classification::Scores(vec![Score { + confidence: 0.5, + target: "orchestrator".to_string(), + }])) + } +} + +fn targets() -> LlmTargetSet { + LlmTargetSet::new( + ["orchestrator", "worker", "reviewer"] + .iter() + .map(|name| LlmTarget { + semantic_name: (*name).to_string(), + llm_client: Some(Arc::new(EchoClient) as Arc), + }) + .collect(), + ) +} + +fn request(headers: &[(&str, &str)]) -> Request { + Request { + llm_request: text_request(Some("auto".to_string()), "hi"), + raw_request: None, + metadata: Some(Metadata::from_headers( + &headers + .iter() + .map(|(name, value)| ((*name).to_string(), (*value).to_string())) + .collect::>(), + )), + } +} + +/// Affinity replays an existing pin, the override seeds one for delegated work, and the +/// terminal classifier serves everything else. +fn router() -> Arc { + Arc::new( + FallThrough::new(targets()) + .with_component(Arc::new(AffinityRouter::for_subagents())) + .with_classifier(Arc::new(SubagentOverride::new("worker"))) + .with_classifier(Arc::new(AlwaysOrchestrator)), + ) +} + +/// Runs one turn on `ctx`, returning the target that served it. +/// +/// `ctx` carries the per-session [`State`] the processor chain folds into. The affinity +/// assignments are held on the `AffinityRouter` instance instead, so replay follows the +/// shared router rather than the context. +async fn turn( + router: &Arc, + ctx: Context, + headers: &[(&str, &str)], +) -> Result { + let (_, response) = router.clone().run(ctx, request(headers)).await?; + Ok(response + .llm_response + .as_agg() + .map(completion_text) + .unwrap_or_default()) +} + +/// A Claude Code child agent's lineage headers. +fn child(agent: &str) -> Vec<(&str, &str)> { + vec![ + ("x-claude-code-session-id", "session-1"), + ("x-claude-code-agent-id", agent), + ] +} + +#[tokio::test] +async fn root_traffic_falls_through_to_the_terminal_classifier() -> Result<()> { + let ctx = Context::::default(); + // No sub-agent lineage: affinity abstains (subagents only), the override abstains, and + // the terminal classifier serves the turn. + let served = turn(&router(), ctx, &[("x-claude-code-session-id", "session-1")]).await?; + assert_eq!(served, "orchestrator"); + Ok(()) +} + +#[tokio::test] +async fn delegated_work_is_routed_to_the_worker() -> Result<()> { + let ctx = Context::::default(); + assert_eq!(turn(&router(), ctx, &child("child-1")).await?, "worker"); + Ok(()) +} + +#[tokio::test] +async fn the_override_seeds_a_pin_that_affinity_replays() -> Result<()> { + let router = router(); + let ctx = Context::::default(); + + // Turn 1: affinity has no assignment, so the override decides and the decision is + // replayed into affinity. + assert_eq!( + turn(&router, ctx.clone(), &child("child-1")).await?, + "worker" + ); + + // Turns 2-3: the lineage header still identifies the child, but affinity — not the + // override — is now the classifier that answers, because it scores first. + assert_eq!( + turn(&router, ctx.clone(), &child("child-1")).await?, + "worker" + ); + assert_eq!(turn(&router, ctx, &child("child-1")).await?, "worker"); + Ok(()) +} + +#[tokio::test] +async fn harness_maintenance_turns_are_not_forced_to_the_worker() -> Result<()> { + let router = router(); + let ctx = Context::::default(); + + // A Codex `compact` turn is sub-agent *lineage* but not delegated *work*: the two + // predicates differ on purpose, so the override abstains and it routes normally. + let served = turn( + &router, + ctx, + &[ + ("x-codex-session-id", "session-1"), + ("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( + FallThrough::new(targets()) + .with_component(affinity) + .with_classifier(Arc::new(SubagentOverride::new(worker))) + .with_classifier(Arc::new(AlwaysOrchestrator)), + ) +} + +#[tokio::test] +async fn the_pin_outlives_the_policy_that_seeded_it() -> Result<()> { + // Two cascades sharing one affinity instance, with overrides that disagree. The first + // seeds "worker"; the second would score "reviewer" but is never consulted, because + // affinity scores earlier in the cascade and replays the existing pin. This is what + // distinguishes a real replay from the override simply repeating itself. + let affinity = Arc::new(AffinityRouter::for_subagents()); + let seed = router_overriding_to(affinity.clone(), "worker"); + let rebound = router_overriding_to(affinity, "reviewer"); + + assert_eq!( + turn(&seed, Context::default(), &child("child-1")).await?, + "worker" + ); + assert_eq!( + turn(&rebound, Context::default(), &child("child-1")).await?, + "worker" + ); + Ok(()) +} + +#[tokio::test] +async fn without_a_shared_pin_the_second_policy_wins() -> Result<()> { + // The negative control for the test above: with independent affinity instances there + // is no pin to replay, so the second cascade's own override decides. Without this, + // the replay assertion could pass for the wrong reason. + let seed = router_overriding_to(Arc::new(AffinityRouter::for_subagents()), "worker"); + let rebound = router_overriding_to(Arc::new(AffinityRouter::for_subagents()), "reviewer"); + + assert_eq!( + turn(&seed, Context::default(), &child("child-1")).await?, + "worker" + ); + assert_eq!( + turn(&rebound, Context::default(), &child("child-1")).await?, + "reviewer" + ); + Ok(()) +} + +#[tokio::test] +async fn distinct_children_are_pinned_independently() -> Result<()> { + // Two cascades sharing one affinity instance but disagreeing on the worker. Routing + // both children through the *same* override would pass whether or not affinity keys + // per agent, so the sibling is sent through the cascade that scores "reviewer": it + // can only come back "reviewer" if it did not inherit child-1's pin. + let affinity = Arc::new(AffinityRouter::for_subagents()); + let seed = router_overriding_to(affinity.clone(), "worker"); + let sibling = router_overriding_to(affinity, "reviewer"); + + assert_eq!( + turn(&seed, Context::default(), &child("child-1")).await?, + "worker" + ); + assert_eq!( + turn(&sibling, Context::default(), &child("child-2")).await?, + "reviewer" + ); + // child-1's own pin is untouched by the sibling's assignment: affinity replays it + // even through the cascade whose override would otherwise score "reviewer". + assert_eq!( + turn(&sibling, Context::default(), &child("child-1")).await?, + "worker" + ); + Ok(()) +} + +#[tokio::test] +async fn a_cascade_without_the_override_still_routes_root_traffic() -> Result<()> { + // Affinity and the override are independent: dropping the override leaves a valid + // cascade, which is the point of composing them rather than nesting one in the other. + let router = Arc::new( + FallThrough::new(targets()) + .with_component(Arc::new(AffinityRouter::for_subagents())) + .with_classifier(Arc::new(AlwaysOrchestrator)), + ); + let ctx = Context::::default(); + assert_eq!(turn(&router, ctx, &child("child-1")).await?, "orchestrator"); + Ok(()) +} diff --git a/crates/switchyard-py/src/core_bindings.rs b/crates/switchyard-py/src/core_bindings.rs index ce183924f..fece0fe6b 100644 --- a/crates/switchyard-py/src/core_bindings.rs +++ b/crates/switchyard-py/src/core_bindings.rs @@ -10,6 +10,7 @@ pub(crate) mod request; pub(crate) mod response; pub(crate) mod roles; pub(crate) mod session; +pub(crate) mod subagent; pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { context::register(module)?; @@ -17,5 +18,6 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { response::register(module)?; roles::register(module)?; session::register(module)?; + subagent::register(module)?; Ok(()) } diff --git a/crates/switchyard-py/src/core_bindings/subagent.rs b/crates/switchyard-py/src/core_bindings/subagent.rs new file mode 100644 index 000000000..2948c8576 --- /dev/null +++ b/crates/switchyard-py/src/core_bindings/subagent.rs @@ -0,0 +1,29 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Python binding for the canonical sub-agent detection policy. +//! +//! Profiles must not sniff lineage headers themselves: the fact (explicit +//! `x-switchyard-is-subagent`, Claude Code agent lineage, Codex/relay markers) and the +//! work-vs-maintenance policy both live in the protocol crate, so every engine — the libsy +//! classifier and the serve-path profile alike — answers "is this delegated work?" +//! identically. + +use std::collections::BTreeMap; + +use pyo3::prelude::*; + +/// Whether `headers` mark this request as delegated sub-agent *work*. +/// +/// Wraps [`switchyard_protocol::Metadata::from_headers`] for the lineage fact and +/// `is_subagent_work` for the kind policy, so harness maintenance turns (Codex `compact`, +/// `memory_consolidation`) stay on normal routing rather than being sent to a worker target. +#[pyfunction] +fn is_subagent_request(headers: BTreeMap) -> bool { + switchyard_protocol::Metadata::from_headers(&headers).is_subagent_work() +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(is_subagent_request, module)?)?; + Ok(()) +} diff --git a/crates/switchyard-py/src/libsy_bindings.rs b/crates/switchyard-py/src/libsy_bindings.rs index a0e21f139..68984c9eb 100644 --- a/crates/switchyard-py/src/libsy_bindings.rs +++ b/crates/switchyard-py/src/libsy_bindings.rs @@ -9,7 +9,7 @@ use async_trait::async_trait; use pyo3::exceptions::{PyTypeError, PyValueError}; use pyo3::prelude::*; use serde_json::{json, Value}; -use switchyard_libsy::algorithms::{Noop, Random, SubagentOverride}; +use switchyard_libsy::algorithms::{Noop, Random}; use switchyard_libsy::{ AggLlmResponse, Algorithm, Context, Decision, LlmClientError, LlmResponse, LlmTarget, LlmTargetSet, Metadata, Request, Response, RoutedLlmClient, @@ -111,8 +111,8 @@ impl PyAlgorithm { /// /// `headers`, when given, is normalized into the request's correlation /// [`Metadata`] exactly as an HTTP host would (`Metadata::from_headers`), - /// so metadata-driven algorithms such as `subagent_override` see the same - /// signals in Python as when served over HTTP. + /// so metadata-driven algorithms see the same signals in Python as when + /// served over HTTP. #[pyo3(signature = (request, headers=None))] fn run<'py>( &self, @@ -175,24 +175,6 @@ fn random_algorithm(py: Python<'_>, targets: Vec>) -> PyResult

, - inner: Py, - worker: Py, -) -> PyResult { - let inner = Arc::clone(&inner.bind(py).get().inner); - let worker = worker.bind(py).try_borrow()?.clone_core(py); - Ok(PyAlgorithm::new(Arc::new(SubagentOverride::new( - inner, worker, - )))) -} - fn other_python_error(error: PyErr) -> LlmClientError { LlmClientError::Ffi { source: Box::new(error), @@ -211,10 +193,6 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { libsy_module.add_class::()?; libsy_module.add_function(wrap_pyfunction!(noop_algorithm, &libsy_module)?)?; libsy_module.add_function(wrap_pyfunction!(random_algorithm, &libsy_module)?)?; - libsy_module.add_function(wrap_pyfunction!( - subagent_override_algorithm, - &libsy_module - )?)?; libsy_module.add("LibsyError", module.getattr("LibsyError")?)?; module.add_submodule(&libsy_module)?; Ok(()) diff --git a/docs/routing_algorithms/overview.md b/docs/routing_algorithms/overview.md index 28ff6b356..75dfcdb40 100644 --- a/docs/routing_algorithms/overview.md +++ b/docs/routing_algorithms/overview.md @@ -89,6 +89,54 @@ routes: Switchyard does not start or manage the model server; it only sends requests to the configured endpoint. +## Sub-agent override (`subagent_target`) + +Any route may name an optional `subagent_target` in its common envelope, +alongside `type`. A request carrying a recognized sub-agent signal — Claude Code +agent-lineage headers, Codex delegated-work kinds (`x-openai-subagent: +collab_spawn` or `review`), or an explicit `x-switchyard-is-subagent: true` — +bypasses the route's own chain and runs as a direct passthrough to that target: + +```yaml +routes: + assistant: + type: model + target: + model: strong-model + subagent_target: + model: cheaper-worker-model +``` + +This keeps a sub-agent loop on one intentional, cache-compatible target instead +of re-routing every worker turn. The worker may live on a different provider +entirely — give it its own `base_url` and `api_key` and one route spans two +upstreams: + +```yaml + subagent_target: + model: my-local-model + base_url: http://localhost:8000/v1 + api_key: dummy + format: openai +``` + +Detection is the protocol crate's canonical policy, shared with the libsy +`SubagentOverride` classifier, so both engines agree on what counts as +delegated work. Harness-maintenance turns (`compact`, `memory_consolidation`) +and unrecognized kinds stay on normal routing, as does everything else when the +field is absent. A worker-target failure surfaces as a normal target error — it +is never silently re-routed through the route's own chain. + +To suppress sub-agent routing for a request that carries a recognized signal, +send `x-switchyard-is-subagent: false`. This explicit header overrides Codex and +Claude Code lineage signals in either direction: `false` keeps the request on +normal routing even when delegated-work headers are present, and `true` marks a +request as a sub-agent even when no harness headers appear. + +The key applies to `model`, `passthrough`, `deterministic`, `escalation_router`, +and `stage_router` routes. `random_routing` expands into its table entries on a +separate path and does not consume it. + ## How session affinity composes Session affinity is configured directly on the LLM classifier router. After diff --git a/switchyard/cli/route_bundle.py b/switchyard/cli/route_bundle.py index a1844900f..727c648b9 100644 --- a/switchyard/cli/route_bundle.py +++ b/switchyard/cli/route_bundle.py @@ -48,6 +48,7 @@ from switchyard.lib.profiles.escalation_router_config import EscalationRouterConfig from switchyard.lib.profiles.random_routing import RandomRoutingConfig from switchyard.lib.profiles.stage_router_config import StageRouterConfig +from switchyard.lib.profiles.subagent_override import SubagentOverrideRuntime from switchyard.lib.route_table import ChainRuntime, RouteTable from switchyard.lib.route_table_builders import ( build_passthrough_table, @@ -141,13 +142,21 @@ def llm_target_to_route_dict(target: LlmTarget) -> dict[str, Any]: "extra_headers", "endpoint", }) +# Envelope keys valid on every route type, whatever its own schema. They configure the +# wrapper around a route's chain rather than the chain itself, so they are accepted by +# _validate_route_keys and excluded from the per-profile config in _route_config. +_ROUTE_ENVELOPE_KEYS = frozenset({ + # Delegated sub-agent work is served by this target instead of the route's own + # chain. Consumed in _build_switchyard_for_route. + "subagent_target", +}) _COMMON_ROUTE_KEYS = frozenset({ "type", "kind", "defaults", "display_name", "description", -}) +}) | _ROUTE_ENVELOPE_KEYS _ROUTE_METADATA_KEYS = frozenset({ "type", "kind", @@ -785,6 +794,66 @@ def _build_switchyard_for_route( stats: StatsAccumulator, pre_routing_request_processors: Sequence[Any] = (), extra_response_processors: Sequence[Any] = (), +) -> ChainRuntime: + """Build the route's chain, wrapping it when the route sets ``subagent_target``.""" + chain = _build_route_chain( + model_id, + route, + route_type=route_type, + target_defaults=target_defaults, + stats=stats, + pre_routing_request_processors=pre_routing_request_processors, + extra_response_processors=extra_response_processors, + ) + worker = _subagent_worker_runtime( + model_id, + route, + target_defaults=target_defaults, + stats=stats, + pre_routing_request_processors=pre_routing_request_processors, + extra_response_processors=extra_response_processors, + ) + return chain if worker is None else SubagentOverrideRuntime(chain, worker) + + +def _subagent_worker_runtime( + model_id: str, + route: Mapping[str, object], + target_defaults: Mapping[str, object], + stats: StatsAccumulator, + pre_routing_request_processors: Sequence[Any] = (), + extra_response_processors: Sequence[Any] = (), +) -> ChainRuntime | None: + """Build the route's sub-agent worker chain, or ``None`` when unset. + + The worker is a plain passthrough to one target: delegated work is served directly + rather than re-routed through the route's own policy. + """ + subagent_raw = route.get("subagent_target") + if subagent_raw is None: + return None + return build_tier_passthrough_switchyard( + _target_value( + subagent_raw, + target_defaults, + default_id=f"{model_id}#subagent", + where=f"route {model_id!r} subagent_target", + ), + stats, + enable_stats=_optional_bool(route.get("enable_stats"), default=True), + extra_request_processors=pre_routing_request_processors, + extra_response_processors=extra_response_processors, + ) + + +def _build_route_chain( + model_id: str, + route: Mapping[str, object], + route_type: str, + target_defaults: Mapping[str, object], + stats: StatsAccumulator, + pre_routing_request_processors: Sequence[Any] = (), + extra_response_processors: Sequence[Any] = (), ) -> ChainRuntime: if route_type in ("model", "passthrough"): # Both kinds resolve to a single-tier passthrough chain — same shape the @@ -1468,7 +1537,7 @@ def _validate_route_keys( route_type: str, ) -> None: where = f"route {model_id!r}" - _validate_allowed_keys(route, _ROUTE_KEYS_BY_TYPE[route_type], where) + _validate_allowed_keys(route, _ROUTE_KEYS_BY_TYPE[route_type] | _ROUTE_ENVELOPE_KEYS, where) if "defaults" in route: defaults = _require_mapping(route["defaults"], f"{where}.defaults") _validate_allowed_keys( diff --git a/switchyard/lib/cost_estimator.py b/switchyard/lib/cost_estimator.py index f47397497..5cdb14888 100644 --- a/switchyard/lib/cost_estimator.py +++ b/switchyard/lib/cost_estimator.py @@ -123,6 +123,19 @@ class ModelPriceData: "openai/nvidia/moonshotai/kimi-k2.5": ModelPriceData( input=0.60, output=2.50, cached=0.15, cache_write=0.60, ), + # Z.ai GLM-5.2 — REFERENCE pricing only. Switchyard reaches this model as + # ``glm-5.2-fp8`` on a self-hosted vLLM deployment, which has no per-token + # billing at all: the real cost is GPU-hours. These rates are + # SiliconFlow's published GLM-5.2 serverless tier (verified 2026-07-27), + # carried so benchmark arms can be compared against commercial models on a + # common basis — "what this traffic would have cost at market rates", not + # what we are billed. Provider rates vary about 2x (OpenRouter lists + # $0.76/$2.42), so treat the absolute number as indicative, not exact. + # OpenAI wire format and no documented cache-write premium, so + # cache_write = input. + "glm-5.2-fp8": ModelPriceData( + input=1.40, output=4.40, cached=0.26, cache_write=1.40, + ), # DeepSeek V4 Flash — official api-docs.deepseek.com standard list # price (post-promo). 284B total / 13B active, 1M-token context # window. Aggressive cache discount (98% off on hits). OpenAI wire @@ -179,7 +192,13 @@ class ModelPriceData: # --- Anthropic Claude on AWS Bedrock (via NVIDIA Inference Hub) --- # 5-minute cache write = 1.25x input; cache read = 0.1x input. # Opus 4.8 keeps the Opus-tier $5/$25 list price (docs.anthropic.com, - # verified 2026-07-11). + # verified 2026-07-11). Opus 5 ships as a drop-in upgrade at the same + # Opus-tier rate (platform.claude.com/docs/en/about-claude/pricing, + # verified 2026-07-27); fast mode ($10/$50) is a separate SKU and is + # not priced here. + "aws/anthropic/bedrock-claude-opus-5": ModelPriceData( + input=5.00, output=25.00, cached=0.50, cache_write=6.25, + ), "aws/anthropic/bedrock-claude-opus-4-8": ModelPriceData( input=5.00, output=25.00, cached=0.50, cache_write=6.25, ), @@ -208,6 +227,9 @@ class ModelPriceData: "azure/anthropic/claude-opus-4-8": ModelPriceData( input=5.00, output=25.00, cached=0.50, cache_write=6.25, ), + "azure/anthropic/claude-opus-5": ModelPriceData( + input=5.00, output=25.00, cached=0.50, cache_write=6.25, + ), "azure/anthropic/claude-sonnet-4-5": ModelPriceData( input=3.00, output=15.00, cached=0.30, cache_write=3.75, ), @@ -221,6 +243,9 @@ class ModelPriceData: input=1.00, output=5.00, cached=0.10, cache_write=1.25, ), # --- Anthropic direct API aliases (no AWS prefix) --- + "claude-opus-5": ModelPriceData( + input=5.00, output=25.00, cached=0.50, cache_write=6.25, + ), "claude-opus-4-8": ModelPriceData( input=5.00, output=25.00, cached=0.50, cache_write=6.25, ), diff --git a/switchyard/lib/profiles/subagent_override.py b/switchyard/lib/profiles/subagent_override.py new file mode 100644 index 000000000..d3c69e42e --- /dev/null +++ b/switchyard/lib/profiles/subagent_override.py @@ -0,0 +1,57 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Route delegated sub-agent requests to a fixed worker runtime.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +from switchyard.lib.request_metadata import CTX_PROFILE_REQUEST_HEADERS +from switchyard_rust.core import is_subagent_request + +if TYPE_CHECKING: + from switchyard.lib.proxy_context import ProxyContext + from switchyard.lib.route_table import ChainRuntime + from switchyard_rust.core import ChatRequest + + +class SubagentOverrideRuntime: + """Wrap a runtime, sending delegated sub-agent work to ``worker`` instead. + + Normal traffic runs the wrapped runtime unchanged. A request whose headers carry a + delegated sub-agent signal runs ``worker`` — a passthrough to the route's configured + ``subagent_target``. The override never rewrites the request or response, and a worker + failure surfaces as a normal target error rather than falling back to the wrapped + runtime. + + Detection is :func:`~switchyard_rust.profiles.is_subagent_request`, the protocol + crate's canonical policy, so harness maintenance turns (Codex ``compact``) stay on + normal routing and every engine agrees on what counts as delegated work. + """ + + def __init__(self, inner: ChainRuntime, worker: ChainRuntime) -> None: + self._inner = inner + self._worker = worker + + def _branch(self, ctx: ProxyContext | None) -> ChainRuntime: + """Select the worker for delegated work, else the wrapped runtime. + + Without a context there are no retained headers to read, so the request cannot be + identified as delegated and routes normally. + """ + if ctx is None: + return self._inner + headers = ctx.metadata.get(CTX_PROFILE_REQUEST_HEADERS) or {} + return self._worker if is_subagent_request(headers) else self._inner + + async def call(self, request: ChatRequest, *, ctx: ProxyContext | None = None) -> Any: + """Execute one request through the branch its headers select.""" + return await self._branch(ctx).call(request, ctx=ctx) + + def iter_components(self) -> list[Any]: + """Return both branches' lifecycle components in startup order.""" + return [*self._inner.iter_components(), *self._worker.iter_components()] + + +__all__ = ["SubagentOverrideRuntime"] diff --git a/switchyard/libsy/algorithms.py b/switchyard/libsy/algorithms.py index d5f2e64a9..434526feb 100644 --- a/switchyard/libsy/algorithms.py +++ b/switchyard/libsy/algorithms.py @@ -5,6 +5,5 @@ from switchyard_rust.libsy import noop as noop from switchyard_rust.libsy import random as random -from switchyard_rust.libsy import subagent_override as subagent_override -__all__ = ["noop", "random", "subagent_override"] +__all__ = ["noop", "random"] diff --git a/switchyard_rust/core.py b/switchyard_rust/core.py index 59cbd0a5a..551fa3aa3 100644 --- a/switchyard_rust/core.py +++ b/switchyard_rust/core.py @@ -582,8 +582,12 @@ def session_key_from_body( body: Mapping[str, Any] | JsonValue, depth: int = 0 ) -> str | None: ... + def is_subagent_request(headers: Mapping[str, str]) -> bool: ... + def __getattr__(name: str) -> object: + if name == "is_subagent_request": + return _load_native().is_subagent_request if name == "SessionCache": return _load_native().SessionCache if name == "session_key_from_body": @@ -801,6 +805,7 @@ def response_is_streaming(response: object) -> bool: "SwitchyardRuntimeError", "SwitchyardUnsupportedRequestTypeError", "SwitchyardUpstreamError", + "is_subagent_request", "request_type_enum", "request_type_matches", "request_type_value", diff --git a/switchyard_rust/libsy.py b/switchyard_rust/libsy.py index ff6c3123e..745309062 100644 --- a/switchyard_rust/libsy.py +++ b/switchyard_rust/libsy.py @@ -10,9 +10,7 @@ from switchyard_rust.core import _load_native -_EXPORTS = frozenset( - {"Algorithm", "LibsyError", "LlmTarget", "noop", "random", "subagent_override"} -) +_EXPORTS = frozenset({"Algorithm", "LibsyError", "LlmTarget", "noop", "random"}) class LlmClient(Protocol): @@ -53,8 +51,6 @@ def noop() -> Algorithm: ... def random(targets: Sequence[LlmTarget]) -> Algorithm: ... - def subagent_override(inner: Algorithm, worker: LlmTarget) -> Algorithm: ... - def __getattr__(name: str) -> object: if name in _EXPORTS: diff --git a/tests/test_libsy_minimal_bindings.py b/tests/test_libsy_minimal_bindings.py index 504ab3cd9..b2014a33a 100644 --- a/tests/test_libsy_minimal_bindings.py +++ b/tests/test_libsy_minimal_bindings.py @@ -99,37 +99,6 @@ async def test_invalid_request_is_rejected_at_the_boundary() -> None: ) -async def test_subagent_override_routes_marked_work_to_the_worker() -> None: - orchestrator = EchoClient("orchestrator") - worker = EchoClient("worker") - algorithm = algorithms.subagent_override( - algorithms.random([LlmTarget("orchestrator", orchestrator)]), - LlmTarget("worker", worker), - ) - - # No headers: the wrapped algorithm serves the request. - _, response = await algorithm.run(request_body()) - assert response["model"] == "orchestrator" - - # Claude Code child-agent lineage routes to the fixed worker target. - decisions, response = await algorithm.run( - request_body(), - headers={ - "x-claude-code-session-id": "root", - "x-claude-code-agent-id": "child-1", - }, - ) - assert response["model"] == "worker" - assert decisions[-1]["selected_model"] == "worker" - - # Harness maintenance stays on the wrapped algorithm. - _, response = await algorithm.run( - request_body(), headers={"x-openai-subagent": "compact"} - ) - assert response["model"] == "orchestrator" - assert len(worker.calls) == 1 - - async def test_client_failure_becomes_libsy_error() -> None: class FailingClient: async def call(self, request: dict[str, Any]) -> dict[str, Any]: diff --git a/tests/test_subagent_routing.py b/tests/test_subagent_routing.py new file mode 100644 index 000000000..edd82e466 --- /dev/null +++ b/tests/test_subagent_routing.py @@ -0,0 +1,137 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Serve-path sub-agent routing: the ``subagent_target`` route envelope key.""" + +from __future__ import annotations + +from typing import Any + +import pytest + +from switchyard.cli.route_bundle import ( + RouteBundle, + RouteBundleConfigError, + build_table_from_bundle, +) +from switchyard.lib.profiles.subagent_override import SubagentOverrideRuntime +from switchyard.lib.request_metadata import CTX_PROFILE_REQUEST_HEADERS +from switchyard_rust.core import is_subagent_request + + +class _RecordingRuntime: + """A ChainRuntime stand-in that records the calls routed to it.""" + + def __init__(self, name: str) -> None: + self.name = name + self.calls = 0 + + async def call(self, request: Any, *, ctx: Any = None) -> str: + self.calls += 1 + return self.name + + def iter_components(self) -> list[Any]: + return [self.name] + + +class _Ctx: + """Minimal ProxyContext stand-in exposing only the retained header map.""" + + def __init__(self, headers: dict[str, str] | None = None) -> None: + self.metadata: dict[str, Any] = {} + if headers is not None: + self.metadata[CTX_PROFILE_REQUEST_HEADERS] = headers + + +# --- detection policy ----------------------------------------------------------------- + + +def test_detection_matches_the_protocol_policy() -> None: + # Claude Code child-agent lineage is delegated work. + assert is_subagent_request( + {"x-claude-code-session-id": "root", "x-claude-code-agent-id": "child-1"} + ) + # Codex delegated-work kinds. + assert is_subagent_request({"x-openai-subagent": "review"}) + # Harness maintenance is sub-agent lineage but NOT delegated work. + assert not is_subagent_request({"x-openai-subagent": "compact"}) + # Explicit opt-out wins. + assert not is_subagent_request({"x-switchyard-is-subagent": "false"}) + # Ordinary traffic. + assert not is_subagent_request({}) + + +# --- runtime wrapper ------------------------------------------------------------------ + + +async def test_delegated_work_goes_to_the_worker() -> None: + inner, worker = _RecordingRuntime("inner"), _RecordingRuntime("worker") + runtime = SubagentOverrideRuntime(inner, worker) + + served = await runtime.call( + object(), + ctx=_Ctx({"x-claude-code-session-id": "s", "x-claude-code-agent-id": "child-1"}), + ) + assert served == "worker" + assert (inner.calls, worker.calls) == (0, 1) + + +async def test_normal_traffic_runs_the_wrapped_runtime() -> None: + inner, worker = _RecordingRuntime("inner"), _RecordingRuntime("worker") + runtime = SubagentOverrideRuntime(inner, worker) + + assert await runtime.call(object(), ctx=_Ctx({})) == "inner" + # A maintenance turn carries lineage but is not delegated work. + assert await runtime.call(object(), ctx=_Ctx({"x-openai-subagent": "compact"})) == "inner" + assert (inner.calls, worker.calls) == (2, 0) + + +async def test_missing_context_routes_normally() -> None: + inner, worker = _RecordingRuntime("inner"), _RecordingRuntime("worker") + runtime = SubagentOverrideRuntime(inner, worker) + + # No context means no retained headers, so nothing identifies the request. + assert await runtime.call(object(), ctx=None) == "inner" + assert worker.calls == 0 + + +def test_components_span_both_branches() -> None: + runtime = SubagentOverrideRuntime(_RecordingRuntime("inner"), _RecordingRuntime("worker")) + assert runtime.iter_components() == ["inner", "worker"] + + +# --- route bundle wiring -------------------------------------------------------------- + + +def _bundle(route: dict[str, Any]) -> RouteBundle: + return RouteBundle( + routes={"my-route": route}, + defaults={"api_key": "sk-test", "base_url": "https://example.invalid/v1"}, + ) + + +def test_subagent_target_wraps_any_route_type() -> None: + table = build_table_from_bundle( + _bundle({ + "type": "model", + "target": {"model": "strong-model"}, + "subagent_target": {"model": "worker-model"}, + }) + ) + assert isinstance(table.lookup_switchyard("my-route"), SubagentOverrideRuntime) + + +def test_routes_without_the_key_are_unwrapped() -> None: + table = build_table_from_bundle(_bundle({"type": "model", "target": {"model": "strong-model"}})) + assert not isinstance(table.lookup_switchyard("my-route"), SubagentOverrideRuntime) + + +def test_unknown_route_keys_are_still_rejected() -> None: + with pytest.raises(RouteBundleConfigError, match="subagent_targets"): + build_table_from_bundle( + _bundle({ + "type": "model", + "target": {"model": "strong-model"}, + "subagent_targets": {"model": "typo"}, + }) + )