diff --git a/Cargo.lock b/Cargo.lock index c1664fc83..c1bc969d2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3440,6 +3440,7 @@ dependencies = [ "reqwest", "serde", "serde_json", + "strum", "strum_macros", "switchyard-libsy", "switchyard-llm-client", diff --git a/Cargo.toml b/Cargo.toml index 2c788b727..b13b91ab4 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -53,6 +53,7 @@ switchyard-protocol = { path = "crates/protocol", version = "0.3.0" } switchyard-runner = { path = "crates/switchyard-runner", version = "0.3.0" } switchyard-server = { path = "crates/switchyard-server", version = "0.3.0" } switchyard-translation = { path = "crates/switchyard-translation", version = "0.3.0" } +strum = "0.28" strum_macros = "0.28" thiserror = "2" tokio = { version = "1", features = ["full"] } diff --git a/crates/libsy/src/core/algorithm.rs b/crates/libsy/src/core/algorithm.rs index 6e45c1267..2b7d79230 100644 --- a/crates/libsy/src/core/algorithm.rs +++ b/crates/libsy/src/core/algorithm.rs @@ -50,6 +50,7 @@ pub struct RuntimeModels { /// When using subagents this is the parent agent category. by_category: HashMap>, subagent: Option>>, + has_target_restriction: bool, } impl RuntimeModels { @@ -58,6 +59,7 @@ impl RuntimeModels { Self { by_category, subagent: None, + has_target_restriction: false, } } @@ -67,6 +69,33 @@ impl RuntimeModels { self } + /// Rejects routing calls and completion candidates outside the configured model groups. + /// Enforced by the driver and [`drive`]; hosts bind IDs to permitted deployments. + pub fn with_target_restriction(mut self) -> Self { + self.has_target_restriction = true; + self + } + + fn check_targets(&self, targets: &[ModelId]) -> Result<()> { + if !self.has_target_restriction { + return Ok(()); + } + // Include delegated groups: the parent driver publishes subagent outcomes too. + for target in targets { + if !self + .by_category + .values() + .chain(self.subagent.iter().flat_map(HashMap::values)) + .any(|models| models.contains(target)) + { + return Err(LibsyError::TargetNotFound { + target: target.clone(), + }); + } + } + Ok(()) + } + /// The models in `category`, ordered best-first. pub fn models_for(&self, category: &Category) -> &[ModelId] { self.by_category.get(category).map_or(&[], Vec::as_slice) @@ -401,6 +430,7 @@ impl Driver { let Some(selected_model_id) = models.first() else { return Err(LibsyError::NoTargets); }; + self.models.check_targets(&models)?; request.llm_request.model = Some(selected_model_id.to_string()); let started = Instant::now(); let (reply, response) = oneshot::channel::>(); @@ -445,6 +475,7 @@ impl Driver { mut request: DecisionRequest, model: ModelId, ) -> Result { + self.models.check_targets(std::slice::from_ref(&model))?; request.model = Some(model.clone()); let started = Instant::now(); let (reply, response) = oneshot::channel(); @@ -563,7 +594,7 @@ where F: Fn(Call) -> Fut, Fut: Future>, { - let stream = algorithm.run_stream(request, models); + let stream = algorithm.run_stream(request, Arc::clone(&models)); tokio::pin!(stream); let mut in_flight = futures::stream::FuturesUnordered::new(); @@ -579,11 +610,18 @@ where match step { None => break, // stream has ended, no more steps Some(item) => match item? { - Step::CallModel(call) => in_flight.push(serve(Call::Model(call))), + Step::CallModel(call) => { + models.check_targets(&call.models)?; + in_flight.push(serve(Call::Model(call))); + } Step::CallDecision(call) => { + models.check_targets(std::slice::from_ref(&call.model))?; in_flight.push(serve(Call::Decision(call))); } Step::Done(outcome) => { + if outcome.response.is_none() { + models.check_targets(&outcome.selected_model_ids)?; + } final_outcome = Some(*outcome); break; } @@ -717,15 +755,19 @@ pub trait Algorithm: Send + Sync + 'static { let algorithm = self.name().to_string(); // Catch a panicking algorithm so the run still publishes a terminal step. let route = AssertUnwindSafe(self.route(driver.clone(), request)).catch_unwind(); - let result = observability::observe_run(&algorithm, async move { - route.await.unwrap_or_else(|payload| { + let result = observability::observe_run(&algorithm, async { + let outcome = route.await.unwrap_or_else(|payload| { Err(LibsyError::AlgorithmError { message: format!( "algorithm task panicked: {}", panic_message(payload.as_ref()) ), }) - }) + })?; + if outcome.response.is_none() { + driver.models.check_targets(&outcome.selected_model_ids)?; + } + Ok(outcome) }) .await; @@ -851,6 +893,188 @@ mod tests { names.iter().map(|name| ModelId::from(*name)).collect() } + enum TargetProbe { + Model(Vec), + Decision(ModelId), + Completion(Vec), + Subagent, + } + + #[async_trait] + impl Algorithm for TargetProbe { + fn name(&self) -> &str { + "target_probe" + } + + async fn route( + self: Arc, + driver: Driver, + request: Request, + ) -> Result { + match &*self { + Self::Model(models) => { + let response = driver.call_model(request.clone(), models.clone()).await?; + Ok(RoutingOutcome::answered( + models[0].clone(), + request, + response, + )) + } + Self::Decision(model) => { + driver + .call_decision( + DecisionRequest { + model: None, + context: Value::Null, + questions: Default::default(), + }, + model.clone(), + ) + .await?; + Ok(RoutingOutcome::route_to("allowed".into(), vec![], request)) + } + Self::Completion(models) => Ok(RoutingOutcome::route_to( + models[0].clone(), + models[1..].to_vec(), + request, + )), + Self::Subagent => Ok(RoutingOutcome::route_to( + driver + .for_subagent()? + .first_model_for(&Category::Any)? + .clone(), + vec![], + request, + )), + } + } + } + + struct RewrittenSteps(Arc); + + #[async_trait] + impl Algorithm for RewrittenSteps { + fn name(&self) -> &str { + "rewritten_steps" + } + + async fn route( + self: Arc, + driver: Driver, + request: Request, + ) -> Result { + Arc::clone(&self.0).route(driver, request).await + } + + fn run_stream(self: Arc, request: Request, models: Arc) -> StepStream { + Box::pin(Arc::clone(&self.0).run_stream(request, models).map(|step| { + step.map(|mut step| { + match &mut step { + Step::CallModel(call) => call.models.push("outside".into()), + Step::CallDecision(call) => call.model = "outside".into(), + Step::Done(outcome) => outcome.selected_model_ids.push("outside".into()), + } + step + }) + })) + } + } + + fn restricted_models() -> RuntimeModels { + RuntimeModels::new(HashMap::from([ + (Category::Any, target_set(&["allowed", "fallback"])), + (Category::Judge, target_set(&["judge"])), + ])) + .with_target_restriction() + .with_subagent(HashMap::from([(Category::Any, target_set(&["subagent"]))])) + } + + // No forbidden call or completion candidate reaches the host. + #[tokio::test] + async fn target_restriction_blocks_undeclared_work() { + for probe in [ + TargetProbe::Model(target_set(&["outside"])), + TargetProbe::Model(target_set(&["allowed", "outside"])), + TargetProbe::Decision("outside".into()), + TargetProbe::Completion(target_set(&["outside"])), + TargetProbe::Completion(target_set(&["allowed", "outside"])), + ] { + let mut stream = Arc::new(probe).run_stream(request(), Arc::new(restricted_models())); + assert!(matches!( + stream.next().await, + Some(Err(LibsyError::TargetNotFound { target })) if target == "outside" + )); + assert!(stream.next().await.is_none()); + } + // A stream override can rewrite valid steps after the default driver checked them. + for probe in [ + TargetProbe::Model(target_set(&["allowed"])), + TargetProbe::Decision("judge".into()), + TargetProbe::Completion(target_set(&["allowed"])), + ] { + let result = drive( + Arc::new(RewrittenSteps(Arc::new(probe))), + request(), + Arc::new(restricted_models()), + |_| async { panic!("undeclared work must not reach the host") }, + ) + .await; + assert!( + matches!(result, Err(LibsyError::TargetNotFound { target }) if target == "outside") + ); + } + } + + // Judges, fallbacks, delegated targets and local answers remain valid. + #[tokio::test] + async fn target_restriction_preserves_allowed_work() -> Result<()> { + for (probe, expected) in [ + ( + TargetProbe::Model(target_set(&["allowed", "fallback"])), + &["allowed"][..], + ), + (TargetProbe::Model(target_set(&["judge"])), &["judge"][..]), + (TargetProbe::Decision("judge".into()), &["allowed"][..]), + ( + TargetProbe::Completion(target_set(&["allowed", "fallback"])), + &["allowed", "fallback"][..], + ), + (TargetProbe::Subagent, &["subagent"][..]), + ] { + let outcome = drive( + Arc::new(probe), + request(), + Arc::new(restricted_models()), + |call| async move { + match call { + Call::Model(call) => { + call.respond(std::future::ready(Ok(reply("ok")))).await + } + Call::Decision(call) => serve_decision(*call).await, + } + }, + ) + .await?; + assert_eq!(outcome.selected_model_ids, target_set(expected)); + } + let outcome = drive( + Arc::new(crate::Noop {}), + request(), + Arc::new(RuntimeModels::default().with_target_restriction()), + |_| async { panic!("Noop must not request model calls") }, + ) + .await?; + assert!(outcome.response.is_some()); + Ok(()) + } + + #[tokio::test] + async fn target_restriction_is_opt_in() { + let mut stream = orch(target_set(&["outside"])) + .run_stream(request(), Arc::new(RuntimeModels::default())); + assert!(matches!(stream.next().await, Some(Ok(Step::CallModel(_))))); + } + #[tokio::test] async fn typed_driver_preserves_call_and_stream_boundaries() -> Result<()> { tokio::time::timeout(std::time::Duration::from_secs(1), async { diff --git a/crates/switchyard-runner/Cargo.toml b/crates/switchyard-runner/Cargo.toml index 3a80b61b3..a330157f5 100644 --- a/crates/switchyard-runner/Cargo.toml +++ b/crates/switchyard-runner/Cargo.toml @@ -24,6 +24,7 @@ serde.workspace = true serde_json.workspace = true switchyard-llm-client.workspace = true switchyard-protocol.workspace = true +strum.workspace = true strum_macros.workspace = true thiserror.workspace = true toml = "1.1" diff --git a/crates/switchyard-runner/src/algorithm.rs b/crates/switchyard-runner/src/algorithm.rs index 646b37a43..e2bad61e2 100644 --- a/crates/switchyard-runner/src/algorithm.rs +++ b/crates/switchyard-runner/src/algorithm.rs @@ -195,6 +195,12 @@ struct CustomClassifierRouteConfig { pub struct CategoryModelConfig(BTreeMap>); impl CategoryModelConfig { + fn restrict_to(&mut self, allowed: &BTreeSet<&str>) { + for models in self.0.values_mut() { + models.retain(|model| allowed.contains(model.as_str())); + } + } + fn validate(&self, route_name: &str, default_target: &Category) -> AlgorithmResult<()> { for category in [Category::Any, Category::Judge] { if self.get(&category).is_empty() { @@ -641,6 +647,81 @@ impl AlgorithmSpec { } } + /// Filters configured pools without replacing required targets or changing their roles. + pub(crate) fn restricted_to(&self, allowed: &BTreeSet<&str>) -> AlgorithmResult { + let mut restricted = self.clone(); + match &mut restricted { + Self::Random { + targets, weights, .. + } => { + if let Some(weights) = weights { + if weights.len() != targets.len() { + return Err(AlgorithmConfigError::new(format!( + "random privacy expected {} weights, got {}", + targets.len(), + weights.len() + ))); + } + // Keep weight positions paired with the original target order. + let mut index = 0; + weights.retain(|_| { + let keep = allowed.contains(targets[index].as_str()); + index += 1; + keep + }); + } + targets.retain(|target| allowed.contains(target.as_str())); + if targets.is_empty() { + return Err(AlgorithmConfigError::new( + "privacy leaves random without a private target".to_string(), + )); + } + } + Self::LlmClassifier { config } => { + if let Some(models) = &mut config.models { + models.restrict_to(allowed); + } + } + _ => {} + } + if let Some((_, decision)) = restricted.decision_judge() + && let Some(target) = decision + .candidates + .values() + .find(|target| !allowed.contains(target.as_str())) + { + return Err(AlgorithmConfigError::new(format!( + "privacy decision candidate {target} must be private" + ))); + } + let subagents = match &mut restricted { + Self::Passthrough { subagents, .. } + | Self::StageRouter { subagents, .. } + | Self::Composite { subagents, .. } => subagents.as_mut(), + _ => None, + }; + if let Some(SubagentRouteConfig::LlmClassifier(config)) = subagents + && let Some(models) = &mut config.models + { + models.restrict_to(allowed); + } + if let Some(target) = restricted + .callable_target_names() + .into_iter() + .find(|target| !allowed.contains(target)) + { + return Err(AlgorithmConfigError::new(format!( + "privacy requires target {target} to be private" + ))); + } + Ok(restricted) + } + + /// Whether the algorithm supports a restricted set of configured targets. + pub(crate) const fn supports_privacy_lanes(&self) -> bool { + !matches!(self, Self::PrefillRouter { .. }) + } + pub(crate) fn decision_judge(&self) -> Option<(&str, &DecisionJudgeRouteConfig)> { match self { Self::LlmClassifier { config, .. } => config @@ -1687,3 +1768,149 @@ fn resolve_runtime_models( }) .collect() } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn restricted_random_preserves_weight_positions() -> AlgorithmResult<()> { + let allowed = BTreeSet::from(["private_a", "private_b"]); + let mut algorithm = AlgorithmSpec::Random { + targets: ["public_a", "private_a", "public_b", "private_b"] + .map(String::from) + .to_vec(), + weights: Some(vec![1.0, 2.0, 3.0, 4.0]), + seed: Some(42), + }; + let AlgorithmSpec::Random { + targets, + weights, + seed, + } = algorithm.restricted_to(&allowed)? + else { + unreachable!(); + }; + assert_eq!(targets, ["private_a", "private_b"]); + assert_eq!(weights, Some(vec![2.0, 4.0])); + assert_eq!(seed, Some(42)); + if let AlgorithmSpec::Random { weights, .. } = &mut algorithm { + *weights = Some(vec![1.0]); + } + assert!(algorithm.restricted_to(&allowed).is_err()); + Ok(()) + } + + fn custom_classifier() -> LlmClassifierRouteConfig { + LlmClassifierRouteConfig { + mode: Some(ClassifierMode::Custom), + models: Some(CategoryModelConfig(BTreeMap::from([ + ("any".into(), vec!["public".into(), "private".into()]), + ("efficient".into(), vec!["public".into(), "private".into()]), + ("judge".into(), vec!["judge".into()]), + ("research".into(), vec!["public".into()]), + ]))), + default_target: Some("efficient".into()), + prompt: Some("Select an available target.".into()), + response_schema: Some(r#"{"type":"object"}"#.into()), + max_output_tokens: default_classifier_max_output_tokens(), + policy: Some(ClassifierPolicyConfig::TargetSelector { + selector: "/target".into(), + }), + ..LlmClassifierRouteConfig::default() + } + } + + #[test] + fn restricted_custom_pools_cover_parent_and_subagents() -> AlgorithmResult<()> { + let allowed = BTreeSet::from(["private", "judge"]); + let targets = BTreeMap::from([ + ("private".into(), ModelId::from("private")), + ("judge".into(), ModelId::from("judge")), + ]); + let parent = AlgorithmSpec::LlmClassifier { + config: custom_classifier(), + }; + let child = AlgorithmSpec::Passthrough { + target: "private".into(), + subagents: Some(SubagentRouteConfig::LlmClassifier(Box::new( + custom_classifier(), + ))), + }; + for algorithm in [parent, child] { + let restricted = algorithm.restricted_to(&allowed)?; + let groups = restricted.runtime_model_names("test")?; + let groups = groups.subagent.unwrap_or(groups.parent); + assert_eq!(groups[&Category::Any], ["private"]); + assert_eq!(groups[&Category::Efficient], ["private"]); + assert_eq!(groups[&Category::Judge], ["judge"]); + restricted.build("test", &targets)?; + let missing_judge = algorithm.restricted_to(&BTreeSet::from(["private"]))?; + assert!(missing_judge.build("test", &targets).is_err()); + } + Ok(()) + } + + #[test] + fn restricted_roles_and_decision_candidates_are_not_replaced() { + let stage = AlgorithmSpec::Auto { + capable_target: "private".into(), + efficient_target: "public".into(), + }; + assert!(stage.restricted_to(&BTreeSet::from(["private"])).is_err()); + let decision = AlgorithmSpec::LlmClassifier { + config: LlmClassifierRouteConfig { + classifier_target: "judge".into(), + strong_target: Some("strong".into()), + weak_target: Some("weak".into()), + max_output_tokens: default_classifier_max_output_tokens(), + decision: Some(DecisionJudgeRouteConfig { + cutoff: 0.5, + instructions: None, + candidates: BTreeMap::from([ + ("strong".into(), "strong".into()), + ("weak".into(), "weak".into()), + ("extra".into(), "public".into()), + ]), + evidence: Value::Null, + }), + ..LlmClassifierRouteConfig::default() + }, + }; + let judge = decision.decision_judge().unwrap().1; + let classifier = serde_json::json!({ + "target": "judge", + "decision": { + "cutoff": judge.cutoff, "candidates": judge.candidates, + "evidence": judge.evidence, + }, + }); + let tiers = serde_json::json!({ + "capable_target": "strong", + "efficient_target": "weak", + "confidence_threshold": 0.6, + }); + let composite = serde_json::from_value(serde_json::json!({ + "type": "composite", "classifier": classifier, "stage": tiers, + })) + .unwrap(); + let stage = serde_json::from_value(serde_json::json!({ + "type": "stage_router", "classifier": classifier, + "capable_target": "strong", "efficient_target": "weak", + "confidence_threshold": 0.6, "picker": "efficient_first", + })) + .unwrap(); + for decision in [decision, composite, stage] { + assert!(decision.runtime_model_names("test").is_ok()); + let error = decision + .restricted_to(&BTreeSet::from(["strong", "weak", "judge"])) + .expect_err( + "an extra public candidate must not be restored during lane construction", + ); + assert_eq!( + error.to_string(), + "privacy decision candidate public must be private" + ); + } + } +} diff --git a/crates/switchyard-runner/src/config.rs b/crates/switchyard-runner/src/config.rs index e8918efe0..c559a5318 100644 --- a/crates/switchyard-runner/src/config.rs +++ b/crates/switchyard-runner/src/config.rs @@ -18,6 +18,8 @@ use switchyard_llm_client::{ }; use switchyard_protocol::{ModelId, RoutedDecisionClient, RoutedLlmClient, WireFormat}; +use crate::privacy::PrivacyPolicy; +use crate::route::ExecutionLane; use crate::{ AlgorithmSpec, AuxiliaryTarget, CallerAuthKind, DecisionTarget, ModelCapabilities, Route, Runner, RunnerError, @@ -94,14 +96,41 @@ struct RouteConfig { tool_calling: Option, reasoning: Option, vision: Option, + privacy: Option, algorithm: AlgorithmSpec, } +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct PrivacyConfig { + #[serde(default)] + accept_external_signal: bool, +} + struct TargetPromptPolicy { prompts: HashMap, routing_answer_target: Option, } +#[derive(Clone, Copy)] +struct ResolvedTarget<'a> { + target_name: &'a str, + config: &'a TargetConfig, +} + +type LaneTargets<'a> = BTreeMap<&'a str, ResolvedTarget<'a>>; + +#[derive(Clone, Copy)] +struct ResolvedDecisionTarget<'a> { + model: &'a ModelId, + client: &'a Arc, +} + +struct BuiltPrivacy { + policy: PrivacyPolicy, + restricted: ExecutionLane, +} + impl<'de> Deserialize<'de> for RouteConfig { fn deserialize(deserializer: D) -> Result where @@ -113,6 +142,7 @@ impl<'de> Deserialize<'de> for RouteConfig { let tool_calling = take_optional(&mut table, "tool_calling")?; let reasoning = take_optional(&mut table, "reasoning")?; let vision = take_optional(&mut table, "vision")?; + let privacy = take_optional(&mut table, "privacy")?; let algorithm = AlgorithmSpec::deserialize(toml::Value::Table(table)) .map_err(serde::de::Error::custom)?; Ok(Self { @@ -121,6 +151,7 @@ impl<'de> Deserialize<'de> for RouteConfig { tool_calling, reasoning, vision, + privacy, algorithm, }) } @@ -155,14 +186,6 @@ impl RouteConfig { vision: self.vision, } } - - fn routing_target_names(&self) -> Vec<&str> { - self.algorithm.routing_target_names() - } - - fn callable_target_names(&self) -> Vec<&str> { - self.algorithm.callable_target_names() - } } impl DeploymentConfig { @@ -180,7 +203,7 @@ impl DeploymentConfig { "route {name}: model selection requires auto, passthrough, or capability llm_classifier with every_request" ))); } - for target_name in route.routing_target_names() { + for target_name in route.algorithm.routing_target_names() { if let Some(target) = self.targets.get(target_name) && (target.system_prompt.is_some() || target.reasoning_effort.is_some() @@ -245,6 +268,11 @@ impl DeploymentConfig { } std::collections::hash_map::Entry::Occupied(slot) => { let (first_name, first) = slot.get(); + if first.is_private != target.is_private { + return Err(RunnerError::configuration(format!( + "targets {first_name} and {target_name} share a model and llm client but disagree on private" + ))); + } if first.reasoning_effort != target.reasoning_effort || first.extra_body != target.extra_body || first.omit_body_fields != target.omit_body_fields @@ -266,6 +294,7 @@ impl DeploymentConfig { let mut provider_api_keys = Vec::new(); let clients = self.build_clients(&mut provider_api_keys)?; let decision_clients = self.build_decision_clients(&mut provider_api_keys)?; + let mut decision_approvals = HashMap::new(); for (name, target) in &self.decision_targets { validate_value("decision target name", name)?; validate_value(&format!("decision target {name} id"), &target.id)?; @@ -280,66 +309,44 @@ impl DeploymentConfig { target.decision_client ))); } + if decision_approvals + .insert((&target.decision_client, &target.id), target.is_private) + .is_some_and(|approved| approved != target.is_private) + { + return Err(RunnerError::configuration(format!( + "decision target {name} shares a model and client with an alias but disagrees on private" + ))); + } } - let targets = self.build_targets(); let fallback_base_url = self.fallback_base_url()?; let mut routes = Vec::with_capacity(self.routes.len()); for (route_name, config) in &self.routes { - let decision_judge = config.algorithm.decision_judge(); - for name in config.callable_target_names() { - let (exists, kind) = if decision_judge.is_some_and(|(judge, _)| judge == name) { - (self.decision_targets.contains_key(name), "decision") - } else { - (self.targets.contains_key(name), "LLM") - }; - if !exists { - return Err(RunnerError::configuration(format!( - "route references unknown target {name}; route {route_name} requires target kind {kind}" - ))); - } - } - for name in config.routing_target_names().into_iter().chain( - decision_judge - .into_iter() - .flat_map(|(_, judge)| judge.candidates.values().map(String::as_str)), - ) { - if !self.targets.contains_key(name) { - return Err(RunnerError::configuration(format!( - "route {route_name} completion and candidate target {name} must be an LLM target" - ))); - } - } + let standard_targets = self.resolve_lane_targets(&config.algorithm)?; + let decision = config + .algorithm + .decision_judge() + .map(|(name, _)| self.resolve_decision_target(route_name, name, &decision_clients)) + .transpose()?; let capabilities = config.capabilities(); if capabilities.context_window == Some(0) { return Err(RunnerError::configuration(format!( "route {route_name} context_window must be greater than zero" ))); } - let (algorithm, models) = config - .algorithm - .build_with_runtime_models(route_name, &targets) - .map_err(|error| RunnerError::configuration_source(error.to_string(), error))?; - let (route_clients, caller_auth) = - self.build_route_clients(route_name, config, &clients, &decision_clients)?; - let anthropic_auxiliary_target = - self.build_anthropic_auxiliary_target(config, &clients); - let responses_auxiliary_target = - self.build_responses_auxiliary_target(config, &clients); - let decision_targets = config - .routing_target_names() - .into_iter() - .filter_map(|name| self.decision_target(name)) - .collect(); - let route = Route::new( - algorithm, - route_clients, - caller_auth, - capabilities, - anthropic_auxiliary_target, - responses_auxiliary_target, - decision_targets, - models, - ); + let privacy = + self.build_privacy(route_name, config, &standard_targets, decision, &clients)?; + let (standard, caller_auth) = self.build_lane( + route_name, + &config.algorithm, + &standard_targets, + decision, + &clients, + config.privacy.is_some(), + )?; + let mut route = Route::from_lane(standard, caller_auth, capabilities); + if let Some(privacy) = privacy { + route = route.with_privacy(privacy.policy, privacy.restricted); + } routes.push((config.id.clone(), route)); } let runner = Runner::new(routes) @@ -348,6 +355,180 @@ impl DeploymentConfig { Ok(runner) } + fn build_privacy( + &self, + route_name: &str, + route_config: &RouteConfig, + standard_targets: &LaneTargets<'_>, + standard_decision_target: Option>, + llm_clients: &BTreeMap>, + ) -> RunnerResult> { + let Some(config) = &route_config.privacy else { + return Ok(None); + }; + if !route_config.algorithm.supports_privacy_lanes() { + return Err(RunnerError::configuration(format!( + "route {route_name} cannot use privacy with prefill_router" + ))); + } + if !config.accept_external_signal { + return Err(RunnerError::configuration(format!( + "route {route_name} privacy must configure at least one request input" + ))); + } + let approved = self + .targets + .iter() + .filter(|(_, target)| target.is_private) + .map(|(name, _)| name.as_str()) + .chain( + self.decision_targets + .iter() + .filter(|(_, target)| target.is_private) + .map(|(name, _)| name.as_str()), + ) + .collect(); + let restricted_algorithm = route_config + .algorithm + .restricted_to(&approved) + .map_err(|error| RunnerError::configuration(format!("route {route_name}: {error}")))?; + let restricted_targets = self.resolve_lane_targets(&restricted_algorithm)?; + if self.uses_forward_auth(standard_targets) { + return Err(RunnerError::configuration(format!( + "route {route_name} cannot use privacy with forward_auth" + ))); + } + let policy = PrivacyPolicy::new(config.accept_external_signal); + let (restricted, _) = self.build_lane( + route_name, + &restricted_algorithm, + &restricted_targets, + standard_decision_target, + llm_clients, + true, + )?; + Ok(Some(BuiltPrivacy { policy, restricted })) + } + + fn uses_forward_auth(&self, targets: &LaneTargets<'_>) -> bool { + targets.values().any(|target| { + self.llm_clients + .get(&target.config.llm_client) + .is_some_and(|client| client.forward_auth) + }) + } + + fn resolve_lane_targets<'a>( + &'a self, + algorithm: &'a AlgorithmSpec, + ) -> RunnerResult> { + let decision = algorithm.decision_judge().map(|(target, _)| target); + let mut callable = algorithm.callable_target_names(); + callable.retain(|name| Some(*name) != decision); + callable.sort_unstable(); + callable.dedup(); + callable + .into_iter() + .map(|name| { + let config = self.targets.get(name).ok_or_else(|| { + RunnerError::configuration(format!("route references unknown target {name}")) + })?; + Ok(( + name, + ResolvedTarget { + target_name: name, + config, + }, + )) + }) + .collect() + } + + fn resolve_decision_target<'a>( + &'a self, + route_name: &str, + name: &str, + clients: &'a BTreeMap>, + ) -> RunnerResult> { + let target = self.decision_targets.get(name).ok_or_else(|| { + RunnerError::configuration(format!( + "route {route_name} requires unknown decision target {name}" + )) + })?; + let client = clients.get(&target.decision_client).ok_or_else(|| { + RunnerError::configuration(format!( + "decision target {name} references unknown decision client {}", + target.decision_client + )) + })?; + Ok(ResolvedDecisionTarget { + model: &target.id, + client, + }) + } + + fn build_lane( + &self, + route_name: &str, + route: &AlgorithmSpec, + targets: &LaneTargets<'_>, + decision: Option>, + clients: &BTreeMap>, + restrict_targets: bool, + ) -> RunnerResult<(ExecutionLane, Option)> { + let mut model_ids = targets + .iter() + .map(|(name, target)| ((*name).to_string(), target.config.id.clone())) + .collect::>(); + if let Some((_, config)) = route.decision_judge() { + for candidate in config.candidates.values() { + if !model_ids.contains_key(candidate) { + let target = self.targets.get(candidate).ok_or_else(|| { + RunnerError::configuration(format!( + "route {route_name} completion and candidate target {candidate} must be an LLM target" + )) + })?; + model_ids.insert(candidate.clone(), target.id.clone()); + } + } + } + if let (Some((name, _)), Some(decision)) = (route.decision_judge(), decision) { + model_ids.insert(name.to_string(), decision.model.clone()); + } + let (algorithm, mut models) = route + .build_with_runtime_models(route_name, &model_ids) + .map_err(|error| RunnerError::configuration_source(error.to_string(), error))?; + let (route_clients, caller_auth) = + self.build_route_clients(route_name, route, targets, decision, clients)?; + let anthropic_auxiliary_target = + self.build_anthropic_auxiliary_target(route, targets, clients); + let responses_auxiliary_target = + self.build_responses_auxiliary_target(route, targets, clients); + let decision_targets = route + .routing_target_names() + .into_iter() + .filter_map(|name| { + targets + .get(name) + .and_then(|target| self.decision_target(target.target_name)) + }) + .collect(); + if restrict_targets { + models = models.with_target_restriction(); + } + Ok(( + ExecutionLane::new( + algorithm, + route_clients, + anthropic_auxiliary_target, + responses_auxiliary_target, + decision_targets, + models, + ), + caller_auth, + )) + } + fn build_clients( &self, provider_api_keys: &mut Vec, @@ -456,69 +637,47 @@ impl DeploymentConfig { .collect() } - fn build_targets(&self) -> BTreeMap { - self.targets - .iter() - .map(|(name, config)| (name.clone(), config.id.clone())) - .chain( - self.decision_targets - .iter() - .map(|(name, config)| (name.clone(), config.id.clone())), - ) - .collect() - } - /// Builds the client router for one route. The second value is the caller credential family /// that the route's forwarding clients need, or `None` when no client forwards the caller's /// credential. A request through the other family's APIs fails before any upstream call. fn build_route_clients( &self, route_name: &str, - route: &RouteConfig, + route: &AlgorithmSpec, + targets: &LaneTargets<'_>, + decision: Option>, clients: &BTreeMap>, - decision_clients: &BTreeMap>, ) -> RunnerResult<(ClientRouter, Option)> { let TargetPromptPolicy { prompts, routing_answer_target, - } = self.build_route_target_prompts(route_name, route)?; + } = self.build_route_target_prompts(route_name, route, targets)?; let mut by_model = HashMap::new(); let mut targets_by_model: HashMap<&str, (&str, &TargetConfig)> = HashMap::new(); let mut caller_auth = None; let mut has_mixed_families = false; let mut forwarding_origins = BTreeSet::new(); - let mut decisions_by_model = HashMap::new(); - for name in route.callable_target_names() { - if route - .algorithm - .decision_judge() - .is_some_and(|(judge, _)| judge == name) - { - let target = &self.decision_targets[name]; - decisions_by_model.insert( - target.id.clone(), - decision_clients[&target.decision_client].clone(), - ); - continue; - } - let target = self.targets.get(name).ok_or_else(|| { - RunnerError::configuration(format!("route references unknown target {name}")) - })?; - if let Some((first_name, first)) = targets_by_model.insert(&target.id, (name, target)) + for resolved in targets.values() { + let target = resolved.config; + if let Some((first_name, first)) = + targets_by_model.insert(&target.id, (resolved.target_name, target)) && first.llm_client != target.llm_client { return Err(RunnerError::configuration(format!( - "route {route_name} targets {first_name} and {name} use model {} on different llm clients; execution is keyed by model id, so use distinct model ids within this route or put these targets in separate routes", - target.id + "route {route_name} targets {first_name} and {} use model {} on different llm clients; execution is keyed by model id, so use distinct model ids within this route or put these targets in separate routes", + resolved.target_name, target.id ))); } let client = clients.get(&target.llm_client).ok_or_else(|| { - RunnerError::configuration(format!("target {name} has no constructed llm client")) + RunnerError::configuration(format!( + "target {} has no constructed llm client", + resolved.target_name + )) })?; let client_config = self.llm_clients.get(&target.llm_client).ok_or_else(|| { RunnerError::configuration(format!( - "target {name} references unknown llm client {}", - target.llm_client + "target {} references unknown llm client {}", + resolved.target_name, target.llm_client )) })?; if client_config.forward_auth { @@ -547,8 +706,11 @@ impl DeploymentConfig { let completion_targets = route .routing_target_names() .into_iter() - .map(|name| self.targets[name].id.clone()) + .map(|name| targets[name].config.id.clone()) .collect::>(); + let decisions_by_model = decision + .map(|decision| HashMap::from([(decision.model.clone(), Arc::clone(decision.client))])) + .unwrap_or_default(); let router = ClientRouter::new_with_decision_clients( by_model, decisions_by_model, @@ -563,14 +725,18 @@ impl DeploymentConfig { fn build_route_target_prompts( &self, route_name: &str, - route: &RouteConfig, + route: &AlgorithmSpec, + targets: &LaneTargets<'_>, ) -> RunnerResult { let mut prompts = HashMap::new(); let mut aliases = HashMap::<&ModelId, Option<&str>>::new(); - for name in route.algorithm.routing_target_names() { - let target = self.targets.get(name).ok_or_else(|| { - RunnerError::configuration(format!("route references unknown target {name}")) - })?; + for name in route.routing_target_names() { + let target = targets + .get(name) + .map(|target| target.config) + .ok_or_else(|| { + RunnerError::configuration(format!("route references unknown target {name}")) + })?; let prompt = target.system_prompt.as_deref(); if aliases .insert(&target.id, prompt) @@ -589,20 +755,28 @@ impl DeploymentConfig { prompts, routing_answer_target: None, }; - let Some((response_name, dependency_name)) = - route.algorithm.routing_response_and_dependency() - else { + let Some((response_name, dependency_name)) = route.routing_response_and_dependency() else { return Ok(policy); }; - let response = self.targets.get(response_name).ok_or_else(|| { - RunnerError::configuration(format!("route references unknown target {response_name}")) - })?; + let response = targets + .get(response_name) + .map(|target| target.config) + .ok_or_else(|| { + RunnerError::configuration(format!( + "route references unknown target {response_name}" + )) + })?; if !policy.prompts.contains_key(&response.id) { return Ok(policy); } - let dependency = self.targets.get(dependency_name).ok_or_else(|| { - RunnerError::configuration(format!("route references unknown target {dependency_name}")) - })?; + let dependency = targets + .get(dependency_name) + .map(|target| target.config) + .ok_or_else(|| { + RunnerError::configuration(format!( + "route references unknown target {dependency_name}" + )) + })?; if response.id == dependency.id { return Err(RunnerError::configuration(format!( "route {route_name} cannot apply system_prompt to target {response_name}: model {} is also used by routing-only target {dependency_name}", @@ -627,7 +801,8 @@ impl DeploymentConfig { fn build_anthropic_auxiliary_target( &self, - route: &RouteConfig, + route: &AlgorithmSpec, + targets: &LaneTargets<'_>, clients: &BTreeMap>, ) -> Option { route @@ -635,12 +810,17 @@ impl DeploymentConfig { .into_iter() .enumerate() .filter_map(|(index, name)| { + let resolved = *targets.get(name)?; let target = self.build_auxiliary_target( - name, + resolved, clients, AuxiliaryOperation::AnthropicCountTokens, )?; - Some((count_tokens_priority(name, &target.model), index, target)) + Some(( + count_tokens_priority(resolved.target_name, &target.model), + index, + target, + )) }) .min_by_key(|(priority, index, _)| (*priority, *index)) .map(|(_, _, target)| target) @@ -648,26 +828,27 @@ impl DeploymentConfig { fn build_responses_auxiliary_target( &self, - route: &RouteConfig, + route: &AlgorithmSpec, + targets: &LaneTargets<'_>, clients: &BTreeMap>, ) -> Option { route.routing_target_names().into_iter().find_map(|name| { - self.build_auxiliary_target(name, clients, AuxiliaryOperation::ResponsesInputTokens) + let target = *targets.get(name)?; + self.build_auxiliary_target(target, clients, AuxiliaryOperation::ResponsesInputTokens) }) } fn build_auxiliary_target( &self, - name: &str, + target: ResolvedTarget<'_>, clients: &BTreeMap>, operation: AuxiliaryOperation, ) -> Option { - let target = self.targets.get(name)?; - let client = clients.get(&target.llm_client)?; + let client = clients.get(&target.config.llm_client)?; client - .supports_auxiliary(&target.id, operation) + .supports_auxiliary(&target.config.id, operation) .then(|| AuxiliaryTarget { - model: target.id.clone(), + model: target.config.id.clone(), client: client.clone(), }) } @@ -747,6 +928,8 @@ enum DecisionClientConfig { struct DecisionTargetConfig { id: ModelId, decision_client: String, + #[serde(default, rename = "private")] + is_private: bool, } #[derive(Debug, Deserialize)] @@ -754,6 +937,8 @@ struct DecisionTargetConfig { struct TargetConfig { id: ModelId, llm_client: String, + #[serde(default, rename = "private")] + is_private: bool, #[serde(default)] extra_body: BTreeMap, /// Top-level request fields dropped from the outbound body before `extra_body` is merged. @@ -935,8 +1120,9 @@ bogus = true #[cfg(test)] mod deployment_tests { use super::*; + use crate::mark_privacy_restricted; use serde_json::json; - use switchyard_protocol::Category; + use switchyard_protocol::{Category, Message, Request, Role}; const VALID_CONFIG: &str = r#" schema_version = 1 @@ -994,6 +1180,62 @@ capable_target = "strong" efficient_target = "weak" "#; + fn decision_privacy_config() -> String { + let config = VALID_CONFIG + .replace( + "id = \"strong/model\"", + "id = \"strong/model\"\nprivate = true", + ) + .replace("id = \"weak/model\"", "id = \"weak/model\"\nprivate = true"); + format!( + r#"{config} + +[decision_clients.privacy] +format = "system_one" +endpoint = "https://example.test/v1/systemone" +api_key_env = "PATH" +timeout_ms = 5000 + +[decision_targets.standard] +id = "decision/standard" +decision_client = "privacy" +private = true + +[routes.typed] +id = "switchyard/typed" +type = "llm_classifier" +mode = "capability" +classifier_target = "standard" +strong_target = "strong" +weak_target = "weak" + +[routes.typed.decision] +cutoff = 0.4 +candidates = {{ a = "strong", b = "weak" }} +evidence = {{}} + +[routes.typed.privacy] +accept_external_signal = true +"# + ) + } + + fn external_privacy_config() -> String { + let config = VALID_CONFIG + .replace("id = \"strong/model\"", "id = \"strong/model\"\nprivate = true") + .replace( + "type = \"passthrough\"\ntarget = \"weak\"", + "type = \"random\"\ntargets = [\"weak\", \"strong\"]\nweights = [3.0, 1.0]\nseed = 7", + ); + format!( + r#"{config} + +[routes.passthrough.privacy] +accept_external_signal = true +"# + ) + } + #[test] fn public_runner_from_toml_builds_a_deployment() -> RunnerResult<()> { let runner = Runner::from_toml(VALID_CONFIG)?; @@ -1033,6 +1275,97 @@ efficient_target = "weak" Ok(()) } + #[test] + fn privacy_requires_approved_roles_and_decision_targets() -> RunnerResult<()> { + let configured = decision_privacy_config(); + let runner = Runner::from_toml(&configured)?; + let route = runner + .route("switchyard/typed") + .expect("typed privacy route should exist"); + assert_eq!( + route.models().models_for(&Category::Judge), + [ModelId::from("decision/standard")] + ); + + for (config, expected) in [ + ( + configured.replace( + "decision_client = \"privacy\"\nprivate = true", + "decision_client = \"privacy\"", + ), + "privacy requires target standard to be private", + ), + ( + configured.replace("id = \"weak/model\"\nprivate = true", "id = \"weak/model\""), + "privacy decision candidate weak must be private", + ), + ] { + let message = error_message(&config); + assert!( + message.contains(expected), + "expected {expected:?}, got {message}" + ); + } + Ok(()) + } + + #[tokio::test] + async fn external_privacy_selects_the_restricted_lane() -> RunnerResult<()> { + let configured = external_privacy_config(); + let runner = Runner::from_toml(&configured)?; + let route = runner + .route("switchyard/passthrough") + .expect("privacy route should exist"); + + let mut request = Request::default(); + request.llm_request.messages.push(Message::text( + Role::User, + "Contact user@example.test. Authorization: Bearer abcdefghijklmnop", + )); + let outcome = route.decide(request.clone()).await?; + assert!( + route + .models() + .models_for(&Category::Any) + .contains(outcome.selected_model_id()?) + ); + assert_eq!( + crate::privacy::selected_lane(&outcome.request).map(|lane| lane.as_str()), + Some("standard") + ); + mark_privacy_restricted(&mut request); + assert_eq!( + route.decide(request).await?.selected_model_id()?, + "strong/model" + ); + + let empty = configured.replace( + "accept_external_signal = true", + "accept_external_signal = false", + ); + assert!(error_message(&empty).contains("must configure at least one request input")); + let obsolete = format!( + "{configured}\n[routes.passthrough.privacy.deterministic]\ndetectors = [\"email\"]\n" + ); + assert!(error_message(&obsolete).contains("unknown field `deterministic`")); + let unmapped = format!( + "{configured}\n[routes.passthrough.privacy.restricted_targets]\nweak = \"strong\"\n" + ); + assert!(error_message(&unmapped).contains("unknown field `restricted_targets`")); + let no_private_target = configured.replace("private = true", "private = false"); + assert!(error_message(&no_private_target).contains("without a private target")); + let missing_tier = + format!("{configured}\n[routes.plan_execute.privacy]\naccept_external_signal = true\n"); + assert!( + error_message(&missing_tier).contains("privacy requires target weak to be private") + ); + let conflicting_alias = format!( + "{configured}\n[targets.alias]\nid = \"strong/model\"\nllm_client = \"responses\"\n" + ); + assert!(error_message(&conflicting_alias).contains("disagree on private")); + Ok(()) + } + #[test] fn duplicate_route_ids_are_rejected() { let config = format!( @@ -1214,11 +1547,15 @@ evidence = { description = "candidate comparison" } (source.replace("b = \"weak\"", "b = \"strong\""), false), ] { let deployment: DeploymentConfig = toml::from_str(&source).unwrap(); + let targets = deployment + .targets + .iter() + .map(|(name, target)| (name.clone(), target.id.clone())) + .collect(); let spec = &deployment.routes["composed"].algorithm; assert_eq!(spec.decision_judge().unwrap().0, "tier_judge"); assert_eq!( - spec.build_with_runtime_models("composed", &deployment.build_targets()) - .is_ok(), + spec.build_with_runtime_models("composed", &targets).is_ok(), valid ); } @@ -1292,7 +1629,7 @@ evidence = { description = "candidate comparison" } let Some(stage_route) = parsed.routes.get("stage") else { return Err(RunnerError::configuration("stage route is missing")); }; - let callable_targets = stage_route.callable_target_names(); + let callable_targets = stage_route.algorithm.callable_target_names(); for expected in ["strong", "weak", "stage_judge", "classifier"] { assert!(callable_targets.contains(&expected)); } diff --git a/crates/switchyard-runner/src/lib.rs b/crates/switchyard-runner/src/lib.rs index 115e8ce83..d8ad1602b 100644 --- a/crates/switchyard-runner/src/lib.rs +++ b/crates/switchyard-runner/src/lib.rs @@ -6,6 +6,7 @@ mod algorithm; mod config; mod failure; +mod privacy; mod provider_key_redactor; mod route; mod runner; @@ -18,6 +19,7 @@ pub use algorithm::{ pub use failure::{RouteErrorKind, RouteErrorPhase, RouteErrorSummary, stream_error_summary}; // Re-exported because `Route::new` takes it, so a host wiring routes does not need a libsy dep. pub use libsy::RuntimeModels; +pub use privacy::mark_privacy_restricted; pub use provider_key_redactor::ProviderKeyRedactor; pub use route::{ AuxiliaryTarget, CallerAuthKind, ModelCapabilities, Route, RunOutput, RunnerError, diff --git a/crates/switchyard-runner/src/privacy.rs b/crates/switchyard-runner/src/privacy.rs new file mode 100644 index 000000000..8a4706ff0 --- /dev/null +++ b/crates/switchyard-runner/src/privacy.rs @@ -0,0 +1,270 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Privacy policy used to select a route's execution lane. + +use serde_json::Value; +use strum_macros::{EnumString, IntoStaticStr}; +use switchyard_protocol::{LlmClientError, Request, WireFormat}; + +const EXTERNAL_RESTRICTION_KEY: &str = "switchyard.internal.external_privacy_restriction"; +pub(crate) const SELECTED_LANE_KEY: &str = "switchyard.internal.privacy_lane"; +const RESPONSES_STATE_FIELDS: [&str; 2] = ["previous_response_id", "conversation"]; + +pub(crate) struct PrivacyPolicy { + accept_external_signal: bool, +} + +/// Target set allowed to serve one request. +#[derive(Clone, Copy, EnumString, IntoStaticStr)] +#[strum(serialize_all = "snake_case")] +pub(crate) enum PrivacyLane { + Standard, + Restricted, +} + +impl PrivacyLane { + pub(crate) fn as_str(self) -> &'static str { + self.into() + } +} + +/// Records the lane after routing so target metadata can resolve duplicate model IDs. +pub(crate) fn record_selected_lane(request: &mut Request, lane: PrivacyLane) { + request + .metadata + .get_or_insert_default() + .extra_metadata + .get_or_insert_default() + .insert(SELECTED_LANE_KEY.to_string(), lane.as_str().to_string()); +} + +/// Reads the runner-owned lane marker from a completed routing request. +pub(crate) fn selected_lane(request: &Request) -> Option { + request + .metadata + .as_ref() + .and_then(|metadata| metadata.extra_metadata.as_ref()) + .and_then(|metadata| metadata.get(SELECTED_LANE_KEY)) + .and_then(|lane| lane.parse().ok()) +} + +/// Bounded explanation of the privacy lane selected for one request. +pub(crate) struct PrivacyDecision { + lane: PrivacyLane, + source: PrivacySource, + reason_code: &'static str, +} + +#[derive(Clone, Copy, IntoStaticStr)] +#[strum(serialize_all = "snake_case")] +pub(crate) enum PrivacySource { + Policy, + ExternalSignal, +} + +impl PrivacySource { + pub(crate) fn as_str(self) -> &'static str { + self.into() + } +} + +impl PrivacyDecision { + /// Returns the selected execution lane. + pub(crate) fn lane(&self) -> &'static str { + self.lane.as_str() + } + + /// Returns the policy component that selected the lane. + pub(crate) fn source(&self) -> &'static str { + self.source.as_str() + } + + /// Returns a bounded reason for the decision. + pub(crate) const fn reason_code(&self) -> &'static str { + self.reason_code + } + + pub(crate) const fn selected_lane(&self) -> PrivacyLane { + self.lane + } + + const fn new(lane: PrivacyLane, source: PrivacySource, reason_code: &'static str) -> Self { + Self { + lane, + source, + reason_code, + } + } + + const fn all_clear() -> Self { + Self::new(PrivacyLane::Standard, PrivacySource::Policy, "all_clear") + } +} + +/// Marks a request as requiring a route that accepts external privacy signals. +/// Execution fails if the selected route has not enabled `accept_external_signal`. +/// This marker is available only to trusted in-process hosts, not HTTP clients. +pub fn mark_privacy_restricted(request: &mut Request) { + request + .metadata + .get_or_insert_default() + .extra_metadata + .get_or_insert_default() + .insert(EXTERNAL_RESTRICTION_KEY.to_string(), String::new()); +} + +pub(crate) fn has_external_restriction(request: &Request) -> bool { + request + .metadata + .as_ref() + .and_then(|metadata| metadata.extra_metadata.as_ref()) + .is_some_and(|metadata| metadata.contains_key(EXTERNAL_RESTRICTION_KEY)) +} + +pub(crate) fn external_signal_not_accepted() -> LlmClientError { + LlmClientError::InvalidRequest { + message: "route does not accept external privacy signals".to_string(), + } +} + +impl PrivacyPolicy { + pub(crate) const fn new(accept_external_signal: bool) -> Self { + Self { + accept_external_signal, + } + } + + pub(crate) async fn decide( + &self, + request: &Request, + ) -> Result { + if has_external_restriction(request) { + if !self.accept_external_signal { + return Err(external_signal_not_accepted()); + } + return Ok(PrivacyDecision::new( + PrivacyLane::Restricted, + PrivacySource::ExternalSignal, + "restricted", + )); + } + Ok(PrivacyDecision::all_clear()) + } +} + +pub(crate) fn validate_mixed_request(request: &Request) -> Result<(), LlmClientError> { + let extensions = &request.llm_request.extensions.fields; + let preserved = &request.llm_request.preservation.requests; + let inspect_extensions = match request + .metadata + .as_ref() + .and_then(|metadata| metadata.wire_format) + { + Some(format) => format == WireFormat::OpenAiResponses, + None => preserved.is_empty(), + }; + let has_state = inspect_extensions + && RESPONSES_STATE_FIELDS + .iter() + .any(|field| extensions.get(*field).is_some_and(|value| !value.is_null())) + || preserved.iter().any(|(format, body)| { + format.as_str() == WireFormat::OpenAiResponses.as_str() && has_responses_state(body) + }); + if has_state { + return Err(LlmClientError::InvalidRequest { + message: + "mixed-traffic privacy routes do not support OpenAI Responses continuation state" + .to_string(), + }); + } + Ok(()) +} + +fn has_responses_state(body: &Value) -> bool { + RESPONSES_STATE_FIELDS + .iter() + .any(|field| body.get(*field).is_some_and(|value| !value.is_null())) +} + +#[cfg(test)] +mod tests { + use super::*; + use switchyard_protocol::{ContentBlock, Message, Role}; + + #[tokio::test] + async fn external_restriction_requires_route_opt_in() { + let mut request = Request::default(); + mark_privacy_restricted(&mut request); + + assert!(PrivacyPolicy::new(false).decide(&request).await.is_err()); + assert!(matches!( + PrivacyPolicy::new(true).decide(&request).await, + Ok(PrivacyDecision { + lane: PrivacyLane::Restricted, + .. + }) + )); + + let mut opaque = Request::default(); + opaque.llm_request.messages.push(Message { + role: Role::User, + content: vec![ContentBlock::Unknown { + provider: "custom".into(), + raw: Value::Null, + }], + }); + let decision = PrivacyPolicy::new(true) + .decide(&opaque) + .await + .expect("unmarked request should remain valid"); + assert!(matches!(decision.lane, PrivacyLane::Standard)); + } + + #[test] + fn responses_continuations_are_rejected_from_normalized_and_preserved_state() { + for field in RESPONSES_STATE_FIELDS { + for preserved in [false, true] { + let mut request = Request::default(); + if preserved { + request.llm_request.preservation.requests.insert( + "openai_responses".into(), + serde_json::json!({field: "state_123"}), + ); + } else { + request.metadata.get_or_insert_default().wire_format = + Some(WireFormat::OpenAiResponses); + request + .llm_request + .extensions + .fields + .insert(field.to_string(), Value::String("state_123".to_string())); + } + + assert!(matches!( + validate_mixed_request(&request), + Err(LlmClientError::InvalidRequest { .. }) + )); + } + } + + let mut chat = Request::default(); + chat.metadata.get_or_insert_default().wire_format = Some(WireFormat::OpenAiChat); + chat.llm_request.preservation.requests.insert( + "openai_chat".into(), + serde_json::json!({"conversation": "chat extension"}), + ); + chat.llm_request.extensions.fields.insert( + "conversation".to_string(), + Value::String("chat extension".to_string()), + ); + assert!(validate_mixed_request(&chat).is_ok()); + + let mut untagged = Request::default(); + untagged.llm_request.extensions.fields.insert( + "conversation".to_string(), + Value::String("responses state".to_string()), + ); + assert!(validate_mixed_request(&untagged).is_err()); + } +} diff --git a/crates/switchyard-runner/src/route.rs b/crates/switchyard-runner/src/route.rs index 9d177c768..069fff6a3 100644 --- a/crates/switchyard-runner/src/route.rs +++ b/crates/switchyard-runner/src/route.rs @@ -1,7 +1,7 @@ // SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. // SPDX-License-Identifier: Apache-2.0 -//! One configured algorithm and the clients that serve its targets. +//! Configured route execution and the clients that serve its targets. use std::error::Error; use std::sync::Arc; @@ -13,6 +13,10 @@ use switchyard_protocol::{LlmClientError, ModelId, Request, Response, WireFormat use thiserror::Error; use crate::DecisionTarget; +use crate::privacy::{ + PrivacyDecision, PrivacyLane, PrivacyPolicy, external_signal_not_accepted, + has_external_restriction, record_selected_lane, selected_lane, validate_mixed_request, +}; /// Capabilities declared for one route. /// @@ -105,21 +109,52 @@ impl RunnerError { } } -/// A configured algorithm and the per-target clients its calls resolve through. -pub struct Route { +// Configured algorithm and targets for one execution lane. +pub(crate) struct ExecutionLane { algorithm: Arc, // Resolves each offloaded call to the client configured for the target the algorithm // selected. A route is a synthetic model with no upstream of its own, so this is a // per-target lookup, never one client serving the whole route. clients: ClientRouter, - caller_auth: Option, - capabilities: ModelCapabilities, anthropic_auxiliary_target: Option, responses_auxiliary_target: Option, decision_targets: Vec, models: Arc, } +impl ExecutionLane { + pub(crate) fn new( + algorithm: Arc, + clients: ClientRouter, + anthropic_auxiliary_target: Option, + responses_auxiliary_target: Option, + decision_targets: Vec, + models: RuntimeModels, + ) -> Self { + Self { + algorithm, + clients, + anthropic_auxiliary_target, + responses_auxiliary_target, + decision_targets, + models: Arc::new(models), + } + } +} + +struct PrivacyExecution { + policy: PrivacyPolicy, + restricted: ExecutionLane, +} + +/// One or more execution lanes behind a configured route. +pub struct Route { + standard: ExecutionLane, + privacy: Option, + caller_auth: Option, + capabilities: ModelCapabilities, +} + /// The selected model and untouched response produced by a route execution. pub struct RunOutput { pub selected_model: ModelId, @@ -138,22 +173,71 @@ impl Route { responses_auxiliary_target: Option, decision_targets: Vec, models: RuntimeModels, + ) -> Self { + Self::from_lane( + ExecutionLane::new( + algorithm, + clients, + anthropic_auxiliary_target, + responses_auxiliary_target, + decision_targets, + models, + ), + caller_auth, + capabilities, + ) + } + + pub(crate) fn from_lane( + lane: ExecutionLane, + caller_auth: Option, + capabilities: ModelCapabilities, ) -> Self { Self { - algorithm, - clients, + standard: lane, + privacy: None, caller_auth, capabilities, - anthropic_auxiliary_target, - responses_auxiliary_target, - decision_targets, - models: Arc::new(models), } } + pub(crate) fn with_privacy(mut self, policy: PrivacyPolicy, restricted: ExecutionLane) -> Self { + self.privacy = Some(PrivacyExecution { policy, restricted }); + self + } + + async fn select_lane( + &self, + request: &Request, + ) -> Result<(&ExecutionLane, Option), RunnerError> { + let Some(privacy) = &self.privacy else { + if has_external_restriction(request) { + return Err(external_signal_not_accepted().into()); + } + return Ok((&self.standard, None)); + }; + validate_mixed_request(request)?; + let route = request.llm_request.model.as_deref().unwrap_or_default(); + let algorithm = self.algorithm_name(); + let decision = privacy.policy.decide(request).await?; + tracing::debug!( + switchyard.route = route, + switchyard.algorithm = algorithm, + privacy.lane = decision.lane(), + privacy.source = decision.source(), + privacy.reason_code = decision.reason_code(), + "privacy lane selected" + ); + let lane = match decision.selected_lane() { + PrivacyLane::Standard => &self.standard, + PrivacyLane::Restricted => &privacy.restricted, + }; + Ok((lane, Some(decision))) + } + /// Returns the configured libsy algorithm name. pub fn algorithm_name(&self) -> &str { - self.algorithm.name() + self.standard.algorithm.name() } /// Returns model-list capability metadata. @@ -167,16 +251,25 @@ impl Route { } /// Resolves a selected model to this route's non-secret target metadata. - pub(crate) fn decision_target(&self, model: &ModelId) -> Option { - self.decision_targets + pub(crate) fn decision_target( + &self, + model: &ModelId, + request: &Request, + ) -> Option { + let targets = match (&self.privacy, selected_lane(request)) { + (None, _) | (Some(_), Some(PrivacyLane::Standard)) => &self.standard.decision_targets, + (Some(privacy), Some(PrivacyLane::Restricted)) => &privacy.restricted.decision_targets, + (Some(_), None) => return None, + }; + targets .iter() .find(|target| target.model == *model) .cloned() } - /// Returns the models grouped for one algorithm execution. + /// Returns the models grouped for the standard execution lane. pub fn models(&self) -> &RuntimeModels { - &self.models + &self.standard.models } /// Rejects a caller format incompatible with forwarded credentials. @@ -195,11 +288,12 @@ impl Route { request: Request, observer: Option, ) -> Result { + let (lane, _) = self.select_lane(&request).await?; let (selected_model, response) = switchyard_llm_client::run( - Arc::clone(&self.algorithm), - self.clients.clone(), + Arc::clone(&lane.algorithm), + lane.clients.clone(), request, - Arc::clone(&self.models), + Arc::clone(&lane.models), observer, ) .await?; @@ -211,14 +305,19 @@ impl Route { /// Completes routing-time calls without serving a post-routing completion. pub async fn decide(&self, request: Request) -> Result { - switchyard_llm_client::decide( - Arc::clone(&self.algorithm), - self.clients.clone(), + let (lane, privacy_decision) = self.select_lane(&request).await?; + let mut outcome = switchyard_llm_client::decide( + Arc::clone(&lane.algorithm), + lane.clients.clone(), request, - Arc::clone(&self.models), + Arc::clone(&lane.models), ) .await - .map_err(Into::into) + .map_err(RunnerError::from)?; + if let Some(decision) = privacy_decision { + record_selected_lane(&mut outcome.request, decision.selected_lane()); + } + Ok(outcome) } /// Executes a model-bearing provider operation through a compatible target. @@ -227,10 +326,11 @@ impl Route { request: Request, operation: AuxiliaryOperation, ) -> Result { + let (lane, _) = self.select_lane(&request).await?; let target = match operation { - AuxiliaryOperation::AnthropicCountTokens => &self.anthropic_auxiliary_target, + AuxiliaryOperation::AnthropicCountTokens => &lane.anthropic_auxiliary_target, AuxiliaryOperation::ResponsesInputTokens | AuxiliaryOperation::ResponsesCompact => { - &self.responses_auxiliary_target + &lane.responses_auxiliary_target } } .as_ref() @@ -242,3 +342,169 @@ impl Route { .map_err(Into::into) } } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + use std::sync::Mutex; + + use libsy::{Driver, Passthrough}; + use reqwest::StatusCode; + use switchyard_protocol::{Category, LlmResponse, RoutedLlmClient, text_response}; + + use super::*; + use crate::privacy::mark_privacy_restricted; + + struct CallThenRoute; + + #[async_trait::async_trait] + impl Algorithm for CallThenRoute { + fn name(&self) -> &str { + "call_then_route" + } + + async fn route( + self: Arc, + driver: Driver, + request: Request, + ) -> libsy::Result { + driver + .call_model( + request.clone(), + driver.models_for(&Category::Judge).to_vec(), + ) + .await?; + let (selected, fallbacks) = driver + .models_for(&Category::Any) + .split_first() + .ok_or(LibsyError::NoTargets)?; + Ok(RoutingOutcome::route_to( + selected.clone(), + fallbacks.to_vec(), + request, + )) + } + } + + struct RecordingClient { + calls: Arc>>, + } + + #[async_trait::async_trait] + impl RoutedLlmClient for RecordingClient { + async fn call(&self, request: Request) -> Result { + let model = request + .llm_request + .model + .map(ModelId::from) + .expect("driver should stamp a model"); + self.calls.lock().expect("calls lock").push(model.clone()); + if model == "restricted-primary" { + return Err(LlmClientError::UpstreamHttp { + status: StatusCode::SERVICE_UNAVAILABLE, + body: "unavailable".to_string(), + }); + } + Ok(Response { + llm_response: LlmResponse::Agg(text_response(Some(model.to_string()), "ok")), + metadata: None, + upstream_headers: Default::default(), + }) + } + } + + fn lane( + algorithm: Arc, + client: Arc, + models: RuntimeModels, + auxiliary_model: &str, + ) -> ExecutionLane { + let routed_models = models + .models_for(&Category::Any) + .iter() + .chain(models.models_for(&Category::Judge)) + .cloned() + .map(|model| (model, Arc::clone(&client))) + .collect(); + let auxiliary_client = + Arc::new(TranslatingLlmClient::new(&[]).expect("empty test client should be valid")); + ExecutionLane::new( + algorithm, + ClientRouter::new(routed_models), + Some(AuxiliaryTarget { + model: auxiliary_model.into(), + client: auxiliary_client, + }), + None, + Vec::new(), + models, + ) + } + + fn restricted_request() -> Request { + let mut request = Request::default(); + mark_privacy_restricted(&mut request); + request + } + + #[tokio::test] + async fn privacy_restricted_lane_covers_judge_fallbacks_and_auxiliary_calls() { + let calls = Arc::new(Mutex::new(Vec::new())); + let client: Arc = Arc::new(RecordingClient { + calls: Arc::clone(&calls), + }); + let standard = lane( + Arc::new(Passthrough), + Arc::clone(&client), + RuntimeModels::new(HashMap::from([( + Category::Any, + vec!["standard-primary".into()], + )])), + "standard-auxiliary", + ); + let restricted = lane( + Arc::new(CallThenRoute), + client, + RuntimeModels::new(HashMap::from([ + (Category::Judge, vec!["restricted-judge".into()]), + ( + Category::Any, + vec!["restricted-primary".into(), "restricted-fallback".into()], + ), + ])), + "restricted-auxiliary", + ); + let route = Route::from_lane(standard, None, ModelCapabilities::default()) + .with_privacy(PrivacyPolicy::new(true), restricted); + + let output = route + .execute(restricted_request(), None) + .await + .expect("restricted fallback should answer"); + assert_eq!(output.selected_model, "restricted-primary"); + assert_eq!( + output.response.served_model().map(ModelId::as_str), + Some("restricted-fallback") + ); + assert_eq!( + *calls.lock().expect("calls lock"), + [ + ModelId::from("restricted-judge"), + ModelId::from("restricted-primary"), + ModelId::from("restricted-fallback"), + ] + ); + + let error = route + .call_auxiliary( + restricted_request(), + AuxiliaryOperation::AnthropicCountTokens, + ) + .await + .expect_err("empty auxiliary client should reject the model"); + assert!( + error.to_string().contains("restricted-auxiliary"), + "{error}" + ); + } +} diff --git a/crates/switchyard-runner/src/runner.rs b/crates/switchyard-runner/src/runner.rs index ffba5120f..1e7086814 100644 --- a/crates/switchyard-runner/src/runner.rs +++ b/crates/switchyard-runner/src/runner.rs @@ -122,7 +122,7 @@ impl Runner { ) -> Option { let route = self.route(model.as_str())?; let resolve = |selected: &ModelId| { - let mut target = route.decision_target(selected)?; + let mut target = route.decision_target(selected, &outcome.request)?; let mut url = reqwest::Url::parse(&target.base_url).ok()?; let query: Vec<_> = url .query_pairs()