diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 17df71e2..716b086d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -416,6 +416,17 @@ jobs: grep -q 'ghcr.io/azure/kars-controller:latest' /tmp/kars-existing-aks.yaml grep -q 'name: agentmesh-relay' /tmp/kars-existing-aks.yaml + - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 + with: + node-version: "22" + - name: Restore existing chart-test runner + run: npm ci --prefix cli + - name: Local inference legacy-values regression + run: | + node cli/node_modules/vitest/vitest.mjs run \ + --config "$GITHUB_WORKSPACE/cli/vitest.config.ts" \ + --root deploy/helm/kars/tests + security-scan: name: Security Scan runs-on: ubuntu-latest diff --git a/controller/src/inference_policy_reconciler.rs b/controller/src/inference_policy_reconciler.rs index d7a9eab6..b0be7942 100644 --- a/controller/src/inference_policy_reconciler.rs +++ b/controller/src/inference_policy_reconciler.rs @@ -743,7 +743,10 @@ async fn ensure_profile_configmap( "app.kubernetes.io/managed-by".into(), "kars-controller".into(), ), - ("kars.azure.com/inferencepolicy".into(), owner.into()), + ( + "kars.azure.com/inferencepolicy".into(), + crate::labels::value(owner), + ), ( "kars.azure.com/artifact".into(), "inference-policy-profile".into(), diff --git a/controller/src/kars_receipt_launch.rs b/controller/src/kars_receipt_launch.rs index ecc97076..b0be01fc 100644 --- a/controller/src/kars_receipt_launch.rs +++ b/controller/src/kars_receipt_launch.rs @@ -160,6 +160,7 @@ mod tests { }], isolation: Some("enhanced".into()), memory: Some("review-memory".into()), + model_fallbacks: Vec::new(), }), display_name: Some("Review".into()), }, diff --git a/controller/src/kars_task.rs b/controller/src/kars_task.rs index d18b485e..6a7ed088 100644 --- a/controller/src/kars_task.rs +++ b/controller/src/kars_task.rs @@ -182,6 +182,11 @@ pub struct TaskBlueprint { #[serde(default, skip_serializing_if = "Option::is_none")] pub model: Option, + /// Ordered alternative inference routes; absent preserves the default route. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + #[schemars(schema_with = "crate::task_models::fallback_schema")] + pub model_fallbacks: Vec, + /// System prompt / standing instructions for the agent, in addition to the /// objective. Drives `KarsSandbox.spec.agent.instructions`. #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/controller/src/kars_task_authorization_tests.rs b/controller/src/kars_task_authorization_tests.rs index 56486f16..1e2a793b 100644 --- a/controller/src/kars_task_authorization_tests.rs +++ b/controller/src/kars_task_authorization_tests.rs @@ -36,6 +36,7 @@ fn spec() -> KarsTaskSpec { }], isolation: Some("standard".into()), memory: Some("team-memory".into()), + model_fallbacks: Vec::new(), }), ..Default::default() } diff --git a/controller/src/kars_task_execution.rs b/controller/src/kars_task_execution.rs index 6c83b563..9e44e0f5 100644 --- a/controller/src/kars_task_execution.rs +++ b/controller/src/kars_task_execution.rs @@ -22,7 +22,7 @@ use kube::api::{Api, DeleteParams, DynamicObject, ObjectMeta, PostParams, Preconditions}; use kube::core::ApiResource; use kube::{Client, ResourceExt}; -use serde_json::json; +use serde_json::{Value, json}; use crate::kars_task::{KarsTask, TaskBlueprint, TaskEnvelope}; @@ -107,6 +107,23 @@ fn network_policy(blueprint: &TaskBlueprint) -> serde_json::Value { }) } +/// Build primary and fallback routes from the shared normalized blueprint. +fn inference_spec(task_name: &str, blueprint: &TaskBlueprint) -> Result { + let primary = blueprint + .model + .as_ref() + .ok_or_else(|| contract_error("effective task model is missing".into()))?; + Ok(json!({ + "appliesTo": { "sandboxName": task_name }, + "modelPreference": { + "primary": primary, + "fallback": crate::task_models::fallback_routes( + &blueprint.model_fallbacks, &primary.provider, &primary.deployment, + ), + }, + })) +} + /// Materialize the InferencePolicy + KarsSandbox for a launched task using /// atomic creation or version-checked owned updates, then read sandbox status. pub async fn materialize( @@ -123,12 +140,7 @@ pub async fn materialize( // 1. InferencePolicy scoped to this sandbox. Model: blueprint wins, else // the controller default (required — without it the sandbox degrades). - let inference_spec = json!({ - "appliesTo": { "sandboxName": task_name }, - "modelPreference": { - "primary": blueprint.model, - }, - }); + let inference_spec = inference_spec(&task_name, &blueprint)?; apply_dynamic( client, namespace, diff --git a/controller/src/kars_task_execution_tests.rs b/controller/src/kars_task_execution_tests.rs index 08ce0c69..e7092230 100644 --- a/controller/src/kars_task_execution_tests.rs +++ b/controller/src/kars_task_execution_tests.rs @@ -5,6 +5,44 @@ use super::*; use wiremock::matchers::{method, path}; use wiremock::{Mock, MockServer, ResponseTemplate}; +#[test] +fn materialization_keeps_fallbacks_from_the_shared_effective_blueprint() { + use crate::kars_task::{KarsTaskSpec, TaskModel, blueprint::effective_blueprint_with_model}; + let primary = TaskModel { + provider: "azure-openai".into(), + deployment: "primary".into(), + }; + let spec = KarsTaskSpec { + blueprint: Some(TaskBlueprint { + model_fallbacks: vec![ + primary.clone(), + TaskModel { + provider: "backup".into(), + deployment: "second".into(), + }, + TaskModel { + provider: "backup".into(), + deployment: "second".into(), + }, + ], + ..Default::default() + }), + ..Default::default() + }; + let effective = effective_blueprint_with_model(&spec, &primary); + let policy = inference_spec("demo", &effective).unwrap(); + assert_eq!( + policy["modelPreference"]["primary"], + serde_json::to_value(primary).unwrap() + ); + assert_eq!( + policy["modelPreference"]["fallback"], + json!([ + {"provider":"backup","deployment":"second"}, + ]) + ); +} + const OBJECT_PATH: &str = "/apis/kars.azure.com/v1alpha1/namespaces/default/karssandboxes/demo"; const COLLECTION_PATH: &str = "/apis/kars.azure.com/v1alpha1/namespaces/default/karssandboxes"; @@ -210,6 +248,10 @@ async fn materialized_resources_match_the_authorization_blueprint() { deployment: "reviewed-model".into(), provider: String::new(), }), + model_fallbacks: vec![TaskModel { + provider: "backup-account".into(), + deployment: "fallback-model".into(), + }], instructions: Some(" Cite evidence. ".into()), tool_policy: Some("read-only".into()), mcp_servers: vec!["docs".into()], @@ -264,6 +306,10 @@ async fn materialized_resources_match_the_authorization_blueprint() { specs[0]["modelPreference"]["primary"], json!(effective.model) ); + assert_eq!( + specs[0]["modelPreference"]["fallback"], + json!(effective.model_fallbacks) + ); assert_eq!(specs[1]["runtime"]["kind"], "MicrosoftAgentFramework"); assert_eq!(specs[1]["sandbox"]["isolation"], json!(effective.isolation)); assert_eq!( diff --git a/controller/src/kars_team_reconciler/tests.rs b/controller/src/kars_team_reconciler/tests.rs index 87662c22..bdb7552f 100644 --- a/controller/src/kars_team_reconciler/tests.rs +++ b/controller/src/kars_team_reconciler/tests.rs @@ -9,6 +9,55 @@ use crate::kars_task::{ use crate::kars_team::{KarsTeamSpec, TeamCadence, TeamRole}; use k8s_openapi::apimachinery::pkg::apis::meta::v1::{Condition, Time}; +#[test] +fn inherited_fallbacks_reach_principal_member_and_run_authority_without_weakening_budget_gate() { + let mut team = team(); + let primary = TaskModel { + provider: "azure-openai".into(), + deployment: "primary".into(), + }; + team.spec.blueprint = Some(TaskBlueprint { + model: Some(primary.clone()), + model_fallbacks: vec![TaskModel { + provider: "backup".into(), + deployment: "secondary".into(), + }], + ..Default::default() + }); + assert!( + crate::kars_task::validate_execution_contract(&specs::run_spec(&team, "").unwrap()) + .unwrap_err() + .contains("UnsupportedLaunchBudget") + ); + team.spec.envelope.budget = None; + let role = TeamRole { + name: "reviewer".into(), + ..Default::default() + }; + for spec in [ + specs::principal_spec(&team), + specs::member_spec(&team, &role), + specs::run_spec(&team, "").unwrap(), + ] { + let effective = + crate::kars_task::blueprint::effective_blueprint_with_model(&spec, &primary); + assert_eq!(effective.model_fallbacks.len(), 1); + assert_eq!(effective.model_fallbacks[0].provider, "backup"); + let mut changed = spec.clone(); + changed.blueprint.as_mut().unwrap().model_fallbacks[0].deployment = "changed".into(); + assert_ne!( + spec.authorization_digest_with_model(&primary), + changed.authorization_digest_with_model(&primary), + ); + } + let explicit = TeamRole { + blueprint: Some(TaskBlueprint::default()), + ..role + }; + let member = specs::member_spec(&team, &explicit); + assert!(member.blueprint.unwrap().model_fallbacks.is_empty()); +} + pub(super) fn team() -> KarsTeam { let mut team = KarsTeam::new( "eng", diff --git a/controller/src/labels.rs b/controller/src/labels.rs new file mode 100644 index 00000000..9c8beec1 --- /dev/null +++ b/controller/src/labels.rs @@ -0,0 +1,43 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +/// Keep short label values readable and long/non-ASCII resource identities stable. +pub(crate) fn value(input: &str) -> String { + let valid = input.len() <= 63 + && input + .bytes() + .all(|c| c.is_ascii_alphanumeric() || matches!(c, b'-' | b'_' | b'.')) + && input + .as_bytes() + .first() + .is_some_and(u8::is_ascii_alphanumeric) + && input + .as_bytes() + .last() + .is_some_and(u8::is_ascii_alphanumeric); + if valid { + return input.to_string(); + } + format!( + "id-{}", + crate::kars_receipt_log::sha256_hex(input.as_bytes()) + ) + .chars() + .take(63) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + #[test] + fn resource_names_fit_label_values_without_colliding() { + assert_eq!(value("short-name"), "short-name"); + let a = format!("{}a", "long".repeat(35)); + let b = format!("{}b", "long".repeat(35)); + assert!(value(&a).len() <= 63); + assert_ne!(value(&a), value(&b)); + assert_eq!(value(&a), value(&a)); + assert!(value("資料").is_ascii()); + } +} diff --git a/controller/src/main.rs b/controller/src/main.rs index 5cdaeb45..0ec22660 100644 --- a/controller/src/main.rs +++ b/controller/src/main.rs @@ -58,6 +58,7 @@ mod kars_task_execution; mod kars_task_reconciler; mod kars_team; mod kars_team_reconciler; +mod labels; mod leader_election; mod mcp_server; mod mcp_server_reconciler; @@ -72,6 +73,7 @@ mod providers; mod reconciler; mod signer_policy; mod status; +mod task_models; mod team_commons; mod team_digest; #[allow(dead_code)] // helpers consumed by tool_policy_reconciler + future slices. diff --git a/controller/src/reconciler/inference.rs b/controller/src/reconciler/inference.rs new file mode 100644 index 00000000..cf4f3a58 --- /dev/null +++ b/controller/src/reconciler/inference.rs @@ -0,0 +1,181 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use k8s_openapi::api::core::v1::Secret; +use kube::{ + Client, + api::{Api, DeleteParams}, +}; +use serde::Deserialize; +use serde_json::{Value, json}; +use std::collections::BTreeMap; + +const PROVIDERS_SECRET: &str = "kars-inference-providers"; + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase", deny_unknown_fields)] +struct LocalTarget { + namespace: String, + match_labels: BTreeMap, + ports: Vec, +} + +fn valid_namespace(namespace: &str) -> bool { + !namespace.is_empty() + && namespace.len() <= 63 + && namespace + .bytes() + .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == b'-') + && namespace + .as_bytes() + .first() + .is_some_and(u8::is_ascii_alphanumeric) + && namespace + .as_bytes() + .last() + .is_some_and(u8::is_ascii_alphanumeric) +} + +pub(super) fn local_egress_rules(targets: &str, namespaces: &str) -> Result, String> { + let targets: Vec = if targets.trim().is_empty() { + Vec::new() + } else { + serde_json::from_str(targets) + .map_err(|e| format!("LOCAL_INFERENCE_TARGETS_JSON is invalid: {e}"))? + }; + if !targets.is_empty() { + return targets.into_iter().map(|target| { + if !valid_namespace(&target.namespace) || target.match_labels.is_empty() + || target.match_labels.keys().any(|key| key.trim().is_empty()) + || target.ports.is_empty() || target.ports.contains(&0) + { + return Err("local inference targets require a namespace, pod labels and TCP ports in 1..65535".into()); + } + Ok(json!({ + "to": [{ + "namespaceSelector": {"matchLabels": {"kubernetes.io/metadata.name": target.namespace}}, + "podSelector": {"matchLabels": target.match_labels} + }], + "ports": target.ports.iter().map(|port| json!({"protocol": "TCP", "port": port})).collect::>() + })) + }).collect(); + } + namespaces.split(',').map(str::trim).filter(|ns| !ns.is_empty()).map(|namespace| { + if !valid_namespace(namespace) { + return Err(format!("invalid local inference namespace: {namespace}")); + } + Ok(json!({ + "to": [{"namespaceSelector": {"matchLabels": {"kubernetes.io/metadata.name": namespace}}}], + "ports": [{"protocol": "TCP"}] + })) + }).collect() +} + +pub(super) fn configured_local_egress_rules() -> Result, String> { + local_egress_rules( + &std::env::var("LOCAL_INFERENCE_TARGETS_JSON").unwrap_or_default(), + &std::env::var("LOCAL_INFERENCE_NAMESPACES").unwrap_or_default(), + ) +} + +pub(super) fn provider_env_from() -> Value { + json!([{"secretRef": {"name": PROVIDERS_SECRET, "optional": true}}]) +} + +/// A resource-version annotation restarts router pods when environment credentials change. +pub(super) async fn mirror_providers( + client: &Client, + source_ns: &str, + target_ns: &str, + sandbox: &str, +) -> Result, kube::Error> { + let source: Api = Api::namespaced(client.clone(), source_ns); + let Some(secret) = source.get_opt(PROVIDERS_SECRET).await? else { + let target: Api = Api::namespaced(client.clone(), target_ns); + if let Some(existing) = target.get_opt(PROVIDERS_SECRET).await? + && existing + .metadata + .annotations + .as_ref() + .is_some_and(|annotations| { + annotations + .get(super::governance_mounts::MIRROR_SOURCE_KIND_ANNOTATION) + .map(String::as_str) + == Some("InferenceProviders") + && annotations + .get(super::governance_mounts::MIRROR_SOURCE_NS_ANNOTATION) + .map(String::as_str) + == Some(source_ns) + }) + { + match target + .delete(PROVIDERS_SECRET, &DeleteParams::default()) + .await + { + Ok(_) => {} + Err(kube::Error::Api(error)) if error.code == 404 => {} + Err(error) => return Err(error), + } + } + return Ok(None); + }; + super::governance_mounts::mirror_secret( + client, + PROVIDERS_SECRET, + source_ns, + target_ns, + sandbox, + "InferenceProviders", + ) + .await?; + Ok(secret.metadata.resource_version) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn local_inference_is_opt_in_and_precise_targets_replace_namespace_allowance() { + assert!(local_egress_rules("", "").unwrap().is_empty()); + let rules = local_egress_rules( + r#"[{"namespace":"models","matchLabels":{"app":"model"},"ports":[5000]}]"#, + "kars-local-inference", + ) + .unwrap(); + assert_eq!(rules.len(), 1); + assert_eq!( + rules[0]["to"][0]["namespaceSelector"]["matchLabels"]["kubernetes.io/metadata.name"], + "models" + ); + assert_eq!( + rules[0]["to"][0]["podSelector"]["matchLabels"]["app"], + "model" + ); + assert_eq!(rules[0]["ports"][0]["port"], 5000); + } + + #[test] + fn malformed_targets_do_not_fall_back_to_broad_allowances() { + for target in [ + "invalid", + "{}", + r#"[{"namespace":"models","ports":[80]}]"#, + r#"[{"namespace":"models","matchLabels":{},"ports":[80]}]"#, + r#"[{"namespace":"models","matchLabels":{"app":"x"},"ports":[0]}]"#, + r#"[{"namespace":"models","matchLabels":{"app":"x"},"ports":[65536]}]"#, + ] { + assert!(local_egress_rules(target, "models").is_err(), "{target}"); + } + assert!(local_egress_rules("[]", "models").is_ok()); + assert!(local_egress_rules("", "INVALID").is_err()); + } + + #[test] + fn provider_credentials_are_optional_router_inputs() { + assert_eq!( + provider_env_from(), + json!([{"secretRef":{"name":"kars-inference-providers","optional":true}}]) + ); + } +} diff --git a/controller/src/reconciler/mod.rs b/controller/src/reconciler/mod.rs index aced8831..574c02bb 100644 --- a/controller/src/reconciler/mod.rs +++ b/controller/src/reconciler/mod.rs @@ -36,6 +36,7 @@ use crate::fedcred::{FedCredConfig, FedCredManager}; pub(crate) mod byo_contract; mod dev_env; pub(crate) mod governance_mounts; +mod inference; mod mcp_egress; pub(crate) mod trustgraph_mount; @@ -976,6 +977,9 @@ async fn reconcile(sandbox: Arc, ctx: Arc) -> Result, ctx: Arc) -> Result, ctx: Arc) -> Result, ctx: Arc) -> Result, ctx: Arc) -> Result (Option<&'static str>, &' } } +pub(crate) fn cluster_default_model() -> Option { + [ + "KARS_TASK_DEFAULT_MODEL", + "AZURE_OPENAI_DEPLOYMENT", + "DEFAULT_MODEL", + ] + .into_iter() + .filter_map(|key| std::env::var(key).ok()) + .map(|value| value.trim().to_string()) + .find(|value| !value.is_empty()) +} + pub(crate) fn sandbox_node_selector_from( raw: &str, default_pool: &str, diff --git a/controller/src/task_models.rs b/controller/src/task_models.rs new file mode 100644 index 00000000..32ef6b61 --- /dev/null +++ b/controller/src/task_models.rs @@ -0,0 +1,156 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use crate::kars_task::TaskModel; +use serde_json::{Value, json}; + +/// Bounds apply to the new field only, not existing primary routes or rosters. +pub fn fallback_schema(_: &mut schemars::SchemaGenerator) -> schemars::Schema { + schemars::json_schema!({ + "type": "array", + "maxItems": 8, + "items": { + "type": "object", + "required": ["deployment", "provider"], + "properties": { + "deployment": {"type": "string", "minLength": 1, "maxLength": 253, "pattern": "\\S"}, + "provider": {"type": "string", "minLength": 1, "maxLength": 253, "pattern": "\\S"} + } + } + }) +} + +pub fn fallback_routes(routes: &[TaskModel], provider: &str, deployment: &str) -> Vec { + let mut seen = + std::collections::HashSet::from([(provider.to_string(), deployment.to_string())]); + routes + .iter() + .filter(|route| !route.provider.trim().is_empty() && !route.deployment.trim().is_empty()) + .filter(|route| seen.insert((route.provider.clone(), route.deployment.clone()))) + .map(|route| json!({"provider": route.provider, "deployment": route.deployment})) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn canonical_blueprint_preserves_fallbacks_and_authorization_binds_every_route() { + use crate::kars_task::{ + KarsTaskSpec, TaskBlueprint, blueprint::effective_blueprint_with_model, + }; + let primary = TaskModel { + provider: "azure-openai".into(), + deployment: "primary".into(), + }; + let spec = KarsTaskSpec { + objective: "Review".into(), + blueprint: Some(TaskBlueprint { + model: Some(primary.clone()), + model_fallbacks: vec![ + TaskModel { + provider: "account-a".into(), + deployment: "model-a".into(), + }, + TaskModel { + provider: "account-b".into(), + deployment: "model-b".into(), + }, + ], + ..Default::default() + }), + ..Default::default() + }; + let original = serde_json::to_value(&spec).unwrap(); + let effective = effective_blueprint_with_model(&spec, &primary); + assert_eq!( + serde_json::to_value(effective).unwrap()["modelFallbacks"], + original["blueprint"]["modelFallbacks"] + ); + assert_eq!( + spec.authorization_configuration_with_model(&primary)["blueprint"]["modelFallbacks"], + original["blueprint"]["modelFallbacks"] + ); + let digest = spec.authorization_digest_with_model(&primary); + for (pointer, value) in [ + ( + "/blueprint/modelFallbacks/0/provider", + json!("other-account"), + ), + ( + "/blueprint/modelFallbacks/0/deployment", + json!("other-model"), + ), + ("/blueprint/modelFallbacks", json!([])), + ( + "/blueprint/modelFallbacks", + json!([ + {"provider":"account-b","deployment":"model-b"}, + {"provider":"account-a","deployment":"model-a"}, + ]), + ), + ] { + let mut changed = original.clone(); + *changed.pointer_mut(pointer).unwrap() = value; + let changed: KarsTaskSpec = serde_json::from_value(changed).unwrap(); + assert_ne!( + digest, + changed.authorization_digest_with_model(&primary), + "{pointer}" + ); + } + } + + #[test] + fn fallback_routes_preserve_order_and_provider_identity() { + let routes = [ + ("a", "primary"), + ("b", "same"), + ("b", "same"), + ("c", "same"), + ("", "invalid"), + ("b", " "), + ("a", "last"), + ] + .map(|(provider, deployment)| TaskModel { + provider: provider.into(), + deployment: deployment.into(), + }); + assert_eq!( + fallback_routes(&routes, "a", "primary"), + vec![ + json!({"provider":"b","deployment":"same"}), + json!({"provider":"c","deployment":"same"}), + json!({"provider":"a","deployment":"last"}), + ] + ); + assert!(fallback_routes(&[], "a", "primary").is_empty()); + } + + #[test] + fn legacy_blueprint_omits_new_field() { + let blueprint: crate::kars_task::TaskBlueprint = serde_json::from_value(json!({})).unwrap(); + assert!(blueprint.model_fallbacks.is_empty()); + assert!( + serde_json::to_value(blueprint) + .unwrap() + .get("modelFallbacks") + .is_none() + ); + } + + #[test] + fn fallback_schema_is_bounded_without_changing_primary_or_roster() { + let schema = + serde_json::to_value(schemars::schema_for!(crate::kars_task::TaskBlueprint)).unwrap(); + assert_eq!(schema["properties"]["modelFallbacks"]["maxItems"], 8); + assert_eq!( + schema["properties"]["modelFallbacks"]["items"]["properties"]["provider"]["maxLength"], + 253 + ); + let team = + serde_json::to_value(schemars::schema_for!(crate::kars_team::KarsTeamSpec)).unwrap(); + assert!(team["properties"]["roster"].get("maxItems").is_none()); + } +} diff --git a/deploy/helm/kars/templates/controller-deployment.yaml b/deploy/helm/kars/templates/controller-deployment.yaml index 48315287..f7337d57 100644 --- a/deploy/helm/kars/templates/controller-deployment.yaml +++ b/deploy/helm/kars/templates/controller-deployment.yaml @@ -1,3 +1,4 @@ +{{- $localInference := .Values.localInference | default dict }} apiVersion: apps/v1 kind: Deployment metadata: @@ -51,6 +52,10 @@ spec: value: {{ .Values.azure.workloadIdentity.clientId | quote }} - name: KARS_SANDBOX_NODE_SELECTOR_JSON value: {{ .Values.sandbox.nodeSelector | default dict | toJson | quote }} + - name: LOCAL_INFERENCE_NAMESPACES + value: {{ join "," ($localInference.namespaces | default list) | quote }} + - name: LOCAL_INFERENCE_TARGETS_JSON + value: {{ ($localInference.targets | default list) | toJson | quote }} {{- if .Values.controller.byoStrict }} - name: BYO_STRICT_MODE value: "1" diff --git a/deploy/helm/kars/templates/crd-karstask.yaml b/deploy/helm/kars/templates/crd-karstask.yaml index 35dbdd41..78fd122f 100644 --- a/deploy/helm/kars/templates/crd-karstask.yaml +++ b/deploy/helm/kars/templates/crd-karstask.yaml @@ -125,6 +125,26 @@ spec: - deployment - provider type: object + modelFallbacks: + description: Ordered alternative inference routes; absent preserves the default route. + type: array + maxItems: 8 + items: + type: object + required: + - deployment + - provider + properties: + deployment: + type: string + minLength: 1 + maxLength: 253 + pattern: '\S' + provider: + type: string + minLength: 1 + maxLength: 253 + pattern: '\S' runtime: description: |- Harness/runtime (`OpenClaw`, `OpenAIAgents`, `MicrosoftAgentFramework`, diff --git a/deploy/helm/kars/templates/crd-karsteam.yaml b/deploy/helm/kars/templates/crd-karsteam.yaml index 29b19291..bbd68ceb 100644 --- a/deploy/helm/kars/templates/crd-karsteam.yaml +++ b/deploy/helm/kars/templates/crd-karsteam.yaml @@ -121,6 +121,26 @@ spec: - deployment - provider type: object + modelFallbacks: + description: Ordered alternative inference routes; absent preserves the default route. + type: array + maxItems: 8 + items: + type: object + required: + - deployment + - provider + properties: + deployment: + type: string + minLength: 1 + maxLength: 253 + pattern: '\S' + provider: + type: string + minLength: 1 + maxLength: 253 + pattern: '\S' runtime: description: |- Harness/runtime (`OpenClaw`, `OpenAIAgents`, `MicrosoftAgentFramework`, @@ -384,6 +404,26 @@ spec: - deployment - provider type: object + modelFallbacks: + description: Ordered alternative inference routes; absent preserves the default route. + type: array + maxItems: 8 + items: + type: object + required: + - deployment + - provider + properties: + deployment: + type: string + minLength: 1 + maxLength: 253 + pattern: '\S' + provider: + type: string + minLength: 1 + maxLength: 253 + pattern: '\S' runtime: description: |- Harness/runtime (`OpenClaw`, `OpenAIAgents`, `MicrosoftAgentFramework`, diff --git a/deploy/helm/kars/tests/src/local-inference.test.ts b/deploy/helm/kars/tests/src/local-inference.test.ts new file mode 100644 index 00000000..b0fcde70 --- /dev/null +++ b/deploy/helm/kars/tests/src/local-inference.test.ts @@ -0,0 +1,83 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { execFileSync } from "node:child_process"; +import { copyFileSync, mkdirSync, readFileSync, rmSync, writeFileSync } from "node:fs"; +import { createRequire } from "node:module"; +import { join } from "node:path"; +import { fileURLToPath } from "node:url"; + +// Reuse the CLI's existing Vitest configuration and YAML dependency, without +// adding a second runner or changing any CLI installation/upgrade code. +const require = createRequire(new URL("../../../../../cli/package.json", import.meta.url)); +const { parse, parseAllDocuments } = require("yaml"); +const chart = fileURLToPath(new URL("../../", import.meta.url)); +let fixtureNumber = 0; + +function renderLegacy(localInference?: unknown) { + const fixture = join(chart, "tests", `.render-fixture-${process.pid}-${++fixtureNumber}`); + mkdirSync(join(fixture, "templates"), { recursive: true }); + try { + copyFileSync(join(chart, "Chart.yaml"), join(fixture, "Chart.yaml")); + copyFileSync( + join(chart, "templates/controller-deployment.yaml"), + join(fixture, "templates/controller-deployment.yaml"), + ); + // No new chart defaults may fill the missing map: this models Helm's + // --reuse-values path, not a normal old-values/new-defaults coalescing. + writeFileSync(join(fixture, "values.yaml"), "{}\n"); + const legacy = parse(readFileSync(join(chart, "values.yaml"), "utf8")); + delete legacy.localInference; + if (localInference !== undefined) legacy.localInference = localInference; + legacy.controller.replicas = 3; + legacy.controller.image.repository = "registry.customer.test/custom-controller"; + legacy.controller.extraEnv = [{ name: "CUSTOMER_SETTING", value: "retained" }]; + legacy.sandbox.nodeSelector = { "customer.example/pool": "isolated" }; + writeFileSync(join(fixture, "legacy-values.json"), JSON.stringify(legacy)); + const rendered = execFileSync("helm", [ + "template", "customer-release", fixture, "--namespace", "customer-system", + "--values", join(fixture, "legacy-values.json"), + ], { encoding: "utf8", timeout: 10_000, stdio: ["ignore", "pipe", "pipe"] }); + const deployment = parseAllDocuments(rendered).map((doc: { toJSON(): unknown }) => doc.toJSON()) + .find((doc: { kind?: string } | null) => doc?.kind === "Deployment"); + return deployment; + } finally { + rmSync(fixture, { recursive: true, force: true }); + } +} + +describe("local inference Helm upgrade compatibility", () => { + for (const [name, value] of [ + ["missing section", undefined], + ["null section", null], + ["empty section", {}], + ["missing targets", { namespaces: [] }], + ["null lists", { namespaces: null, targets: null }], + ]) { + it(`keeps egress disabled and customer settings intact with ${name}`, () => { + const deployment = renderLegacy(value); + const container = deployment.spec.template.spec.containers[0]; + const env = Object.fromEntries(container.env.map((entry: { name: string; value?: string }) => + [entry.name, entry.value])); + expect(env.LOCAL_INFERENCE_NAMESPACES).toBe(""); + expect(env.LOCAL_INFERENCE_TARGETS_JSON).toBe("[]"); + expect(env.CUSTOMER_SETTING).toBe("retained"); + expect(JSON.parse(env.KARS_SANDBOX_NODE_SELECTOR_JSON)).toEqual({ "customer.example/pool": "isolated" }); + expect(deployment.spec.replicas).toBe(3); + expect(deployment.metadata.namespace).toBe("customer-system"); + expect(container.image).toBe("registry.customer.test/custom-controller:latest"); + }); + } + + it("preserves an explicitly configured local namespace and target", () => { + const target = { + namespace: "models", matchLabels: { app: "private-model" }, ports: [8000], + }; + const deployment = renderLegacy({ namespaces: ["models"], targets: [target] }); + const env = Object.fromEntries(deployment.spec.template.spec.containers[0].env.map( + (entry: { name: string; value?: string }) => [entry.name, entry.value], + )); + expect(env.LOCAL_INFERENCE_NAMESPACES).toBe("models"); + expect(JSON.parse(env.LOCAL_INFERENCE_TARGETS_JSON)).toEqual([target]); + }); +}); diff --git a/deploy/helm/kars/values.yaml b/deploy/helm/kars/values.yaml index 4138989d..94b29da8 100644 --- a/deploy/helm/kars/values.yaml +++ b/deploy/helm/kars/values.yaml @@ -2,6 +2,13 @@ # NOTE: For production, replace "latest" tags with specific image digests # (e.g., sha256:abc123...) and set pullPolicy to IfNotPresent. +localInference: + # Opt in to namespace-wide model egress, or prefer precise pod targets below. + namespaces: [] + # Explicit targets replace namespace-wide allowances. Each requires a + # namespace, nonempty matchLabels and destination pod TCP ports. + targets: [] + # Controller configuration controller: image: diff --git a/docs/local-inference.md b/docs/local-inference.md new file mode 100644 index 00000000..f8e10b54 --- /dev/null +++ b/docs/local-inference.md @@ -0,0 +1,134 @@ +# Local inference and model failover + +Kars can route to operator-configured OpenAI-compatible endpoints alongside +existing Azure OpenAI, Foundry, GitHub Models, Copilot, and typed Ollama or +Anthropic providers. No model or infrastructure is deployed automatically. + +## Configure additional providers + +The optional `kars-inference-providers` Secret in the controller namespace +contains router environment variables: + +| Variable | Meaning | +|---|---| +| `KARS_PROVIDER__ENDPOINT` | Provider's complete API base URL | +| `KARS_PROVIDER__API_KEY` | Optional credential for that provider | +| `KARS_PROVIDER__TOKEN` | Alternative provider credential; API key takes precedence | + +For example, `KARS_PROVIDER_LOCAL_ENDPOINT` maps to provider tag `local`. +Tags are case-insensitive; underscores in environment names become hyphens. +The controller mirrors this Secret into sandbox namespaces and exposes it only +to the inference-router container, never the agent. Secret revisions cause the +next sandbox reconciliation to refresh the pod's environment. + +Use a base URL including `/v1` for servers such as vLLM or llama.cpp that serve +`/v1/chat/completions`. Azure-owned endpoints retain the existing `/openai/v1` +prefix behavior. Custom endpoints do not receive that Azure prefix. Typed +Ollama routes retain their `/v1` translation and receive no credential. + +Each named route uses only its own credential, or no authentication when that +credential is absent. It never borrows the default API key, Workload Identity, +IMDS, or sidecar identity. This also applies to model Services in arbitrary +Kubernetes namespaces. The legacy default route's authentication behavior and +its dedicated `kars-local-inference` namespace exclusion remain unchanged. + +For a named **Copilot** endpoint, the credential must be a GitHub OAuth/seat +token accepted by Copilot, not an Azure/OpenAI API key or an already-exchanged +inference JWT. The router performs Copilot's exchange and maintains a separate +cache for each selected provider identity/credential. Missing named Copilot +credentials fail before inference HTTP; they cannot borrow another account. +The global `COPILOT_GITHUB_TOKEN` cache is reserved for the legacy default route. + +## Allow local model traffic + +Local-model egress is opt-in. Prefer a narrowly scoped Helm value: + +```yaml +localInference: + targets: + - namespace: kars-local-inference + matchLabels: + app: inference-model + ports: [5000] +``` + +Use the model pod's destination port: NetworkPolicy is commonly evaluated after +Service DNAT. Targets require a namespace, nonempty pod labels, and TCP ports in +`1..65535`; malformed targets fail reconciliation rather than broadening access. +Explicit targets replace any namespace-wide allowances. + +For operators intentionally trusting all model services in a namespace, +`localInference.namespaces` accepts namespace names instead. Its default is +empty. Existing default-deny, agent UID isolation, Content Safety, prompt +shields, policy floors, and declared guardrail pipelines remain enforced. +Older installations using Helm `--reuse-values` may have no `localInference` +section at all. Missing sections or lists remain empty and keep local egress +disabled, without replacing customer values. + +## Select primary and fallback routes + +An `InferencePolicy` selects an ordered route chain: + +```yaml +modelPreference: + primary: + provider: local + deployment: primary-model + fallback: + - provider: foundry + deployment: fallback-model +``` + +For a `KarsTask` or a team blueprint, use `model` and `modelFallbacks`: + +```yaml +blueprint: + model: + provider: local + deployment: primary-model + modelFallbacks: + - provider: foundry + deployment: fallback-model +``` + +The new fallback list permits up to eight nonblank provider/deployment pairs. +Exact duplicate pairs and repeats of the primary are removed without changing +the remaining order. Existing primary route fields and team roster sizes are +not restricted by the new field's limits. + +An explicit `InferencePolicy.spec.provider` remains authoritative for the +primary candidate, even if `modelPreference.primary.provider` contains a +conflicting informational tag. Each fallback keeps its own provider, and the +true legacy default remains a separate final candidate. + +Without that authoritative field, the existing primary provider label remains +metadata unless its named endpoint is separately registered. For example, +`anthropic` / `claude-prod` can continue using the default Azure/Foundry route. +Merely supplying a global native-provider key or Ollama URL does not change +that intent. Metadata and explicit native routes have separate health identities +when the same label can select different backends. + +The router tracks health per provider/deployment and tries another candidate +on connection failures known to precede acceptance, HTTP 429, or HTTP 5xx. +Authentication/configuration acquisition failures and ambiguous transport +errors are not retry triggers. Ordinary client/auth/policy errors are returned +rather than retried. An unavailable-model response can additionally recover +once to the configured default model; capability caches include the immutable +provider identity as well as endpoint/model, never credential text. Credential +updates roll the sandbox process and its capability caches. + +Known 429/5xx rejections remain retryable if their bodies truncate after headers. +That is distinct from an accepted 2xx response, an ordinary 4xx rejection, or an +ambiguous connection loss; those body/transport failures are never replayed. + +Both buffered and streaming failover end when successful response headers are +accepted. A later body failure is never replayed onto another provider, +including the chat-to-Responses recovery path. Responses-only model recovery +retains the selected provider and effective deployment. + +Provider families must support the request API. The public typed Anthropic +provider continues to use `/anthropic/v1/messages`; Kars does not silently +translate OpenAI chat requests into Anthropic requests. + +No credentials, private registries, customer deployment values, or hardware +requirements are needed in a public model-routing configuration. diff --git a/docs/security-audits/2026-09-07-inference-local-failover.md b/docs/security-audits/2026-09-07-inference-local-failover.md new file mode 100644 index 00000000..0f9a911c --- /dev/null +++ b/docs/security-audits/2026-09-07-inference-local-failover.md @@ -0,0 +1,172 @@ +# Security Audit — Inference routing and local failover + +Date: 2026-09-07 +Review status: Maintainer sign-off received; independent human review and sign-off pending. + +## Scope + +PR4, initially based on `8fe755b4`, extracts inference routing from canonical +reference `ce9044077` and integrates foundation hash centralization `ea6f5789`. +Gated paths include `controller/src/reconciler/`, +`controller/src/kars_task.rs`, `inference-router/src/routes/`, and the +corresponding Helm CRD and configuration templates. + +No runtime orchestration, task delivery, access-request service, GitHub write +service, witness deployment, private cluster values, or publishing workflow is +introduced. No cluster changes, image publication, or pull-request merges were +performed. + +## T1: New capability / attack surface? YES + +- Operator-configured named provider endpoints and per-provider credentials. +- Ordered primary/fallback model routes, including pre-acceptance SSE failover. +- Optional local-inference NetworkPolicy destinations. +- Optional router-only mirroring of the controller namespace's + `kars-inference-providers` Secret. + +The agent container does not receive the new Secret. Policy routing determines +which configured provider is used. Fallback reconstruction starts from the true +default rather than inheriting another candidate's endpoint or credential. + +## T2: Security-control change? YES, bounded and additive + +- Existing Content Safety and prompt-shield defaults remain enabled. +- Existing input/output guardrail and policy-floor enforcement paths remain. +- Host matching uses parsed hostnames. Named routes carry explicit authentication + provenance and use their own credential or no auth; they never borrow the + default API key, WI/IMDS, or sidecar token. Legacy default behavior is retained. +- Named Copilot routes exchange their own GitHub seat token in an isolated + provider cache; only the legacy default uses the global Copilot account. +- Local egress is opt-in. Precise targets require namespace, pod labels, and + valid TCP ports; they replace broad namespace allowances. Invalid targets + fail reconciliation rather than falling back to broader access. +- New fallback bounds apply to the new field only. Existing team roster sizes, + primary model fields, installation profiles, and explicit client model choice + without a model preference remain compatible. +- Secret revisions change a pod-template annotation so the next reconciliation + refreshes environment credentials. Source removal removes only a matching + controller-mirrored copy. + +## T3: Availability / fail-open risk? MIXED, explicitly bounded + +- Connection failures known to precede acceptance and HTTP 429/5xx can use + another configured provider. Typed authentication/configuration failures and + unknown acceptance states fail closed. +- Unavailable-model recovery is bounded to the configured default and cached + per provider identity/endpoint/model; Responses-only recovery remains with the actual selected + provider and deployment. +- No buffered or streaming generation is replayed after accepted 2xx headers, + including a body failure during chat-to-Responses recovery. +- Explicit policy providers retain primary-route precedence; true-default + fallback remains separate. Missing local-inference Helm values stay disabled + on `--reuse-values` upgrades. +- Named-provider Secrets remain optional for installations using only the + existing default route. Mirror errors are propagated, not concealed. +- The existing public typed provider gates remain; no cross-family + OpenAI-to-Anthropic translation is claimed. + +## Verification + +Closure follow-up after `c47418fc` distinguishes legacy primary-model metadata +from explicit native/named routing, independently of native credential presence. +It also permits retry after a known 429/5xx rejection whose body truncates, +without retrying accepted, ambiguous, authentication/configuration, or ordinary +4xx failures. New HTTP regressions cover metadata/default compatibility, +registered versus native intent, health isolation, buffered/streaming truncated +rejections, and chat-to-Responses recovery. + +Current closure qualification after merging immutable `1fd97818`: + +- 84 router tests passed: 37 targeted unit/HTTP regressions and 47 existing + HTTP integration tests across routing, credentials, guardrails, and egress. +- 76 controller tests passed, including the complete shared + `authorization_configuration_with_model` snapshot, fallback digest mutation, + inheritance, materialization, and fail-closed finite-budget behavior. +- The inheritance fixture now consumes the current fallible `run_spec` API; + production objective/history failure propagation is unchanged. +- Six Helm regressions passed; strict all-target Clippy passed for both crates. +- The existing root Cargo target was exclusively leased, with offline/locked + commands and incremental compilation disabled. No duplicate target or + external push was used. + +The earlier qualification below is retained as historical evidence for +`c47418fc`. + +Six independent-review blockers were repaired after `7a2a5d11`. Added +regressions cover ambient API-key/sidecar isolation, the real Copilot-host +exchange branch, provider precedence in buffered and streaming chat, separate +accounts at the same endpoint/model, typed failures, and actual +chat-to-Responses recovery with truncated accepted bodies. + +Repair qualification after merging immutable, parent-qualified `415b53ca`: + +- 76 targeted router unit/HTTP tests passed, including both first-time and + cached Responses recovery, provider credential isolation, and existing + moderation, Foundry-route, egress, and explicit-client-model regressions. +- 70 targeted controller tests passed. New integration cases verify that + `modelFallbacks` survives the shared effective-blueprint normalizer, changes + task authorization when mutated or reordered, reaches the materialized + InferencePolicy, and is inherited by principal/member/run specifications. + Explicit role overrides remain complete overrides. +- `UnsupportedLaunchBudget` remains fail closed; finite total/subtree and + monetary budgets are not restored as misleading daily-token mappings. +- Strict Clippy passed for both crates with all targets and warnings denied. +- Six offline Helm regressions passed using the existing CLI Vitest runner + and cached Vitest 4.1.10. The fixture chart has no new defaults, reproducing + absent legacy maps rather than letting normal Helm coalescing hide the defect. +- Formatting, LOC/crypto gates, Helm lint, test lint, and diff checks passed. + +Rust validation used the exclusively leased existing root target, offline and +locked, with incremental compilation disabled. No duplicate target directory, +dependency installation, live deployment, or external push was performed. + +Local validation at implementation commit `df0540c3`, before foundation +integration and independent review: + +- Controller/router Rust suites: 2,119 tests passed; three existing doctests + ignored. +- New HTTP regressions cover cross-provider credentials, identical endpoints + with distinct keys, default recovery, streaming acceptance, ordinary auth + failures, and preservation of explicit client model selection. +- Controller tests cover bounded fallback schema without roster restriction, + ordered deduplication, stable label values, and precise egress behavior. +- Existing publication/CLI compatibility checks: 94 passed, two skipped. + Used the authorized existing cache with Vitest 4.1.10, not lockfile-exact + 4.1.8. Locked dependency validation remains the responsibility of GitHub CI. +- Helm lint and default, local-dev, generic, and existing-AKS renders passed. +- Strict Clippy passed for all controller/router targets with warnings denied. + +The foundation integration preserves the existing `sha256_hex` re-export used +by PR4's label helper. The repair qualification above supersedes the earlier +pre-integration build evidence for this delta. Review of PR4's +added documentation and Helm example/configuration lines found no private +hardware, registry, subscription, or resource identifiers. + +## Verdict + +The maintainer explicitly signed off and approved on 2026-09-08. This approval +covers the technically qualified source at +`45cfa009ce6b5242ddebc6dc74f0c2acba1668b1` in Azure/kars#547, not later functional +changes, other unresolved slices, a customer deployment or promotion to `main`. + +An independent person's review and sign-off remain required. This single +maintainer sign-off must not be counted twice or treated as a waiver of the +two-person capability-audit gate. + +Signed-off-by: pallakatos <191481949+pallakatos@users.noreply.github.com> + +## Explicit author waiver for integration assembly + +On 2026-09-08, Kars author `pallakatos` explicitly waived the second-person +sign-off for Azure/kars#547: "just push them and say I waived it". +This supersedes the independent-signature landing requirement above only for +assembling the already-qualified source into `Azure/kars:kars-bridge`. +It does not assert that an independent human review occurred. + +All other required technical and security gates must pass on the landing head. +The signature-only branch-protection exception and the already-authorized +account-specific review allowance must be restored immediately after the merge, +including on failure. No CI result is rewritten as successful. +The waiver does not cover functional changes beyond the qualified source, +later unresolved slices, `main` promotion, customer deployments or public +publication of the private Bridge application. diff --git a/inference-router/src/auth.rs b/inference-router/src/auth.rs index 2e08872e..4f84b581 100644 --- a/inference-router/src/auth.rs +++ b/inference-router/src/auth.rs @@ -42,6 +42,19 @@ impl Default for WorkloadIdentityAuth { } impl WorkloadIdentityAuth { + #[cfg(test)] + pub(crate) fn for_test( + api_key: Option<&str>, + sidecar: Option, + ) -> Self { + Self { + client: reqwest::Client::new(), + token_cache: Arc::new(RwLock::new(HashMap::new())), + api_key: api_key.map(str::to_string), + sidecar: sidecar.map(Arc::new), + } + } + pub fn new() -> Self { // Try to load API key from secret mount (dev mode), then env var (sub-agent) let api_key = std::fs::read_to_string("/run/secrets/azure-openai-key") diff --git a/inference-router/src/config.rs b/inference-router/src/config.rs index 3438a106..73599ca6 100644 --- a/inference-router/src/config.rs +++ b/inference-router/src/config.rs @@ -4,6 +4,25 @@ //! Configuration loaded from environment variables. use anyhow::{Context, Result}; +use std::collections::HashMap; + +/// An operator-configured OpenAI-compatible endpoint and its own credential. +#[derive(Clone, PartialEq, Eq)] +pub struct ProviderEndpoint { + pub tag: String, + pub endpoint: String, + pub api_key: Option, +} + +impl std::fmt::Debug for ProviderEndpoint { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ProviderEndpoint") + .field("tag", &self.tag) + .field("endpoint", &self.endpoint) + .field("has_api_key", &self.api_key.is_some()) + .finish() + } +} /// Registry topology mode. /// @@ -111,6 +130,7 @@ pub struct Config { /// Moderation model (`OPENAI_MODERATION_MODEL`, default /// `omni-moderation-latest`). pub openai_moderation_model: String, + pub providers: HashMap, } /// Read a credential from an env var, falling back to the standard @@ -220,6 +240,7 @@ impl Config { .ok() .filter(|s| !s.is_empty()) .unwrap_or_else(|| "omni-moderation-latest".into()), + providers: parse_providers_from_env(std::env::vars()), }) } @@ -234,10 +255,12 @@ impl Config { self.foundry_endpoint.as_deref(), self.foundry_project_endpoint.as_deref(), ]; - candidates - .iter() - .flatten() - .any(|e| e.contains("models.github.ai") || e.contains("models.inference.ai.azure.com")) + candidates.iter().flatten().any(|endpoint| { + matches!( + crate::proxy::endpoint_host(endpoint).as_deref(), + Some("models.github.ai") | Some("models.inference.ai.azure.com") + ) + }) } /// Returns true when the configured endpoint points at the GitHub @@ -264,14 +287,99 @@ impl Config { candidates .iter() .flatten() - .any(|e| e.contains("api.githubcopilot.com")) + .any(|e| crate::proxy::is_copilot_endpoint(e)) + } + + pub fn resolve_provider(&self, tag: &str) -> Option { + if let Some(provider) = self.providers.get(&tag.to_ascii_lowercase()) { + return Some(provider.clone()); + } + if tag.eq_ignore_ascii_case("github-copilot") + && std::env::var("COPILOT_GITHUB_TOKEN") + .ok() + .is_some_and(|token| !token.trim().is_empty()) + { + return Some(ProviderEndpoint { + tag: "github-copilot".into(), + endpoint: "https://api.githubcopilot.com".into(), + api_key: None, + }); + } + None } } +fn parse_providers_from_env( + vars: impl Iterator, +) -> HashMap { + let mut endpoints = HashMap::new(); + let mut keys = HashMap::new(); + let mut tokens = HashMap::new(); + for (name, value) in vars { + if value.trim().is_empty() { + continue; + } + let Some(rest) = name.strip_prefix("KARS_PROVIDER_") else { + continue; + }; + for (suffix, values) in [ + ("_ENDPOINT", &mut endpoints), + ("_API_KEY", &mut keys), + ("_TOKEN", &mut tokens), + ] { + if let Some(tag) = rest.strip_suffix(suffix) + && !tag.is_empty() + { + values.insert(tag.to_ascii_lowercase().replace('_', "-"), value.clone()); + break; + } + } + } + endpoints + .into_iter() + .map(|(tag, endpoint)| { + let api_key = keys.remove(&tag).or_else(|| tokens.remove(&tag)); + ( + tag.clone(), + ProviderEndpoint { + tag, + endpoint, + api_key, + }, + ) + }) + .collect() +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn named_providers_parse_and_prefer_explicit_keys_deterministically() { + let vars = [ + ( + "KARS_PROVIDER_GITHUB_MODELS_ENDPOINT", + "https://models.github.ai/inference", + ), + ("KARS_PROVIDER_GITHUB_MODELS_TOKEN", "token"), + ("KARS_PROVIDER_GITHUB_MODELS_API_KEY", "key"), + ("KARS_PROVIDER_EMPTY_ENDPOINT", " "), + ("UNRELATED", "ignored"), + ] + .map(|(key, value)| (key.to_string(), value.to_string())); + let parsed = parse_providers_from_env(vars.clone().into_iter()); + let reversed = parse_providers_from_env(vars.into_iter().rev()); + assert_eq!(parsed, reversed); + assert_eq!(parsed.len(), 1); + assert_eq!(parsed["github-models"].api_key.as_deref(), Some("key")); + assert!(!format!("{:?}", parsed["github-models"]).contains("\"key\"")); + let mut config = cfg(None); + config.providers = parsed; + assert!(config.resolve_provider("GITHUB-MODELS").is_some()); + assert!(config.resolve_provider("missing").is_none()); + } + fn cfg(endpoint: Option<&str>) -> Config { Config { port: 8443, @@ -293,6 +401,7 @@ mod tests { openai_moderation_endpoint: "https://api.openai.com".into(), openai_moderation_api_key: None, openai_moderation_model: "omni-moderation-latest".into(), + providers: HashMap::new(), } } diff --git a/inference-router/src/copilot_auth.rs b/inference-router/src/copilot_auth.rs index 0dba3b12..7785cba1 100644 --- a/inference-router/src/copilot_auth.rs +++ b/inference-router/src/copilot_auth.rs @@ -8,7 +8,7 @@ //! //! 1. The **GitHub OAuth/PAT** that the user supplies is *not* sent to //! the Copilot inference API. It's only sent to -//! `POST https://api.github.com/copilot_internal/v2/token`, which +//! `GET https://api.github.com/copilot_internal/v2/token`, which //! returns a short-lived **Copilot JWT** (~30 min TTL). //! 2. The Copilot JWT is what gets attached as `Authorization: Bearer` //! on every `api.githubcopilot.com` upstream request. @@ -27,6 +27,7 @@ use anyhow::{Context, Result, bail}; use serde::Deserialize; +use std::collections::HashMap; use std::sync::Arc; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; use tokio::sync::RwLock; @@ -35,6 +36,7 @@ const TOKEN_EXCHANGE_URL: &str = "https://api.github.com/copilot_internal/v2/tok /// Refresh window: ask for a new JWT this long before the cached one expires. const REFRESH_BUFFER: Duration = Duration::from_secs(60); +const DIRECT_REFRESH_SECS: u64 = 1500; /// Static integration headers Copilot expects on every request. /// Without these, Copilot returns 400 "missing required header" or, worse, @@ -74,6 +76,11 @@ struct CachedJwt { expires_at: Instant, } +struct NamedCache { + github_token: String, + cached: Arc>>, +} + /// In-process cache + exchanger for Copilot JWTs. /// /// Cheap to clone; the underlying state is `Arc>`. @@ -86,6 +93,11 @@ pub struct CopilotTokenCache { /// instead of panicking at startup. github_token: Option, cached: Arc>>, + /// Provider IDs are the keys. Credentials live only in private values, + /// never in cache keys or diagnostic output. + named: Arc>>, + #[cfg(test)] + exchange_override: Option, } impl CopilotTokenCache { @@ -121,6 +133,9 @@ impl CopilotTokenCache { .expect("failed to build reqwest client"), github_token, cached: Arc::new(RwLock::new(None)), + named: Arc::new(RwLock::new(HashMap::new())), + #[cfg(test)] + exchange_override: None, } } @@ -135,7 +150,25 @@ impl CopilotTokenCache { .expect("failed to build reqwest client"), github_token: Some(github_token.into()), cached: Arc::new(RwLock::new(None)), + named: Arc::new(RwLock::new(HashMap::new())), + #[cfg(test)] + exchange_override: None, + } + } + + #[cfg(test)] + pub(crate) fn with_test_exchange(token: &str, url: String) -> Self { + let mut cache = Self::with_token(token); + cache.exchange_override = Some(url); + cache + } + + fn exchange_url(&self) -> &str { + #[cfg(test)] + if let Some(url) = self.exchange_override.as_deref() { + return url; } + TOKEN_EXCHANGE_URL } /// True if a GitHub token is configured (i.e. Copilot path is usable). @@ -146,19 +179,70 @@ impl CopilotTokenCache { /// Returns a valid Copilot JWT, exchanging if needed. /// The base URL allows tests to point at a mock server. pub async fn get_jwt(&self) -> Result { - self.get_jwt_with_base(TOKEN_EXCHANGE_URL).await + self.get_jwt_with_base(self.exchange_url()).await } /// Same as [`get_jwt`] but allows overriding the exchange endpoint. /// Internal — exposed for tests; production code should call `get_jwt`. pub async fn get_jwt_with_base(&self, exchange_url: &str) -> Result { + self.exchange_for_scope(self.github_token.as_deref(), &self.cached, exchange_url) + .await + } + + /// Named Copilot credentials are GitHub OAuth/seat tokens, not Azure keys + /// or an already exchanged JWT. Each immutable provider identity owns a + /// separate exchange cache, independent of the legacy default account. + pub async fn get_jwt_for_provider( + &self, + provider_id: &str, + github_token: &str, + ) -> Result { + self.provider_jwt_with_base(provider_id, github_token, self.exchange_url()) + .await + } + + pub(crate) async fn provider_jwt_with_base( + &self, + provider_id: &str, + github_token: &str, + exchange_url: &str, + ) -> Result { + anyhow::ensure!( + !github_token.trim().is_empty(), + "named Copilot provider has no GitHub seat token" + ); + let cached = { + let mut providers = self.named.write().await; + let entry = providers + .entry(provider_id.to_string()) + .or_insert_with(|| NamedCache { + github_token: github_token.to_string(), + cached: Arc::new(RwLock::new(None)), + }); + if entry.github_token != github_token { + // Defensive invalidation if a future config reload reuses an ID. + entry.github_token = github_token.to_string(); + entry.cached = Arc::new(RwLock::new(None)); + } + entry.cached.clone() + }; + self.exchange_for_scope(Some(github_token), &cached, exchange_url) + .await + } + + async fn exchange_for_scope( + &self, + github_token: Option<&str>, + cached: &Arc>>, + exchange_url: &str, + ) -> Result { // Fast path: cached JWT still has runway. Serve only when BOTH // (a) GitHub's refresh window hasn't elapsed AND // (b) Copilot's hard expiry hasn't passed (minus a safety buffer). // The expires_at check is what fixes the "IDE token expired" // 401s a user hit after ~30 min on long-lived sandbox pods. { - let guard = self.cached.read().await; + let guard = cached.read().await; if let Some(c) = guard.as_ref() { let now = Instant::now(); let expiry_safe = c.expires_at > now + REFRESH_BUFFER; @@ -169,7 +253,7 @@ impl CopilotTokenCache { } // Slow path: exchange. - let gh = self.github_token.as_deref().context( + let gh = github_token.context( "no GitHub token configured for Copilot — set COPILOT_GITHUB_TOKEN or mount /run/secrets/copilot-github-token", )?; @@ -187,6 +271,20 @@ impl CopilotTokenCache { let status = resp.status(); if !status.is_success() { let body = resp.text().await.unwrap_or_default(); + // Non-editor OAuth tokens can authenticate directly even when the + // editor-only exchange endpoint rejects their scope. + if status.is_client_error() { + tracing::warn!(%status, "Copilot token exchange rejected; trying configured direct token"); + let now = Instant::now(); + let direct = CachedJwt { + token: gh.to_string(), + refresh_at: now + Duration::from_secs(DIRECT_REFRESH_SECS), + expires_at: now + Duration::from_secs(DIRECT_REFRESH_SECS + 300), + }; + let token = direct.token.clone(); + *cached.write().await = Some(direct); + return Ok(token); + } bail!("Copilot token exchange returned {status}: {body}"); } @@ -210,7 +308,7 @@ impl CopilotTokenCache { }; let now_instant = Instant::now(); - let cached = CachedJwt { + let refreshed = CachedJwt { token: parsed.token.clone(), refresh_at: now_instant + Duration::from_secs(refresh_secs as u64), // Track Copilot's hard expiry independently so the cache @@ -225,7 +323,7 @@ impl CopilotTokenCache { refresh_secs ); - *self.cached.write().await = Some(cached); + *cached.write().await = Some(refreshed); Ok(parsed.token) } } @@ -233,6 +331,67 @@ impl CopilotTokenCache { #[cfg(test)] mod tests { use super::*; + + #[tokio::test] + async fn named_cache_keys_are_provider_ids_and_rotated_credentials_invalidate_only_their_scope() + { + let server = MockServer::start().await; + let expiry = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() + + 3600; + for (seat, jwt) in [("seat-first", "jwt-first"), ("seat-rotated", "jwt-rotated")] { + Mock::given(header("authorization", format!("token {seat}"))) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "token": jwt, "expires_at": expiry, "refresh_in": 1500, + }))) + .expect(1) + .mount(&server) + .await; + } + let cache = CopilotTokenCache::with_token("unrelated-default-seat"); + assert_eq!( + cache + .provider_jwt_with_base("provider-a", "seat-first", &server.uri()) + .await + .unwrap(), + "jwt-first" + ); + assert_eq!( + cache + .provider_jwt_with_base("provider-a", "seat-first", &server.uri()) + .await + .unwrap(), + "jwt-first" + ); + assert_eq!( + cache + .provider_jwt_with_base("provider-a", "seat-rotated", &server.uri()) + .await + .unwrap(), + "jwt-rotated" + ); + let keys: Vec<_> = cache.named.read().await.keys().cloned().collect(); + assert_eq!(keys, ["provider-a"]); + assert!(cache.cached.read().await.is_none()); + } + + #[tokio::test] + async fn direct_token_fallback_preserves_non_editor_oauth_support() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/copilot_internal/v2/token")) + .respond_with(ResponseTemplate::new(404)) + .mount(&server) + .await; + let cache = CopilotTokenCache::with_token("gho_direct"); + let token = cache + .get_jwt_with_base(&format!("{}/copilot_internal/v2/token", server.uri())) + .await + .unwrap(); + assert_eq!(token, "gho_direct"); + } use serde_json::json; use std::time::{SystemTime, UNIX_EPOCH}; use wiremock::matchers::{header, method, path}; @@ -251,6 +410,8 @@ mod tests { client: reqwest::Client::new(), github_token: None, cached: Arc::new(RwLock::new(None)), + named: Arc::new(RwLock::new(HashMap::new())), + exchange_override: None, }; let err = c.get_jwt().await.unwrap_err(); assert!(err.to_string().contains("no GitHub token")); @@ -288,7 +449,7 @@ mod tests { let server = MockServer::start().await; Mock::given(method("GET")) .and(path("/copilot_internal/v2/token")) - .respond_with(ResponseTemplate::new(401).set_body_string("bad credentials")) + .respond_with(ResponseTemplate::new(503).set_body_string("service unavailable")) .mount(&server) .await; @@ -298,7 +459,7 @@ mod tests { .await .unwrap_err(); let msg = err.to_string(); - assert!(msg.contains("401"), "expected 401 in error, got: {msg}"); + assert!(msg.contains("503"), "expected 503 in error, got: {msg}"); } /// Regression: when GitHub returns `refresh_in` LARGER than the actual diff --git a/inference-router/src/failover.rs b/inference-router/src/failover.rs index 7fa2245c..1bf91615 100644 --- a/inference-router/src/failover.rs +++ b/inference-router/src/failover.rs @@ -1,13 +1,12 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -//! Slice 2d.2 — health-aware deployment failover. +//! Health-aware provider and deployment failover. //! //! Wraps [`crate::proxy::forward`] with a candidate-walk that honours //! `InferencePolicy.spec.modelPreference.{primary,fallback[]}.deployment`. -//! Same-provider only — the router still holds a single Foundry/AOAI -//! client at process start (`UpstreamConfig.endpoint`); we only swap -//! the `deployment` field per attempt. +//! Every candidate is resolved from the original default endpoint, never the +//! preceding candidate, so fallback cannot retain another provider's credentials. //! //! Per-attempt outcome feeds [`DeploymentHealthRegistry`]: //! * 2xx ⇒ `record_success` (clears any streak) @@ -30,10 +29,71 @@ use reqwest::Client; use std::sync::Arc; use crate::auth::WorkloadIdentityAuth; +use crate::config::Config; use crate::copilot_auth::CopilotTokenCache; use crate::deployment_health::DeploymentHealthRegistry; use crate::inference_policy_loader::{InferencePolicySnapshot, ModelRef}; -use crate::proxy::{UpstreamConfig, forward}; +use crate::proxy::failure::{ForwardFailure, retryable_failure}; +use crate::proxy::{AuthenticationProvenance, UpstreamConfig, forward}; + +mod stream; +pub use stream::forward_stream_with_failover; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Candidate { + pub provider: Option, + pub deployment: String, + pub routing_intent: RoutingIntent, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RoutingIntent { + /// Historical primary-model metadata routes only when a named endpoint + /// was separately registered. Native key presence is not routing intent. + Metadata, + /// An authoritative policy provider or a newly declared fallback route. + Explicit, +} + +impl RoutingIntent { + pub(crate) fn primary(policy: &InferencePolicySnapshot) -> Self { + if policy + .provider + .as_ref() + .is_some_and(|provider| !provider.trim().is_empty()) + { + Self::Explicit + } else { + Self::Metadata + } + } +} + +// Native families and implicit Copilot selection can change destination with +// intent alone; their metadata route must not poison a native route's health. +fn intent_can_change_backend(provider: Option<&str>) -> bool { + let Some(provider) = provider else { + return false; + }; + match crate::provider::parse_tag(provider) { + Ok(Some(kind)) => kind != crate::provider::ProviderKind::AzureOpenAI, + Err(_) => true, + Ok(None) => provider.trim().eq_ignore_ascii_case("github-copilot"), + } +} + +fn health_key(candidate: &Candidate) -> String { + match candidate.provider.as_deref().filter(|tag| !tag.is_empty()) { + Some(provider) + if candidate.routing_intent == RoutingIntent::Metadata + && intent_can_change_backend(Some(provider)) => + { + format!("metadata:{provider}::{}", candidate.deployment) + } + Some(provider) => format!("{provider}::{}", candidate.deployment), + None => candidate.deployment.clone(), + } +} /// Decide whether an upstream response status is a *retry-worthy* /// failure that should mark the deployment unhealthy and trigger a @@ -44,8 +104,7 @@ use crate::proxy::{UpstreamConfig, forward}; /// must be a deliberate change, not an accident. #[must_use] pub fn is_failover_trigger(status: StatusCode) -> bool { - let code = status.as_u16(); - code == 429 || (500..=599).contains(&code) + crate::proxy::failure::retryable_rejection(status) } /// Build the ordered candidate list the failover walk will try. @@ -55,46 +114,185 @@ pub fn is_failover_trigger(status: StatusCode) -> bool { /// `upstream.deployment` is returned as a single-element list, so the /// caller always has at least one attempt to make. /// -/// Deduplicates while preserving order: if `primary.deployment` and -/// `fallback[0].deployment` happen to be the same, we only try it -/// once. Empty strings are skipped. +/// Deduplicates provider/deployment pairs while preserving order, except when +/// metadata and explicit intent can select different native/default backends. #[must_use] pub fn build_candidates( upstream: &UpstreamConfig, snapshot: &InferencePolicySnapshot, -) -> Vec { - let mut out: Vec = Vec::new(); - let mut push = |dep: &str| { +) -> Vec { + let mut out: Vec = Vec::new(); + let mut push = |dep: &str, provider: Option, routing_intent: RoutingIntent| { if dep.is_empty() { return; } - if !out.iter().any(|d| d == dep) { - out.push(dep.to_string()); + if !out.iter().any(|d| { + d.deployment == dep + && d.provider == provider + && (d.routing_intent == routing_intent + || !intent_can_change_backend(provider.as_deref())) + }) { + out.push(Candidate { + deployment: dep.to_string(), + provider, + routing_intent, + }); } }; if let Some(ref pref) = snapshot.model_preference { - push(&pref.primary.deployment); - for ModelRef { deployment, .. } in &pref.fallback { - push(deployment); + push( + &pref.primary.deployment, + snapshot + .provider + .clone() + .filter(|p| !p.trim().is_empty()) + .or_else(|| Some(pref.primary.provider.clone()).filter(|p| !p.trim().is_empty())), + RoutingIntent::primary(snapshot), + ); + for ModelRef { + deployment, + provider, + } in &pref.fallback + { + push( + deployment, + Some(provider.clone()).filter(|p| !p.is_empty()), + RoutingIntent::Explicit, + ); } + } else if let Some(provider) = snapshot.provider.as_ref().filter(|p| !p.trim().is_empty()) { + push( + &upstream.deployment, + Some(provider.clone()), + RoutingIntent::Explicit, + ); } // Always keep the env-driven default as a final safety net so a // mid-flight policy unload (or a policy with only an empty // primary) never produces a zero-candidate list. - push(&upstream.deployment); + push(&upstream.deployment, None, RoutingIntent::Explicit); if out.is_empty() { // Theoretically unreachable (`upstream.deployment` is set // from `Config::default_model` which has its own default), // but defence-in-depth: empty list ⇒ one attempt at the // caller-supplied upstream as-is. - out.push(upstream.deployment.clone()); + out.push(Candidate { + provider: None, + deployment: upstream.deployment.clone(), + routing_intent: RoutingIntent::Explicit, + }); } out } +pub(crate) fn resolve_candidate( + base: &UpstreamConfig, + config: &Config, + candidate: &Candidate, +) -> std::result::Result { + let mut upstream = base.clone(); + upstream.deployment = candidate.deployment.clone(); + let Some(tag) = candidate.provider.as_deref() else { + return Ok(upstream); + }; + let identity = AuthenticationProvenance::Named { + provider_id: tag.trim().to_ascii_lowercase(), + }; + if candidate.routing_intent == RoutingIntent::Metadata { + if let Some(target) = config.providers.get(&tag.trim().to_ascii_lowercase()) { + upstream.endpoint = target.endpoint.clone(); + upstream.provider = crate::provider::ProviderKind::AzureOpenAI; + upstream.api_key = None; + upstream.provider_api_key = target.api_key.clone(); + upstream.authentication = identity; + } + return Ok(upstream); + } + match crate::provider::parse_tag(tag)? { + Some(crate::provider::ProviderKind::Anthropic) => { + if let crate::provider::ProviderTarget::Anthropic { endpoint, api_key } = + crate::provider::resolve(Some(tag), config)? + { + upstream.endpoint = endpoint; + upstream.provider = crate::provider::ProviderKind::Anthropic; + upstream.api_key = Some(api_key); + upstream.provider_api_key = None; + upstream.authentication = identity; + } + } + Some(crate::provider::ProviderKind::Ollama) => { + if let crate::provider::ProviderTarget::Ollama { endpoint } = + crate::provider::resolve(Some(tag), config)? + { + upstream.endpoint = endpoint; + upstream.provider = crate::provider::ProviderKind::Ollama; + upstream.api_key = None; + upstream.provider_api_key = None; + upstream.authentication = identity; + } + } + _ => { + if let Some(target) = config.resolve_provider(tag.trim()) { + // An informational tag on the unchanged legacy Copilot default + // does not manufacture a second named account. + if !config + .providers + .contains_key(&tag.trim().to_ascii_lowercase()) + && crate::proxy::is_copilot_endpoint(&base.endpoint) + && crate::proxy::is_copilot_endpoint(&target.endpoint) + { + return Ok(upstream); + } + upstream.endpoint = target.endpoint; + upstream.provider = crate::provider::ProviderKind::AzureOpenAI; + upstream.api_key = None; + upstream.provider_api_key = target.api_key; + upstream.authentication = identity; + } + } + } + Ok(upstream) +} + +pub(crate) fn candidates_for_request( + base: &UpstreamConfig, + policy: &InferencePolicySnapshot, + body: &[u8], +) -> Vec { + let mut candidates = build_candidates(base, policy); + if policy.model_preference.is_none() + && let Ok(value) = serde_json::from_slice::(body) + && let Some(model) = value + .get("model") + .and_then(|model| model.as_str()) + .filter(|model| !model.trim().is_empty()) + && model != candidates[0].deployment + { + // No policy-selected model: preserve the pre-existing public API's + // caller-selected model, then retain the configured safety net. + let mut requested = candidates[0].clone(); + requested.deployment = model.to_string(); + candidates.insert(0, requested); + } + candidates +} + +fn request_body_for_candidate(body: &Bytes, deployment: &str) -> Bytes { + let Ok(mut value) = serde_json::from_slice::(body) else { + return body.clone(); + }; + let Some(object) = value.as_object_mut() else { + return body.clone(); + }; + object.insert("model".into(), deployment.into()); + serde_json::to_vec(&value) + .map(Bytes::from) + .unwrap_or_else(|_| body.clone()) +} + /// Walks `build_candidates(...)`, skipping deployments the health /// cache currently flags as unhealthy, and returns the first /// successful (or non-retryable) response. If every candidate either @@ -113,29 +311,31 @@ pub async fn forward_with_failover( client: &Client, health: &Arc, upstream_base: &UpstreamConfig, + config: &Config, snapshot: &InferencePolicySnapshot, method: Method, path: &str, request_headers: &HeaderMap, request_body: Bytes, -) -> Result<(StatusCode, HeaderMap, Bytes)> { - let candidates = build_candidates(upstream_base, snapshot); +) -> Result<(StatusCode, HeaderMap, Bytes, UpstreamConfig)> { + let candidates = candidates_for_request(upstream_base, snapshot, &request_body); // Track the last *actually attempted* response so we can surface // a real upstream error if every candidate fails. - let mut last_result: Option> = None; + let mut last_result = None; // The very first candidate (regardless of health) — used as a // fallback-of-last-resort when every candidate was skipped // because the cache flagged them all unhealthy. let first_candidate = candidates .first() .cloned() - .unwrap_or_else(|| upstream_base.deployment.clone()); + .expect("candidate list is non-empty"); - for (idx, deployment) in candidates.iter().enumerate() { + for (idx, candidate) in candidates.iter().enumerate() { + let deployment = health_key(candidate); // Skip unhealthy candidates *unless* this is the only one // we have left to try (i.e. we've exhausted the list). - if !health.is_healthy(deployment) { + if !health.is_healthy(&deployment) { tracing::info!( sandbox = %upstream_base.sandbox_name, deployment = %deployment, @@ -144,13 +344,13 @@ pub async fn forward_with_failover( continue; } - let mut upstream = upstream_base.clone(); - upstream.deployment = deployment.clone(); + let upstream = resolve_candidate(upstream_base, config, candidate) + .map_err(ForwardFailure::configuration)?; if idx > 0 { tracing::warn!( sandbox = %upstream_base.sandbox_name, - from = %first_candidate, + from = %health_key(&first_candidate), to = %deployment, attempt = idx + 1, digest = %snapshot.digest, @@ -166,13 +366,13 @@ pub async fn forward_with_failover( method.clone(), path, request_headers, - request_body.clone(), + request_body_for_candidate(&request_body, &candidate.deployment), ) .await; match &attempt { Ok((status, _, _)) if is_failover_trigger(*status) => { - health.record_failure(deployment); + health.record_failure(&deployment); tracing::warn!( sandbox = %upstream_base.sandbox_name, deployment = %deployment, @@ -180,27 +380,32 @@ pub async fn forward_with_failover( digest = %snapshot.digest, "InferencePolicy failover: upstream returned retry-worthy status" ); - last_result = Some(attempt); + last_result = + Some(attempt.map(|(status, headers, body)| (status, headers, body, upstream))); continue; } Ok((status, _, _)) => { if status.is_success() { - health.record_success(deployment); + health.record_success(&deployment); } - return attempt; + return attempt.map(|(status, headers, body)| (status, headers, body, upstream)); } - Err(e) => { - health.record_failure(deployment); + Err(e) if retryable_failure(e) => { + health.record_failure(&deployment); tracing::warn!( sandbox = %upstream_base.sandbox_name, deployment = %deployment, error = %format!("{e:#}"), digest = %snapshot.digest, - "InferencePolicy failover: transport error" + "InferencePolicy failover: retryable upstream failure" ); - last_result = Some(attempt); + last_result = + Some(attempt.map(|(status, headers, body)| (status, headers, body, upstream))); continue; } + Err(_) => { + return attempt.map(|(status, headers, body)| (status, headers, body, upstream)); + } } } @@ -215,12 +420,12 @@ pub async fn forward_with_failover( // error that hides the real cause. tracing::warn!( sandbox = %upstream_base.sandbox_name, - deployment = %first_candidate, + deployment = %health_key(&first_candidate), digest = %snapshot.digest, "InferencePolicy failover: all candidates unhealthy, retrying primary anyway" ); - let mut upstream = upstream_base.clone(); - upstream.deployment = first_candidate.clone(); + let upstream = resolve_candidate(upstream_base, config, &first_candidate) + .map_err(ForwardFailure::configuration)?; let attempt = forward( auth, copilot, @@ -229,18 +434,22 @@ pub async fn forward_with_failover( method, path, request_headers, - request_body, + request_body_for_candidate(&request_body, &first_candidate.deployment), ) .await; match &attempt { - Ok((status, _, _)) if status.is_success() => health.record_success(&first_candidate), + Ok((status, _, _)) if status.is_success() => { + health.record_success(&health_key(&first_candidate)) + } Ok((status, _, _)) if is_failover_trigger(*status) => { - health.record_failure(&first_candidate); + health.record_failure(&health_key(&first_candidate)); + } + Err(error) if retryable_failure(error) => { + health.record_failure(&health_key(&first_candidate)) } - Err(_) => health.record_failure(&first_candidate), _ => {} } - attempt + attempt.map(|(status, headers, body)| (status, headers, body, upstream)) } #[cfg(test)] @@ -248,6 +457,13 @@ mod tests { use super::*; use crate::inference_policy_loader::{ModelPreference, ModelRef}; + fn deployments(candidates: &[Candidate]) -> Vec<&str> { + candidates + .iter() + .map(|candidate| candidate.deployment.as_str()) + .collect() + } + fn upstream(dep: &str) -> UpstreamConfig { UpstreamConfig { endpoint: "https://example.openai.azure.com".into(), @@ -255,6 +471,8 @@ mod tests { sandbox_name: "sbx".into(), provider: crate::provider::ProviderKind::AzureOpenAI, api_key: None, + provider_api_key: None, + authentication: AuthenticationProvenance::LegacyDefault, } } @@ -305,28 +523,56 @@ mod tests { fn build_candidates_includes_primary_then_fallback_chain() { let snap = snapshot_with("primary", &["fb-a", "fb-b"]); let c = build_candidates(&upstream("default"), &snap); - assert_eq!(c, vec!["primary", "fb-a", "fb-b", "default"]); + assert_eq!(deployments(&c), vec!["primary", "fb-a", "fb-b", "default"]); } #[test] fn build_candidates_dedups_overlap() { let snap = snapshot_with("primary", &["primary", "fb-a"]); let c = build_candidates(&upstream("primary"), &snap); - assert_eq!(c, vec!["primary", "fb-a"]); + assert_eq!(deployments(&c), vec!["primary", "fb-a", "primary"]); + } + + #[test] + fn legacy_metadata_and_explicit_native_fallback_are_distinct_routing_intents() { + let snapshot = InferencePolicySnapshot { + model_preference: Some(ModelPreference { + primary: ModelRef { + provider: "anthropic".into(), + deployment: "claude-prod".into(), + }, + fallback: vec![ + ModelRef { + provider: "anthropic".into(), + deployment: "claude-prod".into(), + }, + ModelRef { + provider: "anthropic".into(), + deployment: "claude-prod".into(), + }, + ], + }), + ..Default::default() + }; + let candidates = build_candidates(&upstream("default"), &snapshot); + assert_eq!(candidates.len(), 3); + assert_eq!(candidates[0].routing_intent, RoutingIntent::Metadata); + assert_eq!(candidates[1].routing_intent, RoutingIntent::Explicit); + assert_eq!(candidates[2].provider, None); } #[test] fn build_candidates_skips_empty_deployment_strings() { let snap = snapshot_with("", &["", "fb-a"]); let c = build_candidates(&upstream("default"), &snap); - assert_eq!(c, vec!["fb-a", "default"]); + assert_eq!(deployments(&c), vec!["fb-a", "default"]); } #[test] fn build_candidates_no_policy_yields_just_default() { let snap = InferencePolicySnapshot::default(); let c = build_candidates(&upstream("env-default"), &snap); - assert_eq!(c, vec!["env-default"]); + assert_eq!(deployments(&c), vec!["env-default"]); } #[test] @@ -334,6 +580,6 @@ mod tests { // Even with everything blank, we get a one-element list. let snap = snapshot_with("", &[]); let c = build_candidates(&upstream(""), &snap); - assert_eq!(c, vec![""]); + assert_eq!(deployments(&c), vec![""]); } } diff --git a/inference-router/src/failover/stream.rs b/inference-router/src/failover/stream.rs new file mode 100644 index 00000000..85acffd9 --- /dev/null +++ b/inference-router/src/failover/stream.rs @@ -0,0 +1,86 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use futures::{StreamExt, TryStreamExt, stream::BoxStream}; + +type StreamingResult = ( + StatusCode, + HeaderMap, + BoxStream<'static, Result>, + UpstreamConfig, +); + +/// Fail over only before a successful response is accepted. A successful stream +/// is never replayed, even if its first body chunk subsequently fails. +#[allow(clippy::too_many_arguments)] +pub async fn forward_stream_with_failover( + auth: Arc, + copilot: Option>, + client: Client, + health: &Arc, + upstream_base: &UpstreamConfig, + config: &Config, + snapshot: &InferencePolicySnapshot, + path: &str, + request_headers: HeaderMap, + request_body: Bytes, +) -> Result { + let candidates = candidates_for_request(upstream_base, snapshot, &request_body); + let mut eligible: Vec<_> = candidates + .iter() + .filter(|candidate| health.is_healthy(&health_key(candidate))) + .collect(); + if eligible.is_empty() { + eligible.push(&candidates[0]); + } + let mut last_result = None; + for candidate in eligible { + let key = health_key(candidate); + let upstream = resolve_candidate(upstream_base, config, candidate) + .map_err(ForwardFailure::configuration)?; + let body = request_body_for_candidate(&request_body, &candidate.deployment); + let attempt = crate::proxy::forward_stream( + auth.clone(), + copilot.clone(), + client.clone(), + upstream.clone(), + path, + request_headers.clone(), + body, + ) + .await; + match attempt { + Ok((status, headers, stream)) if is_failover_trigger(status) => { + health.record_failure(&key); + tracing::warn!(provider = %key, status = %status, digest = %snapshot.digest, "streaming inference failover"); + let buffered = stream + .try_fold(Vec::new(), |mut bytes, chunk| async move { + bytes.extend_from_slice(&chunk); + Ok(bytes) + }) + .await; + last_result = Some(match buffered { + Ok(bytes) => { + let stream = + futures::stream::once(async move { Ok(Bytes::from(bytes)) }).boxed(); + Ok((status, headers, stream, upstream)) + } + Err(error) => Err(error.into()), + }); + } + Ok((status, headers, stream)) => { + if status.is_success() { + health.record_success(&key); + } + return Ok((status, headers, stream, upstream)); + } + Err(error) if retryable_failure(&error) => { + health.record_failure(&key); + last_result = Some(Err(error)); + } + Err(error) => return Err(error), + } + } + last_result.expect("at least one candidate is attempted") +} diff --git a/inference-router/src/provider.rs b/inference-router/src/provider.rs index 06ea713d..06e9fc93 100644 --- a/inference-router/src/provider.rs +++ b/inference-router/src/provider.rs @@ -190,6 +190,7 @@ mod tests { openai_moderation_endpoint: "https://api.openai.com".into(), openai_moderation_api_key: None, openai_moderation_model: "omni-moderation-latest".into(), + providers: std::collections::HashMap::new(), } } diff --git a/inference-router/src/proxy.rs b/inference-router/src/proxy.rs index 111ad736..0d40e7ed 100644 --- a/inference-router/src/proxy.rs +++ b/inference-router/src/proxy.rs @@ -11,7 +11,7 @@ use anyhow::{Context, Result}; use axum::http::{HeaderMap, HeaderValue, Method, StatusCode}; use bytes::Bytes; use reqwest::Client; -use std::time::Instant; +use std::time::{Duration, Instant}; use crate::auth::WorkloadIdentityAuth; use crate::copilot_auth::{ @@ -21,6 +21,16 @@ use crate::metrics; use crate::provider::ProviderKind; use std::sync::Arc; +const INFERENCE_REQUEST_TIMEOUT: Duration = Duration::from_secs(600); + +mod authentication; +pub mod failure; +pub use authentication::{AuthenticationProvenance, credential_for_upstream, token_for_endpoint}; +use failure::ForwardFailure; + +#[cfg(test)] +mod authentication_tests; + /// Upstream configuration for a single request. #[derive(Clone)] pub struct UpstreamConfig { @@ -35,6 +45,9 @@ pub struct UpstreamConfig { /// Filled by `provider::resolve` from router-side config only — /// never from the inbound request. pub api_key: Option, + /// Credential for this named provider only; never inherited across providers. + pub provider_api_key: Option, + pub authentication: AuthenticationProvenance, } impl UpstreamConfig { @@ -48,6 +61,8 @@ impl UpstreamConfig { sandbox_name, provider: ProviderKind::AzureOpenAI, api_key: None, + provider_api_key: None, + authentication: AuthenticationProvenance::LegacyDefault, } } } @@ -156,65 +171,31 @@ fn build_upstream_headers( /// compatible *but* requires its own short-lived JWT (exchanged from the /// user's GitHub OAuth token) and three static integration headers. pub fn is_copilot_endpoint(endpoint: &str) -> bool { - endpoint.contains("api.githubcopilot.com") + endpoint_host(endpoint).as_deref() == Some("api.githubcopilot.com") } -/// Acquire the right auth token for a given upstream endpoint. -/// -/// - GitHub Copilot endpoints → exchanged Copilot JWT (cached, refreshed proactively). -/// - Everything else → Azure auth (API key in dev mode, WI/IMDS in AKS mode). -/// -/// Returning `Result` lets the caller surface a clean 502 if the -/// Copilot token cache is uninitialised or the GitHub token is missing — -/// rather than panicking inside `forward()`. -pub async fn token_for_endpoint( - auth: &WorkloadIdentityAuth, - copilot: Option<&CopilotTokenCache>, - endpoint: &str, -) -> Result { - if is_copilot_endpoint(endpoint) { - match copilot { - Some(cache) => cache.get_jwt().await, - None => anyhow::bail!( - "Copilot endpoint configured but no CopilotTokenCache available — \ - set COPILOT_GITHUB_TOKEN or mount /run/secrets/copilot-github-token" - ), - } - } else { - auth.get_token(token_audience(endpoint)).await - } +pub(crate) fn endpoint_host(endpoint: &str) -> Option { + reqwest::Url::parse(endpoint) + .ok() + .and_then(|url| url.host_str().map(str::to_ascii_lowercase)) } -/// Provider-aware credential resolution for a single upstream request. -/// -/// - `AzureOpenAI` → the historic [`token_for_endpoint`] path (Azure -/// WI/IMDS, API key, or Copilot JWT depending on endpoint). -/// - `Anthropic` → the static API key `provider::resolve` copied from -/// router-side config onto `UpstreamConfig.api_key`. Its absence -/// here is a programmer error (resolution fails closed earlier), -/// surfaced as a clean 502 rather than a panic. -/// - `Ollama` → no credential. -pub async fn credential_for_upstream( - auth: &WorkloadIdentityAuth, - copilot: Option<&CopilotTokenCache>, - upstream: &UpstreamConfig, -) -> Result { - match upstream.provider { - ProviderKind::AzureOpenAI => token_for_endpoint(auth, copilot, &upstream.endpoint) - .await - .map(UpstreamCredential::Bearer), - ProviderKind::Anthropic => upstream - .api_key - .clone() - .map(UpstreamCredential::AnthropicApiKey) - .ok_or_else(|| { - anyhow::anyhow!( - "Anthropic upstream selected but no API key on UpstreamConfig — \ - provider resolution must run before forward()" - ) - }), - ProviderKind::Ollama => Ok(UpstreamCredential::None), - } +pub(crate) fn is_azure_ai_host(host: &str) -> bool { + [ + ".openai.azure.com", + ".cognitiveservices.azure.com", + ".services.ai.azure.com", + ".openai.azure.us", + ".cognitiveservices.azure.us", + ".openai.azure.cn", + ".cognitiveservices.azure.cn", + ] + .iter() + .any(|suffix| host.ends_with(suffix)) +} + +pub(crate) fn is_local_inference_host(host: &str) -> bool { + host.ends_with(".kars-local-inference.svc.cluster.local") } /// Record Prometheus metrics from a completed request. @@ -275,7 +256,8 @@ pub async fn forward( ) -> Result<(StatusCode, HeaderMap, Bytes)> { let start = Instant::now(); - let (upstream_url, body) = build_upstream_url(auth, upstream, path, request_body)?; + let (upstream_url, body) = build_upstream_url(auth, upstream, path, request_body) + .map_err(ForwardFailure::configuration)?; let mode = match upstream.provider { ProviderKind::Anthropic => "anthropic", @@ -294,9 +276,10 @@ pub async fn forward( let credential = credential_for_upstream(auth, copilot, upstream) .await - .context("Failed to acquire auth token")?; + .map_err(ForwardFailure::authentication)?; - let headers = build_upstream_headers(request_headers, auth, &credential, &upstream.endpoint)?; + let headers = build_upstream_headers(request_headers, auth, &credential, &upstream.endpoint) + .map_err(ForwardFailure::configuration)?; tracing::info!(sandbox = %upstream.sandbox_name, url = %upstream_url, body_len = body.len(), "Sending upstream request"); @@ -332,7 +315,7 @@ pub async fn forward( let response_body = response .bytes() .await - .context("Failed to read Foundry response")?; + .map_err(|error| ForwardFailure::response_body(status, error))?; let latency = start.elapsed(); record_metrics(upstream, status, latency, &response_body); @@ -427,7 +410,7 @@ async fn send_with_retry( .request(method.clone(), url) .headers(headers.clone()) .body(body.clone()) - .timeout(std::time::Duration::from_secs(120)) + .timeout(INFERENCE_REQUEST_TIMEOUT) .send() .await; @@ -467,10 +450,12 @@ async fn send_with_retry( BACKOFF_MS[(attempt - 1) as usize], )) .await; - last_err = Some(anyhow::Error::from(err)); + last_err = Some(ForwardFailure::transport(err)); continue; } - return Err(anyhow::Error::from(err).context("Foundry upstream request failed")); + return Err( + ForwardFailure::transport(err).context("Foundry upstream request failed") + ); } } } @@ -515,14 +500,16 @@ pub async fn forward_stream( } else { inject_stream_usage(request_body) }; - let (upstream_url, body) = build_upstream_url(&auth, &upstream, path, body_with_usage)?; + let (upstream_url, body) = build_upstream_url(&auth, &upstream, path, body_with_usage) + .map_err(ForwardFailure::configuration)?; tracing::info!(sandbox = %upstream.sandbox_name, model = %upstream.deployment, mode = "stream", "Forwarding SSE stream"); let credential = credential_for_upstream(&auth, copilot.as_deref(), &upstream) .await - .context("Failed to acquire auth token")?; - let headers = build_upstream_headers(&request_headers, &auth, &credential, &upstream.endpoint)?; + .map_err(ForwardFailure::authentication)?; + let headers = build_upstream_headers(&request_headers, &auth, &credential, &upstream.endpoint) + .map_err(ForwardFailure::configuration)?; let start = Instant::now(); @@ -530,10 +517,10 @@ pub async fn forward_stream( .post(&upstream_url) .headers(headers) .body(body) - .timeout(std::time::Duration::from_secs(120)) + .timeout(INFERENCE_REQUEST_TIMEOUT) .send() .await - .context("Streaming upstream request failed")?; + .map_err(ForwardFailure::transport)?; let status = StatusCode::from_u16(response.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY); @@ -573,7 +560,10 @@ pub async fn forward_stream( // see "status=413" and have to guess at causes (token cap? bytes cap? // schema break?). Cap at 4 KiB so a misbehaving upstream can't blow logs. if !status.is_success() { - let body_bytes = response.bytes().await.unwrap_or_default(); + let body_bytes = response + .bytes() + .await + .map_err(|error| ForwardFailure::response_body(status, error))?; let preview = String::from_utf8_lossy(&body_bytes); let preview_trimmed: String = preview.chars().take(2048).collect(); tracing::warn!( @@ -664,16 +654,8 @@ fn inject_stream_usage(body: Bytes) -> Bytes { body } -/// Returns true if the endpoint is a GitHub Models endpoint -/// (https://models.github.ai/inference or the legacy -/// https://models.inference.ai.azure.com URL). GitHub Models is OpenAI-API -/// compatible but does NOT use the Azure `/openai/v1/` URL prefix. -fn is_github_models_endpoint(endpoint: &str) -> bool { - endpoint.contains("models.github.ai") || endpoint.contains("models.inference.ai.azure.com") -} - /// Build the upstream URL and optionally inject model into request body. -/// Uses the unified /openai/v1/ format — works with both API-key and Entra auth. +/// Azure hosts use /openai/v1/; custom OpenAI-compatible endpoints retain their base path. /// /// Routing rules: /// - Anthropic: no path rewrite — callers pass Messages-API paths @@ -703,9 +685,7 @@ fn build_upstream_url( path.trim_start_matches('/').trim_start_matches("v1/"), ), ProviderKind::AzureOpenAI => { - if is_github_models_endpoint(&upstream.endpoint) - || is_copilot_endpoint(&upstream.endpoint) - { + if !endpoint_host(&upstream.endpoint).is_some_and(|host| is_azure_ai_host(&host)) { format!( "{}/{}", upstream.endpoint.trim_end_matches('/'), @@ -722,8 +702,10 @@ fn build_upstream_url( }; let body = if let Ok(mut body_json) = serde_json::from_slice::(&request_body) { - if body_json.get("model").is_none() { - body_json.as_object_mut().unwrap().insert( + if body_json.get("model").is_none() + && let Some(object) = body_json.as_object_mut() + { + object.insert( "model".into(), serde_json::Value::String(upstream.deployment.clone()), ); @@ -746,8 +728,7 @@ fn build_upstream_url( // for output and Azure rejects it on the input side anyway. if path.trim_start_matches('/').starts_with("responses") && upstream.provider == ProviderKind::AzureOpenAI - && !is_github_models_endpoint(&upstream.endpoint) - && !is_copilot_endpoint(&upstream.endpoint) + && endpoint_host(&upstream.endpoint).is_some_and(|host| is_azure_ai_host(&host)) && let Some(obj) = body_json.as_object_mut() { if let Some(inputs) = obj.get_mut("input").and_then(|v| v.as_array_mut()) { @@ -758,6 +739,7 @@ fn build_upstream_url( include.retain(|s| s.as_str() != Some("reasoning.encrypted_content")); } } + rewrite_unsupported_thinking(&mut body_json); serde_json::to_vec(&body_json)?.into() } else { request_body @@ -765,6 +747,36 @@ fn build_upstream_url( Ok((url, body)) } +pub(crate) fn model_requires_adaptive_thinking(model: &str) -> bool { + let model = model.to_ascii_lowercase().replace(['.', '_'], "-"); + [ + "opus-4-8", + "opus-4-7", + "sonnet-5", + "fable-5", + "mythos-5", + "mythos-preview", + ] + .iter() + .any(|family| model.contains(family)) +} + +pub(crate) fn rewrite_unsupported_thinking(body: &mut serde_json::Value) { + if !body + .get("model") + .and_then(|v| v.as_str()) + .is_some_and(model_requires_adaptive_thinking) + { + return; + } + if let Some(thinking) = body.get_mut("thinking").and_then(|v| v.as_object_mut()) + && thinking.get("type").and_then(|v| v.as_str()) == Some("enabled") + { + thinking.clear(); + thinking.insert("type".into(), "adaptive".into()); + } +} + // ── Retry-logic unit tests (R3) ────────────────────────────────────────────── // // Full retry behaviour is exercised end-to-end in @@ -784,6 +796,96 @@ mod retry_tests { assert!(is_idempotent(&Method::HEAD, "/anything")); } + #[cfg(test)] + mod provider_routing_tests { + use super::super::*; + use serde_json::json; + + #[test] + fn host_classification_rejects_substring_spoofing() { + assert!(is_copilot_endpoint("https://api.githubcopilot.com")); + for endpoint in [ + "https://api.githubcopilot.com.evil.example", + "https://evil.example/api.githubcopilot.com", + "https://evil.example?api.githubcopilot.com", + ] { + assert!(!is_copilot_endpoint(endpoint)); + } + assert!(is_azure_ai_host("account.openai.azure.com")); + assert!(!is_azure_ai_host("account.openai.azure.com.evil.example")); + assert!(is_local_inference_host( + "model.kars-local-inference.svc.cluster.local" + )); + assert!(!is_local_inference_host( + "model.kars-local-inference.svc.cluster.local.evil.example" + )); + } + + #[tokio::test] + async fn local_service_never_receives_credentials_even_when_configured() { + let mut upstream = UpstreamConfig::azure( + "http://model.kars-local-inference.svc.cluster.local/v1".into(), + "model".into(), + "sandbox".into(), + ); + upstream.provider_api_key = Some("must-not-leak".into()); + let credential = credential_for_upstream(&WorkloadIdentityAuth::new(), None, &upstream) + .await + .unwrap(); + assert!(matches!(credential, UpstreamCredential::None)); + } + + #[test] + fn azure_and_custom_urls_keep_their_own_api_prefix() { + for (endpoint, expected) in [ + ( + "https://account.openai.azure.com", + "https://account.openai.azure.com/openai/v1/chat/completions", + ), + ( + "http://model.kars-local-inference.svc.cluster.local/v1", + "http://model.kars-local-inference.svc.cluster.local/v1/chat/completions", + ), + ( + "https://custom.example/v1", + "https://custom.example/v1/chat/completions", + ), + ( + "https://models.github.ai/inference", + "https://models.github.ai/inference/chat/completions", + ), + ] { + let upstream = + UpstreamConfig::azure(endpoint.into(), "model".into(), "sandbox".into()); + let (url, _) = build_upstream_url( + &WorkloadIdentityAuth::new(), + &upstream, + "chat/completions", + Bytes::from("{}"), + ) + .unwrap(); + assert_eq!(url, expected); + } + } + + #[test] + fn thinking_migration_preserves_legacy_models() { + for model in ["claude-opus-4.8", "claude-opus-4-7", "claude-sonnet-5"] { + let mut body = + json!({"model":model,"thinking":{"type":"enabled","budget_tokens":12000}}); + rewrite_unsupported_thinking(&mut body); + assert_eq!(body["thinking"], json!({"type":"adaptive"})); + } + for model in ["claude-opus-4.6", "claude-sonnet-4.5", "gpt-5.4"] { + let mut body = + json!({"model":model,"thinking":{"type":"enabled","budget_tokens":12000}}); + let original = body.clone(); + rewrite_unsupported_thinking(&mut body); + assert_eq!(body, original); + } + } + } + #[test] fn post_embeddings_is_idempotent() { assert!(is_idempotent(&Method::POST, "/openai/v1/embeddings")); diff --git a/inference-router/src/proxy/authentication.rs b/inference-router/src/proxy/authentication.rs new file mode 100644 index 00000000..8611fb85 --- /dev/null +++ b/inference-router/src/proxy/authentication.rs @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! A named route cannot inherit credentials from the legacy default route. + +use anyhow::{Context, Result}; + +use super::{ + UpstreamConfig, UpstreamCredential, endpoint_host, is_azure_ai_host, is_copilot_endpoint, + is_local_inference_host, token_audience, +}; +use crate::{auth::WorkloadIdentityAuth, copilot_auth::CopilotTokenCache, provider::ProviderKind}; + +#[derive(Clone, Debug, Default, PartialEq, Eq, serde::Serialize)] +pub enum AuthenticationProvenance { + #[default] + LegacyDefault, + /// An immutable Config provider ID, never its credential. + Named { provider_id: String }, +} + +pub async fn token_for_endpoint( + auth: &WorkloadIdentityAuth, + copilot: Option<&CopilotTokenCache>, + upstream: &UpstreamConfig, +) -> Result { + if let AuthenticationProvenance::Named { provider_id } = &upstream.authentication { + if is_copilot_endpoint(&upstream.endpoint) { + let token = upstream + .provider_api_key + .as_deref() + .filter(|token| !token.trim().is_empty()) + .context("named Copilot provider requires its own GitHub seat token")?; + return copilot + .context("Copilot token cache is unavailable")? + .get_jwt_for_provider(provider_id, token) + .await; + } + return Ok(upstream.provider_api_key.clone().unwrap_or_default()); + } + let endpoint = upstream.endpoint.as_str(); + if is_copilot_endpoint(endpoint) { + return copilot + .context("default Copilot endpoint requires COPILOT_GITHUB_TOKEN")? + .get_jwt() + .await; + } + if let Some(key) = upstream.provider_api_key.as_deref() { + return Ok(key.to_string()); + } + if is_local_inference_host(&endpoint_host(endpoint).unwrap_or_default()) { + return Ok(String::new()); + } + if !auth.is_api_key_mode() && !auth.is_sidecar_mode() { + let host = endpoint_host(endpoint).unwrap_or_default(); + anyhow::ensure!( + is_azure_ai_host(&host), + "Refusing to send a Workload Identity / IMDS token to '{host}': configure the provider's own credential" + ); + } + auth.get_token(token_audience(endpoint)).await +} + +pub async fn credential_for_upstream( + auth: &WorkloadIdentityAuth, + copilot: Option<&CopilotTokenCache>, + upstream: &UpstreamConfig, +) -> Result { + match upstream.provider { + ProviderKind::AzureOpenAI => { + // Preserve the historical default local-service exclusion. + if upstream.authentication == AuthenticationProvenance::LegacyDefault + && is_local_inference_host(&endpoint_host(&upstream.endpoint).unwrap_or_default()) + { + return Ok(UpstreamCredential::None); + } + let token = token_for_endpoint(auth, copilot, upstream).await?; + Ok(if token.is_empty() { + UpstreamCredential::None + } else { + UpstreamCredential::Bearer(token) + }) + } + ProviderKind::Anthropic => upstream + .api_key + .clone() + .map(UpstreamCredential::AnthropicApiKey) + .context("Anthropic provider requires its own configured credential"), + ProviderKind::Ollama => Ok(UpstreamCredential::None), + } +} diff --git a/inference-router/src/proxy/authentication_tests.rs b/inference-router/src/proxy/authentication_tests.rs new file mode 100644 index 00000000..27e722fc --- /dev/null +++ b/inference-router/src/proxy/authentication_tests.rs @@ -0,0 +1,224 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::{ + config::{Config, ProviderEndpoint}, + failover::{Candidate, resolve_candidate}, +}; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{header, method, path}, +}; + +fn named_target(endpoint: String, credential: Option<&str>, id: &str) -> UpstreamConfig { + let mut config = Config::from_env().unwrap(); + config.providers.insert( + id.into(), + ProviderEndpoint { + tag: id.into(), + endpoint, + api_key: credential.map(str::to_string), + }, + ); + resolve_candidate( + &UpstreamConfig::azure( + "https://default.openai.azure.com".into(), + "model".into(), + "test".into(), + ), + &config, + &Candidate { + provider: Some(id.into()), + deployment: "model".into(), + routing_intent: crate::failover::RoutingIntent::Explicit, + }, + ) + .unwrap() +} + +#[tokio::test] +async fn keyless_named_local_and_remote_routes_never_consult_default_api_key_or_sidecar() { + let inference = MockServer::start().await; + let sidecar = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .mount(&sidecar) + .await; + Mock::given(|request: &wiremock::Request| { + !request.headers.contains_key("authorization") + && !request.headers.contains_key("api-key") + && !request.headers.contains_key("x-api-key") + }) + .respond_with(ResponseTemplate::new(200)) + .expect(4) + .mount(&inference) + .await; + + for host in ["model.models.svc.cluster.local", "model.example.test"] { + let client = Client::builder() + .no_proxy() + .resolve(host, *inference.address()) + .build() + .unwrap(); + for auth in [ + WorkloadIdentityAuth::for_test(Some("ambient-default-key"), None), + WorkloadIdentityAuth::for_test( + None, + Some(crate::sidecar_client::SidecarClient::for_test( + sidecar.uri(), + )), + ), + ] { + let upstream = named_target( + format!("http://{host}:{}", inference.address().port()), + None, + "keyless", + ); + let headers = HeaderMap::from_iter([ + ( + "authorization".parse().unwrap(), + HeaderValue::from_static("Bearer inbound-secret"), + ), + ( + "api-key".parse().unwrap(), + HeaderValue::from_static("inbound-key"), + ), + ]); + let (status, _, _) = forward( + &auth, + None, + &client, + &upstream, + Method::POST, + "chat/completions", + &headers, + Bytes::from("{}"), + ) + .await + .unwrap(); + assert_eq!(status, StatusCode::OK); + } + } +} + +#[tokio::test] +async fn named_route_uses_only_its_own_key_and_legacy_default_still_uses_ambient_key() { + let server = MockServer::start().await; + for key in ["named-key", "ambient-default-key"] { + Mock::given(header("authorization", format!("Bearer {key}"))) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + } + let auth = WorkloadIdentityAuth::for_test(Some("ambient-default-key"), None); + let named = named_target(server.uri(), Some("named-key"), "named"); + let default = UpstreamConfig::azure(server.uri(), "model".into(), "test".into()); + for upstream in [named, default] { + assert_eq!( + forward( + &auth, + None, + &Client::new(), + &upstream, + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await + .unwrap() + .0, + StatusCode::OK + ); + } +} + +#[tokio::test] +async fn real_copilot_host_selects_the_named_exchange_cache_not_the_default_account() { + let exchange = MockServer::start().await; + let inference = MockServer::start().await; + let expires_at = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_secs() + + 3600; + for (seat, jwt) in [ + ("default-seat", "default-jwt"), + ("seat-a", "jwt-a"), + ("seat-b", "jwt-b"), + ] { + Mock::given(header("authorization", format!("token {seat}"))) + .and(path("/exchange")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "token": jwt, "expires_at": expires_at, "refresh_in": 1500, + }))) + .expect(1) + .mount(&exchange) + .await; + Mock::given(header("authorization", format!("Bearer {jwt}"))) + .and(header("copilot-integration-id", COPILOT_INTEGRATION_ID)) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&inference) + .await; + } + let url = format!("{}/exchange", exchange.uri()); + let cache = CopilotTokenCache::with_test_exchange("default-seat", url); + let endpoint = format!( + "http://api.githubcopilot.com:{}", + inference.address().port() + ); + assert!(is_copilot_endpoint(&endpoint)); + let client = Client::builder() + .no_proxy() + .resolve("api.githubcopilot.com", *inference.address()) + .build() + .unwrap(); + let auth = WorkloadIdentityAuth::for_test(Some("unrelated-azure-key"), None); + for upstream in [ + named_target(endpoint.clone(), Some("seat-a"), "account-a"), + named_target(endpoint.clone(), Some("seat-b"), "account-b"), + UpstreamConfig::azure(endpoint, "model".into(), "test".into()), + ] { + assert_eq!( + forward( + &auth, + Some(&cache), + &client, + &upstream, + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await + .unwrap() + .0, + StatusCode::OK + ); + } +} + +#[tokio::test] +async fn named_copilot_without_a_seat_token_cannot_borrow_the_default_account() { + let exchange = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(500)) + .expect(0) + .mount(&exchange) + .await; + let upstream = named_target("https://api.githubcopilot.com".into(), None, "missing-seat"); + let result = credential_for_upstream( + &WorkloadIdentityAuth::for_test(Some("azure-key"), None), + Some(&CopilotTokenCache::with_test_exchange( + "other-account", + exchange.uri(), + )), + &upstream, + ) + .await; + assert!(result.is_err()); +} diff --git a/inference-router/src/proxy/failure.rs b/inference-router/src/proxy/failure.rs new file mode 100644 index 00000000..7ed700e8 --- /dev/null +++ b/inference-router/src/proxy/failure.rs @@ -0,0 +1,164 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Failover needs an explicit cause and acceptance boundary, not just `Err`. + +use axum::http::StatusCode; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FailureCategory { + Configuration, + Authentication, + Transport, + ResponseBody, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Acceptance { + NotAccepted, + Accepted(StatusCode), + Rejected(StatusCode), + Unknown, +} + +#[derive(Debug, thiserror::Error)] +#[error("{category:?} upstream failure ({acceptance:?}): {source}")] +pub struct ForwardFailure { + pub category: FailureCategory, + pub acceptance: Acceptance, + #[source] + source: anyhow::Error, +} + +impl ForwardFailure { + pub fn configuration(error: impl Into) -> anyhow::Error { + Self { + category: FailureCategory::Configuration, + acceptance: Acceptance::NotAccepted, + source: error.into(), + } + .into() + } + + pub fn authentication(error: impl Into) -> anyhow::Error { + Self { + category: FailureCategory::Authentication, + acceptance: Acceptance::NotAccepted, + source: error.into(), + } + .into() + } + + pub fn transport(error: reqwest::Error) -> anyhow::Error { + if error.is_builder() { + return Self::configuration(error); + } + // Only a failed connection proves the request was not accepted. + // A timeout/reset after sending could already have started generation. + let acceptance = if error.is_connect() { + Acceptance::NotAccepted + } else { + Acceptance::Unknown + }; + Self { + category: FailureCategory::Transport, + acceptance, + source: error.into(), + } + .into() + } + + pub fn response_body(status: StatusCode, error: reqwest::Error) -> anyhow::Error { + Self { + category: FailureCategory::ResponseBody, + acceptance: if status.is_success() { + Acceptance::Accepted(status) + } else { + Acceptance::Rejected(status) + }, + source: error.into(), + } + .into() + } +} + +pub fn retryable_rejection(status: StatusCode) -> bool { + status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error() +} + +pub fn retryable_failure(error: &anyhow::Error) -> bool { + error + .downcast_ref::() + .is_some_and(|failure| match (failure.category, failure.acceptance) { + (FailureCategory::Transport, Acceptance::NotAccepted) => true, + (FailureCategory::ResponseBody, Acceptance::Rejected(status)) => { + retryable_rejection(status) + } + _ => false, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn body_failures_retry_only_known_retryable_rejections() { + for code in [ + 200, 201, 204, 400, 401, 403, 404, 422, 429, 500, 502, 503, 504, 599, + ] { + let status = StatusCode::from_u16(code).unwrap(); + let acceptance = if status.is_success() { + Acceptance::Accepted(status) + } else { + Acceptance::Rejected(status) + }; + let error: anyhow::Error = ForwardFailure { + category: FailureCategory::ResponseBody, + acceptance, + source: anyhow::anyhow!("truncated body after headers"), + } + .into(); + assert_eq!( + retryable_failure(&error), + code == 429 || code >= 500, + "{code}" + ); + } + for category in [ + FailureCategory::Authentication, + FailureCategory::Configuration, + ] { + let error: anyhow::Error = ForwardFailure { + category, + acceptance: Acceptance::Rejected(StatusCode::SERVICE_UNAVAILABLE), + source: anyhow::anyhow!("not an inference rejection"), + } + .into(); + assert!(!retryable_failure(&error)); + } + } + + #[test] + fn unknown_auth_config_and_accepted_body_errors_cannot_trigger_failover() { + for error in [ + anyhow::anyhow!("untyped failure"), + ForwardFailure::authentication(anyhow::anyhow!("missing credential")), + ForwardFailure::configuration(anyhow::anyhow!("invalid endpoint")), + ForwardFailure { + category: FailureCategory::ResponseBody, + acceptance: Acceptance::Accepted(StatusCode::OK), + source: anyhow::anyhow!("truncated response"), + } + .into(), + ForwardFailure { + category: FailureCategory::Transport, + acceptance: Acceptance::Unknown, + source: anyhow::anyhow!("timeout after sending"), + } + .into(), + ] { + assert!(!retryable_failure(&error)); + } + } +} diff --git a/inference-router/src/routes/anthropic_messages.rs b/inference-router/src/routes/anthropic_messages.rs index 0e939f24..de479901 100644 --- a/inference-router/src/routes/anthropic_messages.rs +++ b/inference-router/src/routes/anthropic_messages.rs @@ -342,10 +342,12 @@ pub(super) async fn anthropic_messages( let mut upstream = state.upstream_config(sandbox_name); // Slice 2d.1: honour `InferencePolicy.modelPreference.primary.deployment`. - crate::routes::apply_model_preference_override(&mut upstream, &policy); + let provider_resolution = + crate::routes::apply_model_preference_override(&mut upstream, &policy, &state.config) + .and_then(|_| crate::routes::apply_provider_resolution(&state, &mut upstream, &policy)); // Retarget at the policy-selected provider (fails closed). - if let Err(e) = crate::routes::apply_provider_resolution(&state, &mut upstream, &policy) { + if let Err(e) = provider_resolution { tracing::warn!( target: "inference.audit", sandbox = %sandbox_name, diff --git a/inference-router/src/routes/chat_completions.rs b/inference-router/src/routes/chat_completions.rs index 62d539a9..89ac9a0f 100644 --- a/inference-router/src/routes/chat_completions.rs +++ b/inference-router/src/routes/chat_completions.rs @@ -19,10 +19,10 @@ use futures::stream::StreamExt; use super::AppState; use super::inference_translate::{chat_to_responses_body, responses_to_chat_body}; +use super::model_routing::{is_responses_only_error, model_capability_key, override_model_in_body}; use crate::errors; use crate::guardrails::{self, Direction, GuardrailError, GuardrailPipeline, GuardrailViolation}; use crate::provider::{ProviderError, ProviderKind}; -use crate::proxy; use crate::safety; use std::sync::Arc; @@ -424,11 +424,14 @@ pub(super) async fn chat_completions( // Forward to Foundry let mut upstream = state.upstream_config(sandbox_name); + let true_default_upstream = upstream.clone(); // Slice 2d.1: honour `InferencePolicy.modelPreference.primary.deployment`. - crate::routes::apply_model_preference_override(&mut upstream, &policy); + let provider_resolution = + crate::routes::apply_model_preference_override(&mut upstream, &policy, &state.config) + .and_then(|_| crate::routes::apply_provider_resolution(&state, &mut upstream, &policy)); // Retarget at the policy-selected provider; fails closed. - if let Err(e) = crate::routes::apply_provider_resolution(&state, &mut upstream, &policy) { + if let Err(e) = provider_resolution { tracing::warn!( target: "inference.audit", sandbox = %sandbox_name, @@ -538,11 +541,28 @@ pub(super) async fn chat_completions( .ok() .and_then(|v| v.get("model")?.as_str().map(String::from)) .unwrap_or_else(|| upstream.deployment.clone()); + upstream = match super::model_routing::effective_primary( + &state, + &true_default_upstream, + &policy, + &body, + ) { + Ok(upstream) => upstream, + Err(error) => { + tracing::warn!(%error, "Inference provider configuration is unavailable"); + return errors::openai( + StatusCode::BAD_GATEWAY, + "Inference provider configuration is unavailable", + errors::PROXY_ERROR, + ) + .into_response(); + } + }; let is_responses_only = state .responses_only_models .read() .ok() - .map(|set| set.contains(&model_name)) + .map(|set| set.contains(&model_capability_key(&upstream))) .unwrap_or(false); if is_responses_only { @@ -550,7 +570,8 @@ pub(super) async fn chat_completions( // Use a streaming response body with SSE keepalive comments to prevent // client timeouts while the Responses API processes (30-50s for reasoning models). tracing::info!(sandbox = %sandbox_name, model = %model_name, "Using cached Responses API path"); - let responses_body = chat_to_responses_body(&body); + let responses_body = + chat_to_responses_body(&override_model_in_body(&body, &upstream.deployment)); let is_stream = serde_json::from_slice::(&body) .ok() @@ -561,9 +582,9 @@ pub(super) async fn chat_completions( // Stream keepalive comments while the Responses API call is in progress, // then send the converted result as a single SSE data frame. let (tx, rx) = tokio::sync::mpsc::channel::>(16); - let auth = state.auth.clone(); - let copilot = state.copilot.clone(); - let client = state.client.clone(); + let response_state = state.clone(); + let response_base = true_default_upstream.clone(); + let response_policy = policy.clone(); let upstream = upstream.clone(); let headers = headers.clone(); let budget = state.budget.clone(); @@ -573,13 +594,11 @@ pub(super) async fn chat_completions( tokio::spawn(async move { // Send keepalive comments every 5 seconds while waiting - let forward_fut = proxy::forward( - &auth, - Some(&copilot), - &client, + let forward_fut = super::model_routing::forward_responses( + &response_state, + &response_base, + &response_policy, &upstream, - axum::http::Method::POST, - "responses", &headers, responses_body, ); @@ -599,7 +618,7 @@ pub(super) async fn chat_completions( }; match result { - Ok((_resp_status, _, resp_body)) => { + Ok((_resp_status, _, resp_body, _selected)) => { let chat_body = responses_to_chat_body(&resp_body); if let Ok(bj) = serde_json::from_slice::(&chat_body) && let Some(total) = bj @@ -655,19 +674,17 @@ pub(super) async fn chat_completions( } // Non-streaming: buffered request/response (no timeout concern) - match proxy::forward( - &state.auth, - Some(&state.copilot), - &state.client, + match super::model_routing::forward_responses( + &state, + &true_default_upstream, + &policy, &upstream, - axum::http::Method::POST, - "responses", &headers, responses_body, ) .await { - Ok((resp_status, resp_hdrs, resp_body)) => { + Ok((resp_status, resp_hdrs, resp_body, _selected)) => { let chat_body = responses_to_chat_body(&resp_body); if let Ok(bj) = serde_json::from_slice::(&chat_body) && let Some(total) = bj @@ -726,18 +743,18 @@ pub(super) async fn chat_completions( // rarely vs. single-request lifetime. let stream_floor = policy.content_safety.clone(); let stream_policy_digest = policy.digest.clone(); - match proxy::forward_stream( - state.auth.clone(), - Some(state.copilot.clone()), - state.client.clone(), - upstream.clone(), - "chat/completions", + match super::model_routing::forward_stream_chat( + &state, + &true_default_upstream, + &policy, headers.clone(), body.clone(), ) .await { - Ok((status, _resp_headers, stream)) if status == StatusCode::BAD_REQUEST => { + Ok((status, _resp_headers, stream, selected_upstream)) + if status == StatusCode::BAD_REQUEST => + { // Might be a Responses-only model — buffer the error and check use futures::TryStreamExt; let err_bytes: Vec = stream @@ -747,38 +764,31 @@ pub(super) async fn chat_completions( }) .await .unwrap_or_default(); - let is_unsupported = serde_json::from_slice::(&err_bytes) - .ok() - .and_then(|v| { - v.get("error")? - .get("message")? - .as_str() - .map(|s| s.contains("unsupported")) - }) - .unwrap_or(false); + let is_unsupported = is_responses_only_error(&err_bytes); if is_unsupported { // Cache this model as Responses-only to skip future chat/completions attempts if let Ok(mut set) = state.responses_only_models.write() { - set.insert(model_name.clone()); + set.insert(model_capability_key(&selected_upstream)); tracing::info!(model = %model_name, "Cached as Responses-only model"); } // Fallback: convert to Responses API, return result as single SSE frame tracing::info!(sandbox = %sandbox_name, "Streaming chat/completions unsupported, falling back to Responses API"); - let responses_body = chat_to_responses_body(&body); - match proxy::forward( - &state.auth, - Some(&state.copilot), - &state.client, - &upstream, - axum::http::Method::POST, - "responses", + let responses_body = chat_to_responses_body(&override_model_in_body( + &body, + &selected_upstream.deployment, + )); + match super::model_routing::forward_responses( + &state, + &true_default_upstream, + &policy, + &selected_upstream, &headers, responses_body, ) .await { - Ok((resp_status, _, resp_body)) => { + Ok((resp_status, _, resp_body, _selected)) => { let chat_body = responses_to_chat_body(&resp_body); if let Ok(bj) = serde_json::from_slice::(&chat_body) && let Some(total) = bj @@ -825,7 +835,7 @@ pub(super) async fn chat_completions( (StatusCode::BAD_REQUEST, Body::from(err_bytes)).into_response() } } - Ok((status, resp_headers, stream)) => { + Ok((status, resp_headers, stream, _selected_upstream)) => { tracing::info!(sandbox = %sandbox_owned, status = %status.as_u16(), "Stream response status"); // Wrap stream to intercept the first SSE chunk for guardrail // annotations and the final chunk for token usage. @@ -971,55 +981,42 @@ pub(super) async fn chat_completions( // retries against `fallback[N].deployment`. The 400-→- // Responses-API recovery further down still runs against the // *successful* upstream's deployment. - let result = crate::failover::forward_with_failover( - &state.auth, - Some(&state.copilot), - &state.client, - &state.deployment_health, - &upstream, + let result = super::model_routing::forward_chat( + &state, + &true_default_upstream, &policy, - axum::http::Method::POST, - "chat/completions", &headers, body.clone(), ) .await; match result { - Ok((status, _resp_headers, resp_body)) - if status == StatusCode::BAD_REQUEST - && serde_json::from_slice::(&resp_body) - .ok() - .and_then(|v| { - v.get("error")? - .get("message")? - .as_str() - .map(|s| s.contains("unsupported")) - }) - .unwrap_or(false) => + Ok((status, _resp_headers, resp_body, selected_upstream)) + if status == StatusCode::BAD_REQUEST && is_responses_only_error(&resp_body) => { // Model doesn't support chat/completions — auto-fallback to Responses API. // Cache this model to skip future chat/completions attempts. if let Ok(mut set) = state.responses_only_models.write() { - set.insert(model_name.clone()); + set.insert(model_capability_key(&selected_upstream)); tracing::info!(model = %model_name, "Cached as Responses-only model"); } // Convert messages → input and proxy to /openai/v1/responses. tracing::info!(sandbox = %sandbox_name, "chat/completions unsupported, falling back to Responses API"); - let responses_body = chat_to_responses_body(&body); - match proxy::forward( - &state.auth, - Some(&state.copilot), - &state.client, - &upstream, - axum::http::Method::POST, - "responses", + let responses_body = chat_to_responses_body(&override_model_in_body( + &body, + &selected_upstream.deployment, + )); + match super::model_routing::forward_responses( + &state, + &true_default_upstream, + &policy, + &selected_upstream, &headers, responses_body, ) .await { - Ok((resp_status, resp_hdrs, resp_body)) => { + Ok((resp_status, resp_hdrs, resp_body, _selected)) => { // Convert Responses API output back to chat/completions format let chat_body = responses_to_chat_body(&resp_body); if let Ok(body_json) = @@ -1058,7 +1055,7 @@ pub(super) async fn chat_completions( } } } - Ok((status, resp_headers, mut resp_body)) => { + Ok((status, resp_headers, mut resp_body, _selected_upstream)) => { // Record token usage from response for budget tracking if let Ok(body_json) = serde_json::from_slice::(&resp_body) && let Some(total) = body_json diff --git a/inference-router/src/routes/inference.rs b/inference-router/src/routes/inference.rs index d0979b72..d7dce2e2 100644 --- a/inference-router/src/routes/inference.rs +++ b/inference-router/src/routes/inference.rs @@ -382,18 +382,21 @@ async fn responses( // the byte stream to flow through unchanged. use axum::body::Body; use futures::TryStreamExt; - match proxy::forward_stream( + match crate::failover::forward_stream_with_failover( state.auth.clone(), Some(state.copilot.clone()), state.client.clone(), - upstream, + &state.deployment_health, + &upstream, + &state.config, + &policy, "responses", headers.clone(), body, ) .await { - Ok((status, resp_headers, stream)) => { + Ok((status, resp_headers, stream, _selected_upstream)) => { // Surface usage tokens by buffering only the very last chunk // is impossible without breaking streaming. We accept that // the budget tracker won't see /v1/responses usage in diff --git a/inference-router/src/routes/inference_translate.rs b/inference-router/src/routes/inference_translate.rs index 8582535b..f61915b8 100644 --- a/inference-router/src/routes/inference_translate.rs +++ b/inference-router/src/routes/inference_translate.rs @@ -243,6 +243,7 @@ pub(super) fn chat_to_responses_body(chat_body: &Bytes) -> Bytes { // Remove chat-specific fields that Responses API doesn't accept obj.remove("stream"); + obj.remove("stream_options"); obj.remove("stop"); obj.remove("frequency_penalty"); obj.remove("presence_penalty"); diff --git a/inference-router/src/routes/mod.rs b/inference-router/src/routes/mod.rs index bcc92eca..ff8d134a 100644 --- a/inference-router/src/routes/mod.rs +++ b/inference-router/src/routes/mod.rs @@ -45,6 +45,7 @@ mod mesh; pub use mesh::mesh_routes; mod mesh_token; +mod model_routing; pub use mesh_token::mesh_token_routes; mod egress; @@ -110,6 +111,7 @@ pub struct AppState { /// Models that don't support chat/completions (need Responses API). /// Populated on first 400 "unsupported" — avoids redundant round-trips. pub responses_only_models: Arc>>, + pub unavailable_models: Arc>>, /// Handoff token store (in-memory, TTL-based, one-at-a-time). pub handoff_tokens: HandoffTokenStore, /// Handoff session tracker (phase, direction, progress). @@ -327,6 +329,7 @@ impl AppState { responses_only_models: Arc::new(std::sync::RwLock::new( std::collections::HashSet::new(), )), + unavailable_models: Arc::new(std::sync::RwLock::new(std::collections::HashSet::new())), admin_token: std::fs::read_to_string("/etc/kars/secrets/admin-token") .or_else(|_| std::fs::read_to_string("/run/secrets/admin-token")) .or_else(|_| std::env::var("ADMIN_TOKEN")) @@ -389,6 +392,10 @@ pub(crate) fn apply_provider_resolution( upstream.endpoint = endpoint; upstream.provider = crate::provider::ProviderKind::Anthropic; upstream.api_key = Some(api_key); + upstream.provider_api_key = None; + upstream.authentication = crate::proxy::AuthenticationProvenance::Named { + provider_id: "anthropic".into(), + }; } crate::provider::ProviderTarget::Ollama { endpoint } => { tracing::info!( @@ -400,41 +407,29 @@ pub(crate) fn apply_provider_resolution( upstream.endpoint = endpoint; upstream.provider = crate::provider::ProviderKind::Ollama; upstream.api_key = None; + upstream.provider_api_key = None; + upstream.authentication = crate::proxy::AuthenticationProvenance::Named { + provider_id: "ollama".into(), + }; } } Ok(()) } -/// Slice 2d.1 — apply `modelPreference.primary.deployment` from a -/// loaded `InferencePolicy` snapshot as a deployment override. -/// -/// Mutates `upstream.deployment` in place when the policy carries a -/// non-empty `primary.deployment` that differs from the current -/// deployment. Logs an `info!` event on every effective override so -/// operators can correlate router-level traffic shaping against the -/// policy bytes (digest is included). -/// -/// Fail-open by design: -/// * `None` snapshot ⇒ no-op (back-compat for sandboxes without an -/// `InferencePolicy`). -/// * Empty-string `primary.deployment` ⇒ no-op (defence-in-depth even -/// though the controller schema rejects empty strings). -/// * Same-deployment override ⇒ no-op + no log spam. -/// -/// **Slice 2d.1 deliberately ignores `primary.provider`** — provider- -/// tagged routing requires a per-provider client registry the router -/// doesn't carry today; Slice 2d.2 will pick that up. Until then the -/// provider tag is informational-only. +/// Apply the policy's primary model and configured provider together. +/// An absent preference or empty deployment leaves the default unchanged; +/// an unconfigured informational provider tag retains the legacy endpoint. pub(crate) fn apply_model_preference_override( upstream: &mut UpstreamConfig, policy: &crate::inference_policy_loader::InferencePolicySnapshot, -) { + config: &crate::config::Config, +) -> Result<(), crate::provider::ProviderError> { let Some(ref pref) = policy.model_preference else { - return; + return Ok(()); }; let target = pref.primary.deployment.as_str(); - if target.is_empty() || target == upstream.deployment { - return; + if target.is_empty() { + return Ok(()); } tracing::info!( sandbox = %upstream.sandbox_name, @@ -444,7 +439,20 @@ pub(crate) fn apply_model_preference_override( digest = %policy.digest, "InferencePolicy modelPreference: overriding deployment" ); - upstream.deployment = target.to_string(); + *upstream = crate::failover::resolve_candidate( + upstream, + config, + &crate::failover::Candidate { + provider: policy + .provider + .clone() + .filter(|p| !p.trim().is_empty()) + .or_else(|| Some(pref.primary.provider.clone()).filter(|p| !p.trim().is_empty())), + deployment: target.to_string(), + routing_intent: crate::failover::RoutingIntent::primary(policy), + }, + )?; + Ok(()) } /// Extract the admin bearer token from either `Authorization: Bearer ` diff --git a/inference-router/src/routes/model_routing.rs b/inference-router/src/routes/model_routing.rs new file mode 100644 index 00000000..61aa5d1c --- /dev/null +++ b/inference-router/src/routes/model_routing.rs @@ -0,0 +1,574 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::AppState; +use crate::{ + failover, + inference_policy_loader::InferencePolicySnapshot, + proxy::{self, UpstreamConfig}, +}; +use anyhow::Result; +use axum::http::{HeaderMap, Method, StatusCode}; +use bytes::Bytes; +use futures::{StreamExt, TryStreamExt, stream::BoxStream}; + +type BufferedResult = (StatusCode, HeaderMap, Bytes, UpstreamConfig); +type StreamingResult = ( + StatusCode, + HeaderMap, + BoxStream<'static, Result>, + UpstreamConfig, +); + +pub(super) fn model_capability_key(upstream: &UpstreamConfig) -> String { + // Config is immutable for this AppState. Its explicit provider ID separates + // accounts even when endpoint/model match; credentials never enter keys. + serde_json::to_string(&( + &upstream.authentication, + upstream.provider.as_tag(), + upstream.endpoint.trim_end_matches('/'), + &upstream.deployment, + )) + .expect("provider capability key always serializes") +} + +pub(super) fn is_responses_only_error(body: &[u8]) -> bool { + let Ok(value) = serde_json::from_slice::(body) else { + return false; + }; + let error = value.get("error").unwrap_or(&value); + let code = error + .get("code") + .and_then(|v| v.as_str()) + .unwrap_or_default(); + let message = error + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_ascii_lowercase(); + code.eq_ignore_ascii_case("unsupported_api_for_model") + || message.contains("not accessible via the /chat/completions endpoint") + || (message.contains("unsupported") + && (message.contains("chat") || message.contains("api"))) +} + +fn is_model_unavailable_error(status: StatusCode, body: &[u8]) -> bool { + if !matches!(status, StatusCode::BAD_REQUEST | StatusCode::NOT_FOUND) + || is_responses_only_error(body) + { + return false; + } + let Ok(value) = serde_json::from_slice::(body) else { + return false; + }; + let error = value.get("error").unwrap_or(&value); + let code = error + .get("code") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_ascii_lowercase(); + let message = error + .get("message") + .and_then(|v| v.as_str()) + .unwrap_or_default() + .to_ascii_lowercase(); + matches!( + code.as_str(), + "deploymentnotfound" | "model_not_found" | "model_not_supported" | "invalid_model" + ) || ((message.contains("model") || message.contains("deployment")) + && [ + "does not exist", + "not found", + "unknown model", + "invalid model", + "not supported", + "is not a valid model", + "no deployment", + ] + .iter() + .any(|phrase| message.contains(phrase))) +} + +pub(super) fn override_model_in_body(body: &[u8], model: &str) -> Bytes { + match serde_json::from_slice::(body) { + Ok(mut value) if value.is_object() => { + value["model"] = model.into(); + serde_json::to_vec(&value) + .map(Bytes::from) + .unwrap_or_else(|_| Bytes::copy_from_slice(body)) + } + _ => Bytes::copy_from_slice(body), + } +} + +fn default_target(state: &AppState, base: &UpstreamConfig) -> UpstreamConfig { + let mut target = base.clone(); + target.deployment = state.config.default_model.clone(); + target +} + +fn primary_target( + state: &AppState, + base: &UpstreamConfig, + policy: &InferencePolicySnapshot, + body: &[u8], +) -> Result { + let candidates = failover::candidates_for_request(base, policy, body); + failover::resolve_candidate(base, &state.config, &candidates[0]) + .map_err(proxy::failure::ForwardFailure::configuration) +} + +fn cached_unavailable( + state: &AppState, + selected: &UpstreamConfig, + fallback: &UpstreamConfig, +) -> bool { + model_capability_key(selected) != model_capability_key(fallback) + && state + .unavailable_models + .read() + .is_ok_and(|cache| cache.contains(&model_capability_key(selected))) +} + +fn remember_unavailable(state: &AppState, selected: &UpstreamConfig) { + if let Ok(mut cache) = state.unavailable_models.write() { + cache.insert(model_capability_key(selected)); + } +} + +pub(super) fn effective_primary( + state: &AppState, + base: &UpstreamConfig, + policy: &InferencePolicySnapshot, + body: &[u8], +) -> Result { + let primary = primary_target(state, base, policy, body)?; + let fallback = default_target(state, base); + if cached_unavailable(state, &primary, &fallback) { + Ok(fallback) + } else { + Ok(primary) + } +} + +/// Recovery starts with the provider that actually answered, not the original +/// primary that may already have failed. Remaining policy fallbacks still apply. +pub(super) async fn forward_responses( + state: &AppState, + base: &UpstreamConfig, + policy: &InferencePolicySnapshot, + selected: &UpstreamConfig, + headers: &HeaderMap, + body: Bytes, +) -> Result { + use crate::inference_policy_loader::{ModelPreference, ModelRef}; + let candidates = failover::build_candidates(base, policy); + let mut start = None; + for (index, candidate) in candidates.iter().enumerate() { + let target = failover::resolve_candidate(base, &state.config, candidate) + .map_err(proxy::failure::ForwardFailure::configuration)?; + if model_capability_key(&target) == model_capability_key(selected) { + start = Some(index); + break; + } + } + if let Some(start) = start { + let models: Vec<_> = candidates[start..] + .iter() + .map(|candidate| ModelRef { + provider: candidate.provider.clone().unwrap_or_default(), + deployment: candidate.deployment.clone(), + }) + .collect(); + let mut response_policy = policy.clone(); + response_policy.provider = + if candidates[start].routing_intent == failover::RoutingIntent::Explicit { + candidates[start].provider.clone() + } else { + None + }; + response_policy.model_preference = Some(ModelPreference { + primary: models[0].clone(), + fallback: models[1..].to_vec(), + }); + return failover::forward_with_failover( + &state.auth, + Some(&state.copilot), + &state.client, + &state.deployment_health, + base, + &state.config, + &response_policy, + Method::POST, + "responses", + headers, + body, + ) + .await; + } + proxy::forward( + &state.auth, + Some(&state.copilot), + &state.client, + selected, + Method::POST, + "responses", + headers, + override_model_in_body(&body, &selected.deployment), + ) + .await + .map(|(status, headers, bytes)| (status, headers, bytes, selected.clone())) +} + +pub(super) async fn forward_chat( + state: &AppState, + base: &UpstreamConfig, + policy: &InferencePolicySnapshot, + headers: &HeaderMap, + body: Bytes, +) -> Result { + let fallback = default_target(state, base); + let primary = primary_target(state, base, policy, &body)?; + let result = if cached_unavailable(state, &primary, &fallback) { + proxy::forward( + &state.auth, + Some(&state.copilot), + &state.client, + &fallback, + Method::POST, + "chat/completions", + headers, + override_model_in_body(&body, &fallback.deployment), + ) + .await + .map(|(status, headers, bytes)| (status, headers, bytes, fallback.clone())) + } else { + failover::forward_with_failover( + &state.auth, + Some(&state.copilot), + &state.client, + &state.deployment_health, + base, + &state.config, + policy, + Method::POST, + "chat/completions", + headers, + body.clone(), + ) + .await + }; + if let Ok((status, _, response, selected)) = &result + && is_model_unavailable_error(*status, response) + && !fallback.deployment.is_empty() + && model_capability_key(selected) != model_capability_key(&fallback) + { + remember_unavailable(state, selected); + return proxy::forward( + &state.auth, + Some(&state.copilot), + &state.client, + &fallback, + Method::POST, + "chat/completions", + headers, + override_model_in_body(&body, &fallback.deployment), + ) + .await + .map(|(status, headers, bytes)| (status, headers, bytes, fallback)); + } + result +} + +pub(super) async fn forward_stream_chat( + state: &AppState, + base: &UpstreamConfig, + policy: &InferencePolicySnapshot, + headers: HeaderMap, + body: Bytes, +) -> Result { + let fallback = default_target(state, base); + let primary = primary_target(state, base, policy, &body)?; + let (status, response_headers, stream, selected) = + if cached_unavailable(state, &primary, &fallback) { + proxy::forward_stream( + state.auth.clone(), + Some(state.copilot.clone()), + state.client.clone(), + fallback.clone(), + "chat/completions", + headers.clone(), + override_model_in_body(&body, &fallback.deployment), + ) + .await + .map(|(status, headers, stream)| (status, headers, stream, fallback.clone()))? + } else { + failover::forward_stream_with_failover( + state.auth.clone(), + Some(state.copilot.clone()), + state.client.clone(), + &state.deployment_health, + base, + &state.config, + policy, + "chat/completions", + headers.clone(), + body.clone(), + ) + .await? + }; + if !matches!(status, StatusCode::BAD_REQUEST | StatusCode::NOT_FOUND) { + return Ok((status, response_headers, stream, selected)); + } + let bytes = stream + .try_fold(Vec::new(), |mut bytes, chunk| async move { + bytes.extend_from_slice(&chunk); + Ok(bytes) + }) + .await?; + if is_model_unavailable_error(status, &bytes) + && !fallback.deployment.is_empty() + && model_capability_key(&selected) != model_capability_key(&fallback) + { + remember_unavailable(state, &selected); + return proxy::forward_stream( + state.auth.clone(), + Some(state.copilot.clone()), + state.client.clone(), + fallback.clone(), + "chat/completions", + headers, + override_model_in_body(&body, &fallback.deployment), + ) + .await + .map(|(status, headers, stream)| (status, headers, stream, fallback)); + } + Ok(( + status, + response_headers, + futures::stream::once(async move { Ok(Bytes::from(bytes)) }).boxed(), + selected, + )) +} + +#[cfg(test)] +#[path = "model_routing_regressions.rs"] +mod regressions; + +#[cfg(test)] +#[path = "model_routing_closure_tests.rs"] +mod closure_tests; + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + use std::sync::Arc; + use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_partial_json, header, path}, + }; + + pub(super) fn test_state(config: crate::config::Config) -> AppState { + let policy_status = Arc::new(crate::policy_status::PolicyStatusRegistry::new()); + let governance = Arc::new(crate::governance::Governance::new_with_status( + "test", + policy_status.clone(), + )); + AppState { + auth: Arc::new(crate::auth::WorkloadIdentityAuth::new()), + copilot: Arc::new(crate::copilot_auth::CopilotTokenCache::from_env()), + client: reqwest::Client::new(), + config: Arc::new(config), + budget: crate::budget::TokenBudgetTracker::new(0, 0), + policy_provider: governance.clone(), + audit_sink: governance.clone(), + signing_provider: governance.clone(), + governance, + blocklist: crate::blocklist::Blocklist::disabled(), + blocked_egress: Arc::new(crate::egress_blocked::BlockedBuffer::with_defaults()), + sandbox_name: Arc::new("test".into()), + inbox: Arc::new(crate::mesh::MeshInbox::new()), + mesh_metrics: Arc::new(crate::mesh::MeshMetrics::new()), + model_override: Default::default(), + responses_only_models: Default::default(), + unavailable_models: Default::default(), + admin_token: None, + handoff_tokens: crate::handoff::HandoffTokenStore::new(), + handoff_session: crate::handoff::HandoffSession::new(), + drain_state: crate::handoff::DrainState::new(), + pending_handoff: crate::handoff::PendingHandoffStore::new(), + policy_status, + inference_policy: crate::inference_policy_loader::empty_handle(), + memory_binding: crate::memory_binding_loader::empty_handle(), + egress_allowlist: crate::egress_allowlist_loader::empty_handle(), + deployment_health: Arc::new(crate::deployment_health::DeploymentHealthRegistry::new()), + } + } + + #[tokio::test] + async fn unavailable_provider_model_recovers_once_then_uses_scoped_cache() { + use crate::{ + config::{Config, ProviderEndpoint}, + inference_policy_loader::{ModelPreference, ModelRef}, + }; + let primary = MockServer::start().await; + let default = MockServer::start().await; + Mock::given(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(404) + .set_body_json(json!({"error":{"code":"model_not_found"}})), + ) + .expect(1) + .mount(&primary) + .await; + Mock::given(header("authorization", "Bearer default-key")) + .and(body_partial_json(json!({"model":"default-model"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"choices":[]}))) + .expect(2) + .mount(&default) + .await; + let mut config = Config::from_env().unwrap(); + config.default_model = "default-model".into(); + config.providers.insert( + "primary".into(), + ProviderEndpoint { + tag: "primary".into(), + endpoint: primary.uri(), + api_key: Some("primary-key".into()), + }, + ); + let state = test_state(config); + let mut base = UpstreamConfig::azure(default.uri(), "default-model".into(), "test".into()); + base.provider_api_key = Some("default-key".into()); + let policy = InferencePolicySnapshot { + model_preference: Some(ModelPreference { + primary: ModelRef { + provider: "primary".into(), + deployment: "missing-model".into(), + }, + fallback: vec![], + }), + ..Default::default() + }; + let (_, _, _, selected) = + forward_chat(&state, &base, &policy, &HeaderMap::new(), Bytes::from("{}")) + .await + .unwrap(); + assert_eq!(selected.endpoint, default.uri()); + let (status, _, stream, selected) = + forward_stream_chat(&state, &base, &policy, HeaderMap::new(), Bytes::from("{}")) + .await + .unwrap(); + assert_eq!(status, StatusCode::OK); + assert_eq!(selected.deployment, "default-model"); + let _: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!(state.unavailable_models.read().unwrap().len(), 1); + } + + #[tokio::test] + async fn responses_recovery_starts_at_the_provider_that_answered() { + use crate::{ + config::{Config, ProviderEndpoint}, + inference_policy_loader::{ModelPreference, ModelRef}, + }; + let first = MockServer::start().await; + let second = MockServer::start().await; + Mock::given(path("/chat/completions")) + .respond_with(ResponseTemplate::new(503)) + .expect(1) + .mount(&first) + .await; + Mock::given(path("/responses")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&first) + .await; + Mock::given(path("/chat/completions")) + .respond_with( + ResponseTemplate::new(400) + .set_body_json(json!({"error":{"code":"unsupported_api_for_model"}})), + ) + .expect(1) + .mount(&second) + .await; + Mock::given(path("/responses")).and(header("authorization", "Bearer second-key")) + .and(body_partial_json(json!({"model":"second-model"}))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"actual deliverable"}]}]}))) + .expect(1).mount(&second).await; + let mut config = Config::from_env().unwrap(); + for (tag, endpoint) in [("first", first.uri()), ("second", second.uri())] { + config.providers.insert( + tag.into(), + ProviderEndpoint { + tag: tag.into(), + endpoint, + api_key: Some(format!("{tag}-key")), + }, + ); + } + let state = test_state(config); + let mut base = UpstreamConfig::azure(first.uri(), "default".into(), "test".into()); + base.provider_api_key = Some("default-key".into()); + let policy = InferencePolicySnapshot { + model_preference: Some(ModelPreference { + primary: ModelRef { + provider: "first".into(), + deployment: "first-model".into(), + }, + fallback: vec![ModelRef { + provider: "second".into(), + deployment: "second-model".into(), + }], + }), + ..Default::default() + }; + let (_, _, _, selected) = + forward_chat(&state, &base, &policy, &HeaderMap::new(), Bytes::from("{}")) + .await + .unwrap(); + assert_eq!(selected.endpoint, second.uri()); + let (_, _, body, winner) = forward_responses( + &state, + &base, + &policy, + &selected, + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await + .unwrap(); + assert_eq!(winner.endpoint, second.uri()); + assert!(String::from_utf8_lossy(&body).contains("actual deliverable")); + } + + #[test] + fn model_errors_do_not_confuse_auth_policy_or_protocol_failures() { + for code in ["DeploymentNotFound", "model_not_found", "invalid_model"] { + let body = serde_json::to_vec(&json!({"error": {"code": code}})).unwrap(); + assert!(is_model_unavailable_error(StatusCode::NOT_FOUND, &body)); + assert!(!is_model_unavailable_error(StatusCode::FORBIDDEN, &body)); + assert!(!is_model_unavailable_error( + StatusCode::TOO_MANY_REQUESTS, + &body + )); + } + let body = br#"{"error":{"code":"unsupported_api_for_model","message":"model not supported via chat API"}}"#; + assert!(is_responses_only_error(body)); + assert!(!is_model_unavailable_error(StatusCode::BAD_REQUEST, body)); + assert!(!is_responses_only_error( + br#"{"error":{"message":"unsupported tool argument"}}"# + )); + } + + #[test] + fn capability_cache_separates_endpoints_and_models() { + let a = UpstreamConfig::azure("https://a.example/v1".into(), "model".into(), "test".into()); + let mut b = a.clone(); + b.endpoint = "https://b.example/v1".into(); + assert_ne!(model_capability_key(&a), model_capability_key(&b)); + b = a.clone(); + b.deployment = "Model".into(); + assert_ne!(model_capability_key(&a), model_capability_key(&b)); + } +} diff --git a/inference-router/src/routes/model_routing_closure_tests.rs b/inference-router/src/routes/model_routing_closure_tests.rs new file mode 100644 index 00000000..b2d2be0f --- /dev/null +++ b/inference-router/src/routes/model_routing_closure_tests.rs @@ -0,0 +1,311 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::regressions::{config, policy, post_chat, read_request, router, state}; +use super::*; +use crate::inference_policy_loader::ModelRef; +use crate::provider::ProviderKind; +use serde_json::json; +use tokio::io::AsyncWriteExt; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_partial_json, header, method, path}, +}; + +fn chat_response(stream: bool, text: &str) -> ResponseTemplate { + if stream { + ResponseTemplate::new(200).set_body_raw( + format!( + "data: {}\n\ndata: [DONE]\n\n", + json!({"choices":[{"delta":{"content":text}}]}) + ), + "text/event-stream", + ) + } else { + ResponseTemplate::new(200).set_body_json(json!({"choices":[{"message":{"content":text}}]})) + } +} + +fn responses_response(text: &str) -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":text}]}], + })) +} + +#[tokio::test] +async fn legacy_primary_metadata_keeps_foundry_with_or_without_native_credentials() { + for tag in ["anthropic", "ollama", "bedrock", "unknown-informational"] { + for native_key in [None, Some("native-key")] { + for stream in [false, true] { + let default = MockServer::start().await; + let native = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&native) + .await; + let responses_only = tag == "anthropic" && native_key.is_some(); + let chat = if responses_only { + ResponseTemplate::new(400) + .set_body_json(json!({"error":{"code":"unsupported_api_for_model"}})) + } else { + chat_response(stream, "legacy-default") + }; + Mock::given(path("/openai/v1/chat/completions")) + .and(header("authorization", "Bearer ambient-default-key")) + .and(body_partial_json(json!({"model":"claude-prod"}))) + .respond_with(chat) + .expect(1) + .mount(&default) + .await; + if responses_only { + Mock::given(path("/openai/v1/responses")) + .and(header("authorization", "Bearer ambient-default-key")) + .and(body_partial_json(json!({"model":"claude-prod"}))) + .respond_with(responses_response("legacy-default")) + .expect(1) + .mount(&default) + .await; + } + let mut config = config(&format!("{}/openai/v1", default.uri()), &[]); + config.anthropic_endpoint = native.uri(); + config.anthropic_api_key = native_key.map(str::to_string); + config.ollama_endpoint = Some(native.uri()); + let policy = policy(tag, "claude-prod"); + let (status, body) = post_chat(router(state(config), &policy).await, stream).await; + assert_eq!( + status, + StatusCode::OK, + "{tag}, native key {}, stream {stream}", + native_key.is_some() + ); + assert!(String::from_utf8_lossy(&body).contains("legacy-default")); + } + } + } +} + +#[tokio::test] +async fn registered_named_intent_and_explicit_native_intent_remain_separate() { + let default = MockServer::start().await; + let registered = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&default) + .await; + Mock::given(path("/chat/completions")) + .and(header("authorization", "Bearer registered-key")) + .respond_with(chat_response(false, "registered")) + .expect(1) + .mount(®istered) + .await; + let mut config = config( + &default.uri(), + &[("anthropic", ®istered.uri(), Some("registered-key"))], + ); + config.anthropic_endpoint = registered.uri(); + config.anthropic_api_key = Some("native-key".into()); + let state = state(config); + let metadata = policy("anthropic", "claude-prod"); + let mut explicit = metadata.clone(); + explicit.provider = Some("anthropic".into()); + let base = UpstreamConfig::azure(default.uri(), "default".into(), "test".into()); + let named = primary_target(&state, &base, &metadata, b"{}").unwrap(); + let native = primary_target(&state, &base, &explicit, b"{}").unwrap(); + assert_eq!(named.provider, ProviderKind::AzureOpenAI); + assert_eq!(named.provider_api_key.as_deref(), Some("registered-key")); + assert_eq!(native.provider, ProviderKind::Anthropic); + assert_eq!(native.api_key.as_deref(), Some("native-key")); + assert_ne!(model_capability_key(&named), model_capability_key(&native)); + let (status, body) = post_chat(router(state, &metadata).await, false).await; + assert_eq!(status, StatusCode::OK); + assert!(String::from_utf8_lossy(&body).contains("registered")); +} + +#[tokio::test] +async fn explicit_native_intent_still_requires_native_credentials() { + let default = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&default) + .await; + let mut config = config( + &default.uri(), + &[("anthropic", &default.uri(), Some("registered-key"))], + ); + config.anthropic_api_key = None; + let state = state(config); + let mut policy = policy("azure-openai", "claude-prod"); + policy.provider = Some("anthropic".into()); + let (status, _) = post_chat(router(state, &policy).await, false).await; + assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE); +} + +#[tokio::test] +async fn unhealthy_legacy_metadata_does_not_suppress_an_explicit_native_fallback() { + let default = MockServer::start().await; + let native = MockServer::start().await; + Mock::given(path("/chat/completions")) + .respond_with(ResponseTemplate::new(503)) + .expect(3) + .mount(&default) + .await; + Mock::given(path("/v1/chat/completions")) + .and(body_partial_json(json!({"model":"llama-prod"}))) + .respond_with(chat_response(false, "native-fallback")) + .expect(4) + .mount(&native) + .await; + let mut config = config(&default.uri(), &[]); + config.ollama_endpoint = Some(native.uri()); + let state = state(config); + let mut policy = policy("ollama", "llama-prod"); + policy + .model_preference + .as_mut() + .unwrap() + .fallback + .push(ModelRef { + provider: "ollama".into(), + deployment: "llama-prod".into(), + }); + let app = router(state, &policy).await; + for _ in 0..4 { + let (status, body) = post_chat(app.clone(), false).await; + assert_eq!(status, StatusCode::OK); + assert!(String::from_utf8_lossy(&body).contains("native-fallback")); + } +} + +async fn rejecting_backend( + status: u16, + responses_recovery: bool, +) -> (String, tokio::task::JoinHandle<()>) { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + assert!( + read_request(&mut socket) + .await + .starts_with("POST /chat/completions ") + ); + if responses_recovery { + let body = r#"{"error":{"code":"unsupported_api_for_model"}}"#; + socket.write_all(format!( + "HTTP/1.1 400 Bad Request\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len(), + ).as_bytes()).await.unwrap(); + socket.shutdown().await.unwrap(); + let (next, _) = listener.accept().await.unwrap(); + socket = next; + assert!( + read_request(&mut socket) + .await + .starts_with("POST /responses ") + ); + } + socket.write_all(format!( + "HTTP/1.1 {status} Rejected\r\nContent-Type: application/json\r\nContent-Length: 10000\r\nConnection: close\r\n\r\n{{\"error\":" + ).as_bytes()).await.unwrap(); + socket.shutdown().await.unwrap(); + }); + (endpoint, server) +} + +#[tokio::test] +async fn truncated_known_503_and_429_rejections_reach_healthy_fallback_in_both_modes() { + for status in [503, 429] { + for stream in [false, true] { + for responses_recovery in [false, true] { + let (endpoint, server) = rejecting_backend(status, responses_recovery).await; + let backup = MockServer::start().await; + let path_name = if responses_recovery { + "/responses" + } else { + "/chat/completions" + }; + let reply = if responses_recovery { + responses_response("healthy-fallback") + } else { + chat_response(stream, "healthy-fallback") + }; + Mock::given(path(path_name)) + .and(header("authorization", "Bearer backup-key")) + .and(body_partial_json(json!({"model":"backup-model"}))) + .respond_with(reply) + .expect(1) + .mount(&backup) + .await; + let state = state(config( + &backup.uri(), + &[ + ("primary", &endpoint, Some("primary-key")), + ("backup", &backup.uri(), Some("backup-key")), + ], + )); + let mut policy = policy("primary", "primary-model"); + policy + .model_preference + .as_mut() + .unwrap() + .fallback + .push(ModelRef { + provider: "backup".into(), + deployment: "backup-model".into(), + }); + let (actual, body) = post_chat(router(state, &policy).await, stream).await; + assert_eq!( + actual, + StatusCode::OK, + "{status}, stream {stream}, Responses {responses_recovery}" + ); + assert!(String::from_utf8_lossy(&body).contains("healthy-fallback")); + tokio::time::timeout(std::time::Duration::from_secs(5), server) + .await + .unwrap() + .unwrap(); + } + } + } +} + +#[tokio::test] +async fn truncated_ordinary_client_rejections_are_never_replayed() { + for status in [400, 401, 403, 404, 422] { + for stream in [false, true] { + let (endpoint, server) = rejecting_backend(status, false).await; + let backup = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&backup) + .await; + let state = state(config( + &backup.uri(), + &[ + ("primary", &endpoint, Some("primary-key")), + ("backup", &backup.uri(), Some("backup-key")), + ], + )); + let mut policy = policy("primary", "primary-model"); + policy + .model_preference + .as_mut() + .unwrap() + .fallback + .push(ModelRef { + provider: "backup".into(), + deployment: "backup-model".into(), + }); + let (actual, _) = post_chat(router(state, &policy).await, stream).await; + assert_eq!(actual, StatusCode::BAD_GATEWAY); + tokio::time::timeout(std::time::Duration::from_secs(5), server) + .await + .unwrap() + .unwrap(); + } + } +} diff --git a/inference-router/src/routes/model_routing_regressions.rs b/inference-router/src/routes/model_routing_regressions.rs new file mode 100644 index 00000000..d8e568d8 --- /dev/null +++ b/inference-router/src/routes/model_routing_regressions.rs @@ -0,0 +1,456 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::{ + auth::WorkloadIdentityAuth, + config::{Config, ProviderEndpoint}, + inference_policy_loader::{LoadedInferencePolicy, ModelPreference, ModelRef}, + proxy::failure::{Acceptance, FailureCategory, ForwardFailure}, +}; +use axum::{Router, body::Body, http::Request}; +use serde_json::json; +use std::sync::Arc; +use tower::ServiceExt; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_partial_json, header, method}, +}; + +pub(super) fn config(default: &str, providers: &[(&str, &str, Option<&str>)]) -> Config { + let mut config = Config::from_env().unwrap(); + config.azure_openai_endpoint = Some(default.into()); + config.default_model = "true-default-model".into(); + config.content_safety_enabled = true; + config.prompt_shields_enabled = true; + config.providers = providers + .iter() + .map(|(id, endpoint, credential)| { + ( + id.to_string(), + ProviderEndpoint { + tag: id.to_string(), + endpoint: endpoint.to_string(), + api_key: credential.map(str::to_string), + }, + ) + }) + .collect(); + config +} + +pub(super) fn policy(provider: &str, model: &str) -> InferencePolicySnapshot { + InferencePolicySnapshot { + model_preference: Some(ModelPreference { + primary: ModelRef { + provider: provider.into(), + deployment: model.into(), + }, + fallback: vec![], + }), + ..Default::default() + } +} + +pub(super) fn state(config: Config) -> AppState { + let mut state = super::tests::test_state(config); + state.auth = Arc::new(WorkloadIdentityAuth::for_test( + Some("ambient-default-key"), + None, + )); + state.copilot = Arc::new(crate::copilot_auth::CopilotTokenCache::with_test_exchange( + "test-default-seat", + "http://127.0.0.1:1/unexpected-exchange".into(), + )); + state.client = reqwest::Client::builder().no_proxy().build().unwrap(); + state.budget = crate::budget::TokenBudgetTracker::new(1_000_000, 1_000_000); + state +} + +pub(super) async fn router(state: AppState, policy: &InferencePolicySnapshot) -> Router { + *state.inference_policy.write().await = Some(LoadedInferencePolicy { + digest: "routing-regression".into(), + source_path: "routing-regression".into(), + per_request_tokens: None, + daily_tokens: None, + monthly_tokens: None, + content_safety: policy.content_safety.clone(), + model_preference: policy.model_preference.clone(), + provider: policy.provider.clone(), + guardrails: vec![], + raw: json!({}), + }); + Router::new() + .merge(crate::routes::inference_routes()) + .with_state(state) +} + +pub(super) async fn post_chat(app: Router, stream: bool) -> (StatusCode, Bytes) { + let response = app.oneshot(Request::builder().method("POST") + .uri("/v1/chat/completions").header("content-type", "application/json") + .body(Body::from(json!({ + "model": "client-requested", "messages": [{"role":"user","content":"hello"}], + "stream": stream, + }).to_string())).unwrap()).await.unwrap(); + let status = response.status(); + let body = axum::body::to_bytes(response.into_body(), 1_048_576) + .await + .unwrap(); + (status, body) +} + +#[tokio::test] +async fn explicit_provider_wins_over_informational_primary_for_buffered_and_streaming_chat() { + for stream in [false, true] { + let local = MockServer::start().await; + let default = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&default) + .await; + let response = if stream { + ResponseTemplate::new(200).set_body_raw( + "data: {\"choices\":[{\"delta\":{\"content\":\"local\"}}]}\n\ndata: [DONE]\n\n", + "text/event-stream", + ) + } else { + ResponseTemplate::new(200) + .set_body_json(json!({"choices":[{"message":{"content":"local"}}]})) + }; + Mock::given(body_partial_json(json!({"model":"policy-model"}))) + .and(|request: &wiremock::Request| { + request.url.path() == "/v1/chat/completions" + && !request.headers.contains_key("authorization") + }) + .respond_with(response) + .expect(1) + .mount(&local) + .await; + let mut config = config(&default.uri(), &[]); + config.ollama_endpoint = Some(local.uri()); + let mut policy = policy("azure-openai", "policy-model"); + policy.provider = Some("ollama".into()); + let base = UpstreamConfig::azure(default.uri(), "true-default-model".into(), "test".into()); + let candidates = failover::build_candidates(&base, &policy); + assert_eq!(candidates[0].provider.as_deref(), Some("ollama")); + assert_eq!(candidates.last().unwrap().provider, None); + assert_eq!(candidates.last().unwrap().deployment, "true-default-model"); + let (status, body) = post_chat(router(state(config), &policy).await, stream).await; + assert_eq!(status, StatusCode::OK); + assert!(String::from_utf8_lossy(&body).contains("local")); + } +} + +#[tokio::test] +async fn availability_cache_does_not_skip_another_credential_on_the_same_endpoint_and_model() { + let shared = MockServer::start().await; + let default = MockServer::start().await; + Mock::given(header("authorization", "Bearer credential-a")) + .respond_with( + ResponseTemplate::new(404).set_body_json(json!({"error":{"code":"model_not_found"}})), + ) + .expect(1) + .mount(&shared) + .await; + Mock::given(header("authorization", "Bearer credential-b")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"choices":[]}))) + .expect(2) + .mount(&shared) + .await; + Mock::given(header("authorization", "Bearer ambient-default-key")) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&default) + .await; + let state = state(config( + &default.uri(), + &[ + ("account-a", &shared.uri(), Some("credential-a")), + ("account-b", &shared.uri(), Some("credential-b")), + ], + )); + let base = UpstreamConfig::azure(default.uri(), "true-default-model".into(), "test".into()); + let a = policy("account-a", "same-model"); + let b = policy("account-b", "same-model"); + forward_chat(&state, &base, &a, &HeaderMap::new(), Bytes::from("{}")) + .await + .unwrap(); + let (status, _, _, selected) = + forward_chat(&state, &base, &b, &HeaderMap::new(), Bytes::from("{}")) + .await + .unwrap(); + assert_eq!(status, StatusCode::OK); + assert_eq!(selected.endpoint, shared.uri()); + let (_, _, stream, selected) = + forward_stream_chat(&state, &base, &b, HeaderMap::new(), Bytes::from("{}")) + .await + .unwrap(); + let _: Vec<_> = stream.try_collect().await.unwrap(); + assert_eq!(selected.endpoint, shared.uri()); + let a = primary_target(&state, &base, &a, b"{}").unwrap(); + let b = primary_target(&state, &base, &b, b"{}").unwrap(); + assert_ne!(model_capability_key(&a), model_capability_key(&b)); + let cache = state.unavailable_models.read().unwrap(); + assert!(cache.contains(&model_capability_key(&a))); + assert!(!cache.contains(&model_capability_key(&b))); + for key in cache.iter() { + assert!(!key.contains("credential-a") && !key.contains("credential-b")); + } +} + +pub(super) async fn read_request(socket: &mut tokio::net::TcpStream) -> String { + use tokio::io::AsyncReadExt; + let mut bytes = Vec::new(); + loop { + let mut chunk = [0; 4096]; + let size = socket.read(&mut chunk).await.unwrap(); + assert!(size > 0, "request closed before its body was sent"); + bytes.extend_from_slice(&chunk[..size]); + if let Some(end) = bytes.windows(4).position(|window| window == b"\r\n\r\n") { + let headers = String::from_utf8_lossy(&bytes[..end]); + let length: usize = headers + .lines() + .find_map(|line| { + line.split_once(':') + .filter(|(name, _)| name.eq_ignore_ascii_case("content-length")) + .map(|(_, value)| value.trim().parse().unwrap()) + }) + .unwrap_or(0); + if bytes.len() >= end + 4 + length { + return String::from_utf8(bytes).unwrap(); + } + } + } +} + +#[tokio::test] +async fn actual_chat_to_responses_recovery_never_replays_after_success_headers() { + use tokio::io::AsyncWriteExt; + for stream in [false, true] { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let (mut chat, _) = listener.accept().await.unwrap(); + assert!( + read_request(&mut chat) + .await + .starts_with("POST /chat/completions ") + ); + let body = r#"{"error":{"code":"unsupported_api_for_model"}}"#; + chat.write_all(format!( + "HTTP/1.1 400 Bad Request\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len(), + ).as_bytes()).await.unwrap(); + chat.shutdown().await.unwrap(); + for _ in 0..2 { + let (mut responses, _) = listener.accept().await.unwrap(); + assert!( + read_request(&mut responses) + .await + .starts_with("POST /responses ") + ); + responses.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: 10000\r\nConnection: close\r\n\r\n{\"output\":[").await.unwrap(); + responses.shutdown().await.unwrap(); + } + }); + let fallback = MockServer::start().await; + Mock::given(method("POST")).respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "output": [{"type":"message","role":"assistant","content":[{"type":"output_text","text":"replayed"}]}], + }))).expect(0).mount(&fallback).await; + let state = state(config( + &fallback.uri(), + &[ + ("primary", &endpoint, Some("primary-key")), + ("fallback", &fallback.uri(), Some("fallback-key")), + ], + )); + let mut policy = policy("primary", "primary-model"); + policy + .model_preference + .as_mut() + .unwrap() + .fallback + .push(ModelRef { + provider: "fallback".into(), + deployment: "fallback-model".into(), + }); + let app = router(state, &policy).await; + let (status, body) = post_chat(app.clone(), stream).await; + assert_eq!(status, StatusCode::BAD_GATEWAY); + let body = String::from_utf8_lossy(&body); + assert!(body.contains("error")); + assert!(!body.contains("replayed")); + let (status, body) = post_chat(app, true).await; + assert_eq!(status, StatusCode::OK); + let body = String::from_utf8_lossy(&body); + assert!(body.contains("error")); + assert!(!body.contains("replayed")); + tokio::time::timeout(std::time::Duration::from_secs(5), server) + .await + .expect("both chat and Responses requests must reach the test backend") + .unwrap(); + } +} + +#[tokio::test] +async fn auth_and_configuration_acquisition_failures_do_not_attempt_fallback() { + for provider in ["copilot", "copilot-exchange", "ollama", "bedrock"] { + let fallback = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&fallback) + .await; + let mut config = config( + &fallback.uri(), + &[ + ("copilot", "https://api.githubcopilot.com", None), + ( + "copilot-exchange", + "https://api.githubcopilot.com", + Some("test-seat"), + ), + ("fallback", &fallback.uri(), Some("fallback-key")), + ], + ); + config.ollama_endpoint = None; + let state = state(config); + let mut policy = policy(provider, "primary-model"); + if matches!(provider, "ollama" | "bedrock") { + policy.provider = Some(provider.into()); + } + policy + .model_preference + .as_mut() + .unwrap() + .fallback + .push(ModelRef { + provider: "fallback".into(), + deployment: "fallback-model".into(), + }); + let base = + UpstreamConfig::azure(fallback.uri(), "true-default-model".into(), "test".into()); + let buffered = forward_chat(&state, &base, &policy, &HeaderMap::new(), Bytes::from("{}")) + .await + .err() + .expect("auth/configuration must fail before inference HTTP"); + let streamed = forward_stream_chat( + &state, + &base, + &policy, + HeaderMap::new(), + Bytes::from(r#"{"stream":true}"#), + ) + .await + .err() + .expect("streaming auth/configuration must fail before inference HTTP"); + for error in [buffered, streamed] { + let failure = error.downcast_ref::().unwrap(); + assert!(matches!( + failure.category, + FailureCategory::Authentication | FailureCategory::Configuration + )); + assert_eq!(failure.acceptance, Acceptance::NotAccepted); + } + assert!( + state + .deployment_health + .snapshot() + .iter() + .all(|entry| entry.failure_streak == 0) + ); + } +} + +#[tokio::test] +async fn connection_failure_before_acceptance_can_still_fail_over() { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let unavailable = format!("http://{}", listener.local_addr().unwrap()); + drop(listener); + let fallback = MockServer::start().await; + Mock::given(header("authorization", "Bearer fallback-key")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"choices":[]}))) + .expect(2) + .mount(&fallback) + .await; + let state = state(config( + &fallback.uri(), + &[ + ("unavailable", &unavailable, None), + ("fallback", &fallback.uri(), Some("fallback-key")), + ], + )); + let mut policy = policy("unavailable", "model"); + policy + .model_preference + .as_mut() + .unwrap() + .fallback + .push(ModelRef { + provider: "fallback".into(), + deployment: "model".into(), + }); + let base = UpstreamConfig::azure(fallback.uri(), "true-default-model".into(), "test".into()); + let (status, _, _, selected) = + forward_chat(&state, &base, &policy, &HeaderMap::new(), Bytes::from("{}")) + .await + .unwrap(); + assert_eq!(status, StatusCode::OK); + assert_eq!(selected.endpoint, fallback.uri()); + let (status, _, stream, selected) = + forward_stream_chat(&state, &base, &policy, HeaderMap::new(), Bytes::from("{}")) + .await + .unwrap(); + assert_eq!(status, StatusCode::OK); + assert_eq!(selected.endpoint, fallback.uri()); + let _: Vec<_> = stream.try_collect().await.unwrap(); +} + +#[tokio::test] +async fn connection_closed_after_sending_has_unknown_acceptance_and_is_not_replayed() { + use tokio::io::AsyncWriteExt; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let endpoint = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + read_request(&mut socket).await; + socket.shutdown().await.unwrap(); + }); + let fallback = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&fallback) + .await; + let state = state(config( + &fallback.uri(), + &[ + ("primary", &endpoint, None), + ("fallback", &fallback.uri(), Some("fallback-key")), + ], + )); + let mut policy = policy("primary", "model"); + policy + .model_preference + .as_mut() + .unwrap() + .fallback + .push(ModelRef { + provider: "fallback".into(), + deployment: "model".into(), + }); + let base = UpstreamConfig::azure(fallback.uri(), "true-default-model".into(), "test".into()); + let result = forward_chat(&state, &base, &policy, &HeaderMap::new(), Bytes::from("{}")).await; + let error = match result { + Err(error) => error, + Ok(_) => panic!("lost response must fail closed"), + }; + let failure = error.downcast_ref::().unwrap(); + assert_eq!(failure.category, FailureCategory::Transport); + assert_eq!(failure.acceptance, Acceptance::Unknown); + tokio::time::timeout(std::time::Duration::from_secs(5), server) + .await + .expect("the request must reach the test backend") + .unwrap(); +} diff --git a/inference-router/src/sidecar_client.rs b/inference-router/src/sidecar_client.rs index 5e9ac61a..7575ccd0 100644 --- a/inference-router/src/sidecar_client.rs +++ b/inference-router/src/sidecar_client.rs @@ -116,6 +116,17 @@ pub struct SidecarClient { } impl SidecarClient { + #[cfg(test)] + pub(crate) fn for_test(base_url: String) -> Self { + Self { + base_url, + pinned_agent_id: "test-agent".into(), + expected_tenant_id: Some("test-tenant".into()), + client: reqwest::Client::new(), + cache: Arc::new(RwLock::new(HashMap::new())), + } + } + /// Construct from `AUTH_SIDECAR_URL` + `PINNED_AGENT_IDENTITY_APP_ID`. /// Returns: /// - `Ok(None)` when BOTH env vars are absent — sidecar mode is diff --git a/inference-router/tests/agt_governance_integration.rs b/inference-router/tests/agt_governance_integration.rs index 190a1dde..76c4b15e 100644 --- a/inference-router/tests/agt_governance_integration.rs +++ b/inference-router/tests/agt_governance_integration.rs @@ -61,6 +61,7 @@ fn test_state(sandbox: &str, admin_token: Option<&str>) -> AppState { openai_moderation_endpoint: "https://api.openai.com".into(), openai_moderation_api_key: None, openai_moderation_model: "omni-moderation-latest".into(), + providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), policy_provider: Arc::clone(&governance) as Arc, @@ -77,6 +78,7 @@ fn test_state(sandbox: &str, admin_token: Option<&str>) -> AppState { model_override: Arc::new(std::sync::RwLock::new(None)), admin_token: admin_token.map(|t| Arc::new(t.to_string())), responses_only_models: Arc::new(std::sync::RwLock::new(Default::default())), + unavailable_models: Default::default(), handoff_tokens: HandoffTokenStore::new(), handoff_session: HandoffSession::new(), drain_state: DrainState::new(), diff --git a/inference-router/tests/anthropic_buffered_guardrail.rs b/inference-router/tests/anthropic_buffered_guardrail.rs index d23733e3..750fe690 100644 --- a/inference-router/tests/anthropic_buffered_guardrail.rs +++ b/inference-router/tests/anthropic_buffered_guardrail.rs @@ -72,6 +72,7 @@ fn test_state(anthropic_endpoint: String, moderation_endpoint: String) -> AppSta openai_moderation_endpoint: moderation_endpoint, openai_moderation_api_key: Some("sk-mod-test".into()), openai_moderation_model: "omni-moderation-latest".into(), + providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000_000, 1_000_000_000), policy_provider: Arc::clone(&governance) as Arc, @@ -86,6 +87,7 @@ fn test_state(anthropic_endpoint: String, moderation_endpoint: String) -> AppSta model_override: Arc::new(std::sync::RwLock::new(None)), admin_token: None, responses_only_models: Arc::new(std::sync::RwLock::new(Default::default())), + unavailable_models: Default::default(), handoff_tokens: HandoffTokenStore::new(), handoff_session: HandoffSession::new(), drain_state: DrainState::new(), diff --git a/inference-router/tests/chat_output_guardrail_nonjson.rs b/inference-router/tests/chat_output_guardrail_nonjson.rs index 2e49d415..99db094e 100644 --- a/inference-router/tests/chat_output_guardrail_nonjson.rs +++ b/inference-router/tests/chat_output_guardrail_nonjson.rs @@ -70,6 +70,7 @@ fn test_state(ollama_endpoint: String, moderation_endpoint: String) -> AppState openai_moderation_endpoint: moderation_endpoint, openai_moderation_api_key: Some("sk-mod-test".into()), openai_moderation_model: "omni-moderation-latest".into(), + providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000_000, 1_000_000_000), policy_provider: Arc::clone(&governance) as Arc, @@ -84,6 +85,7 @@ fn test_state(ollama_endpoint: String, moderation_endpoint: String) -> AppState model_override: Arc::new(std::sync::RwLock::new(None)), admin_token: None, responses_only_models: Arc::new(std::sync::RwLock::new(Default::default())), + unavailable_models: Default::default(), handoff_tokens: HandoffTokenStore::new(), handoff_session: HandoffSession::new(), drain_state: DrainState::new(), diff --git a/inference-router/tests/egress_blocked_endpoint.rs b/inference-router/tests/egress_blocked_endpoint.rs index 125e79e6..843a1777 100644 --- a/inference-router/tests/egress_blocked_endpoint.rs +++ b/inference-router/tests/egress_blocked_endpoint.rs @@ -59,6 +59,7 @@ fn test_state() -> AppState { openai_moderation_endpoint: "https://api.openai.com".into(), openai_moderation_api_key: None, openai_moderation_model: "omni-moderation-latest".into(), + providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), policy_provider: Arc::clone(&governance) as Arc, @@ -73,6 +74,7 @@ fn test_state() -> AppState { model_override: Arc::new(std::sync::RwLock::new(None)), admin_token: None, responses_only_models: Arc::new(std::sync::RwLock::new(Default::default())), + unavailable_models: Default::default(), handoff_tokens: HandoffTokenStore::new(), handoff_session: HandoffSession::new(), drain_state: DrainState::new(), diff --git a/inference-router/tests/failover_walk.rs b/inference-router/tests/failover_walk.rs index cb7e2579..b499862c 100644 --- a/inference-router/tests/failover_walk.rs +++ b/inference-router/tests/failover_walk.rs @@ -136,15 +136,18 @@ async fn primary_503_falls_through_to_fallback_200() { sandbox_name: "sbx".into(), provider: ProviderKind::AzureOpenAI, api_key: None, + provider_api_key: None, + authentication: Default::default(), }; let snap = snapshot("primary-down", &["fallback-up"]); - let (status, _hdrs, body) = forward_with_failover( + let (status, _hdrs, body, _selected) = forward_with_failover( &auth, None, &client, &health, &upstream, + &kars_inference_router::config::Config::from_env().unwrap(), &snap, Method::POST, "chat/completions", @@ -163,11 +166,11 @@ async fn primary_503_falls_through_to_fallback_200() { let snaps = health.snapshot(); let primary = snaps .iter() - .find(|s| s.deployment == "primary-down") + .find(|s| s.deployment == "Foundry::primary-down") .expect("primary health entry missing"); let fallback = snaps .iter() - .find(|s| s.deployment == "fallback-up") + .find(|s| s.deployment == "Foundry::fallback-up") .expect("fallback health entry missing"); assert_eq!(primary.failure_streak, 1); assert!(primary.healthy, "single 503 under 3-strike threshold"); @@ -190,9 +193,9 @@ async fn unhealthy_primary_is_skipped_in_second_pass() { // Pre-mark primary unhealthy (3 failures = at threshold). for _ in 0..3 { - health.record_failure("primary-down"); + health.record_failure("Foundry::primary-down"); } - assert!(!health.is_healthy("primary-down")); + assert!(!health.is_healthy("Foundry::primary-down")); let upstream = UpstreamConfig { endpoint: base, @@ -200,15 +203,18 @@ async fn unhealthy_primary_is_skipped_in_second_pass() { sandbox_name: "sbx".into(), provider: ProviderKind::AzureOpenAI, api_key: None, + provider_api_key: None, + authentication: Default::default(), }; let snap = snapshot("primary-down", &["fallback-up"]); - let (status, _hdrs, _body) = forward_with_failover( + let (status, _hdrs, _body, _selected) = forward_with_failover( &auth, None, &client, &health, &upstream, + &kars_inference_router::config::Config::from_env().unwrap(), &snap, Method::POST, "chat/completions", @@ -223,7 +229,7 @@ async fn unhealthy_primary_is_skipped_in_second_pass() { let snaps = health.snapshot(); let primary = snaps .iter() - .find(|s| s.deployment == "primary-down") + .find(|s| s.deployment == "Foundry::primary-down") .unwrap(); assert_eq!(primary.failure_streak, 3, "no new failure recorded"); } @@ -243,11 +249,12 @@ async fn all_unhealthy_still_punches_primary_for_last_resort() { // Mark BOTH unhealthy. for _ in 0..3 { + health.record_failure("Foundry::primary-down"); + health.record_failure("Foundry::fallback-up"); health.record_failure("primary-down"); - health.record_failure("fallback-up"); } - assert!(!health.is_healthy("primary-down")); - assert!(!health.is_healthy("fallback-up")); + assert!(!health.is_healthy("Foundry::primary-down")); + assert!(!health.is_healthy("Foundry::fallback-up")); let upstream = UpstreamConfig { endpoint: base, @@ -255,15 +262,18 @@ async fn all_unhealthy_still_punches_primary_for_last_resort() { sandbox_name: "sbx".into(), provider: ProviderKind::AzureOpenAI, api_key: None, + provider_api_key: None, + authentication: Default::default(), }; let snap = snapshot("primary-down", &["fallback-up"]); - let (status, _hdrs, _body) = forward_with_failover( + let (status, _hdrs, _body, _selected) = forward_with_failover( &auth, None, &client, &health, &upstream, + &kars_inference_router::config::Config::from_env().unwrap(), &snap, Method::POST, "chat/completions", @@ -279,7 +289,7 @@ async fn all_unhealthy_still_punches_primary_for_last_resort() { let snaps = health.snapshot(); let primary = snaps .iter() - .find(|s| s.deployment == "primary-down") + .find(|s| s.deployment == "Foundry::primary-down") .unwrap(); assert_eq!(primary.failure_streak, 4); } diff --git a/inference-router/tests/foundry_route_guard.rs b/inference-router/tests/foundry_route_guard.rs index e8e3eab8..76d5e3f0 100644 --- a/inference-router/tests/foundry_route_guard.rs +++ b/inference-router/tests/foundry_route_guard.rs @@ -80,6 +80,7 @@ fn test_state() -> AppState { openai_moderation_endpoint: "https://api.openai.com".into(), openai_moderation_api_key: Some("sk-mod-test".into()), openai_moderation_model: "omni-moderation-latest".into(), + providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), policy_provider: Arc::clone(&governance) as Arc, @@ -94,6 +95,7 @@ fn test_state() -> AppState { model_override: Arc::new(std::sync::RwLock::new(None)), admin_token: None, responses_only_models: Arc::new(std::sync::RwLock::new(Default::default())), + unavailable_models: Default::default(), handoff_tokens: HandoffTokenStore::new(), handoff_session: HandoffSession::new(), drain_state: DrainState::new(), diff --git a/inference-router/tests/multi_provider_guardrails.rs b/inference-router/tests/multi_provider_guardrails.rs index e839db6c..61601ecd 100644 --- a/inference-router/tests/multi_provider_guardrails.rs +++ b/inference-router/tests/multi_provider_guardrails.rs @@ -41,6 +41,7 @@ fn config_with_moderation(endpoint: &str, api_key: Option<&str>) -> Config { openai_moderation_endpoint: endpoint.to_string(), openai_moderation_api_key: api_key.map(String::from), openai_moderation_model: "omni-moderation-latest".into(), + providers: Default::default(), } } @@ -65,6 +66,8 @@ async fn ollama_provider_forwards_openai_compat_without_auth() { sandbox_name: "test-sandbox".into(), provider: ProviderKind::Ollama, api_key: None, + provider_api_key: None, + authentication: Default::default(), }; let (status, _headers, resp) = forward( @@ -120,6 +123,8 @@ async fn anthropic_provider_forwards_messages_with_router_held_key() { sandbox_name: "test-sandbox".into(), provider: ProviderKind::Anthropic, api_key: Some("sk-ant-router-held".into()), + provider_api_key: None, + authentication: Default::default(), }; // The inbound request carries an agent-supplied x-api-key that diff --git a/inference-router/tests/policy_status_endpoint.rs b/inference-router/tests/policy_status_endpoint.rs index b9355f27..b40f633b 100644 --- a/inference-router/tests/policy_status_endpoint.rs +++ b/inference-router/tests/policy_status_endpoint.rs @@ -70,6 +70,7 @@ fn test_state() -> (AppState, Arc) { openai_moderation_endpoint: "https://api.openai.com".into(), openai_moderation_api_key: None, openai_moderation_model: "omni-moderation-latest".into(), + providers: Default::default(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), policy_provider: Arc::clone(&governance) as Arc, @@ -84,6 +85,7 @@ fn test_state() -> (AppState, Arc) { model_override: Arc::new(std::sync::RwLock::new(None)), admin_token: None, responses_only_models: Arc::new(std::sync::RwLock::new(Default::default())), + unavailable_models: Default::default(), handoff_tokens: HandoffTokenStore::new(), handoff_session: HandoffSession::new(), drain_state: DrainState::new(), diff --git a/inference-router/tests/provider_failover.rs b/inference-router/tests/provider_failover.rs new file mode 100644 index 00000000..437c7cea --- /dev/null +++ b/inference-router/tests/provider_failover.rs @@ -0,0 +1,357 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use axum::http::{HeaderMap, Method, StatusCode}; +use bytes::Bytes; +use futures::TryStreamExt; +use kars_inference_router::{ + auth::WorkloadIdentityAuth, + config::{Config, ProviderEndpoint}, + deployment_health::DeploymentHealthRegistry, + failover::{forward_stream_with_failover, forward_with_failover}, + inference_policy_loader::{InferencePolicySnapshot, ModelPreference, ModelRef}, + proxy::UpstreamConfig, +}; +use serde_json::json; +use std::sync::Arc; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{body_partial_json, header, method, path}, +}; + +fn config(providers: &[(&str, &str, &str)]) -> Config { + let mut config = Config::from_env().unwrap(); + config.providers = providers + .iter() + .map(|(tag, endpoint, key)| { + ( + tag.to_string(), + ProviderEndpoint { + tag: tag.to_string(), + endpoint: endpoint.to_string(), + api_key: Some(key.to_string()), + }, + ) + }) + .collect(); + config +} + +fn base(endpoint: &str) -> UpstreamConfig { + let mut base = UpstreamConfig::azure( + endpoint.into(), + "default-model".into(), + "test-sandbox".into(), + ); + base.provider_api_key = Some("default-key".into()); + base +} + +fn policy() -> InferencePolicySnapshot { + InferencePolicySnapshot { + digest: "sha256:routing-test".into(), + model_preference: Some(ModelPreference { + primary: ModelRef { + provider: "first".into(), + deployment: "shared-model".into(), + }, + fallback: vec![ModelRef { + provider: "second".into(), + deployment: "shared-model".into(), + }], + }), + ..Default::default() + } +} + +#[tokio::test] +async fn no_model_preference_preserves_an_explicit_client_model() { + let server = MockServer::start().await; + Mock::given(body_partial_json(json!({"model":"client-selected-model"}))) + .respond_with(ResponseTemplate::new(200)) + .expect(1) + .mount(&server) + .await; + let (_, _, _, selected) = forward_with_failover( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::new(), + &Arc::new(DeploymentHealthRegistry::new()), + &base(&server.uri()), + &config(&[]), + &InferencePolicySnapshot::default(), + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from(r#"{"model":"client-selected-model"}"#), + ) + .await + .unwrap(); + assert_eq!(selected.deployment, "client-selected-model"); +} + +#[tokio::test] +async fn buffered_failover_retains_same_model_on_distinct_providers_and_uses_own_keys() { + let first = MockServer::start().await; + let second = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .and(header("authorization", "Bearer first-key")) + .and(body_partial_json(json!({"model":"shared-model"}))) + .respond_with(ResponseTemplate::new(503)) + .expect(1) + .mount(&first) + .await; + Mock::given(method("POST")) + .and(path("/chat/completions")) + .and(header("authorization", "Bearer second-key")) + .and(body_partial_json(json!({"model":"shared-model"}))) + .respond_with( + ResponseTemplate::new(200) + .set_body_json(json!({"choices":[{"message":{"content":"fallback"}}]})), + ) + .expect(1) + .mount(&second) + .await; + let config = config(&[ + ("first", &first.uri(), "first-key"), + ("second", &second.uri(), "second-key"), + ]); + let health = Arc::new(DeploymentHealthRegistry::new()); + let (status, _, body, selected) = forward_with_failover( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::new(), + &health, + &base(&first.uri()), + &config, + &policy(), + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from(r#"{"model":"caller-model","messages":[]}"#), + ) + .await + .unwrap(); + assert_eq!(status, StatusCode::OK); + assert_eq!(selected.endpoint, second.uri()); + assert_eq!(selected.deployment, "shared-model"); + assert!(String::from_utf8_lossy(&body).contains("fallback")); + assert_eq!( + health + .snapshot() + .iter() + .find(|s| s.deployment == "first::shared-model") + .unwrap() + .failure_streak, + 1 + ); +} + +#[tokio::test] +async fn same_endpoint_different_provider_credentials_are_not_reused() { + let server = MockServer::start().await; + Mock::given(header("authorization", "Bearer first-key")) + .respond_with(ResponseTemplate::new(503)) + .expect(1) + .mount(&server) + .await; + Mock::given(header("authorization", "Bearer second-key")) + .respond_with(ResponseTemplate::new(200).set_body_string("second")) + .expect(1) + .mount(&server) + .await; + let config = config(&[ + ("first", &server.uri(), "first-key"), + ("second", &server.uri(), "second-key"), + ]); + let (status, _, _, selected) = forward_with_failover( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::new(), + &Arc::new(DeploymentHealthRegistry::new()), + &base(&server.uri()), + &config, + &policy(), + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await + .unwrap(); + assert_eq!(status, StatusCode::OK); + assert_eq!(selected.provider_api_key.as_deref(), Some("second-key")); +} + +#[tokio::test] +async fn streaming_failover_retains_the_winning_provider_and_headers() { + let first = MockServer::start().await; + let second = MockServer::start().await; + Mock::given(header("authorization", "Bearer first-key")) + .respond_with(ResponseTemplate::new(429).set_body_string("rate limited")) + .expect(1) + .mount(&first) + .await; + Mock::given(header("authorization", "Bearer second-key")) + .and(body_partial_json(json!({"model":"shared-model"}))) + .respond_with(ResponseTemplate::new(200).set_body_raw( + "data: {\"choices\":[]}\n\ndata: [DONE]\n\n", + "text/event-stream", + )) + .expect(1) + .mount(&second) + .await; + let config = config(&[ + ("first", &first.uri(), "first-key"), + ("second", &second.uri(), "second-key"), + ]); + let (status, headers, stream, selected) = forward_stream_with_failover( + Arc::new(WorkloadIdentityAuth::new()), + None, + reqwest::Client::new(), + &Arc::new(DeploymentHealthRegistry::new()), + &base(&first.uri()), + &config, + &policy(), + "chat/completions", + HeaderMap::new(), + Bytes::from(r#"{"stream":true}"#), + ) + .await + .unwrap(); + assert_eq!(status, StatusCode::OK); + assert_eq!(headers["content-type"], "text/event-stream"); + assert_eq!(selected.endpoint, second.uri()); + let chunks: Vec<_> = stream.try_collect().await.unwrap(); + assert!( + chunks + .iter() + .any(|chunk| String::from_utf8_lossy(chunk).contains("[DONE]")) + ); +} + +#[tokio::test] +async fn accepted_stream_failure_is_never_replayed() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let first = format!("http://{}", listener.local_addr().unwrap()); + let server = tokio::spawn(async move { + let (mut socket, _) = listener.accept().await.unwrap(); + let mut request = [0_u8; 8192]; + let _ = socket.read(&mut request).await.unwrap(); + socket.write_all(b"HTTP/1.1 200 OK\r\nContent-Type: text/event-stream\r\nContent-Length: 10000\r\nConnection: close\r\n\r\ndata: partial\n\n").await.unwrap(); + socket.shutdown().await.unwrap(); + }); + let second = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&second) + .await; + let config = config(&[ + ("first", &first, "first-key"), + ("second", &second.uri(), "second-key"), + ]); + let (status, _, stream, selected) = forward_stream_with_failover( + Arc::new(WorkloadIdentityAuth::new()), + None, + reqwest::Client::new(), + &Arc::new(DeploymentHealthRegistry::new()), + &base(&first), + &config, + &policy(), + "chat/completions", + HeaderMap::new(), + Bytes::from(r#"{"stream":true}"#), + ) + .await + .unwrap(); + assert_eq!(status, StatusCode::OK); + assert_eq!(selected.endpoint, first); + assert!(stream.try_collect::>().await.is_err()); + server.await.unwrap(); +} + +#[tokio::test] +async fn auth_errors_are_not_failover_triggers() { + let first = MockServer::start().await; + let second = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(403).set_body_string("denied")) + .expect(1) + .mount(&first) + .await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&second) + .await; + let config = config(&[ + ("first", &first.uri(), "first-key"), + ("second", &second.uri(), "second-key"), + ]); + let (status, _, body, selected) = forward_with_failover( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::new(), + &Arc::new(DeploymentHealthRegistry::new()), + &base(&first.uri()), + &config, + &policy(), + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await + .unwrap(); + assert_eq!(status, StatusCode::FORBIDDEN); + assert_eq!(body, "denied"); + assert_eq!(selected.endpoint, first.uri()); +} + +#[tokio::test] +async fn exhausted_named_providers_restore_true_default_endpoint_model_and_key() { + let first = MockServer::start().await; + let second = MockServer::start().await; + let default = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(503)) + .expect(1) + .mount(&first) + .await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(503)) + .expect(1) + .mount(&second) + .await; + Mock::given(header("authorization", "Bearer default-key")) + .and(body_partial_json(json!({"model":"default-model"}))) + .respond_with(ResponseTemplate::new(200).set_body_string("default")) + .expect(1) + .mount(&default) + .await; + let config = config(&[ + ("first", &first.uri(), "first-key"), + ("second", &second.uri(), "second-key"), + ]); + let (status, _, _, selected) = forward_with_failover( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::new(), + &Arc::new(DeploymentHealthRegistry::new()), + &base(&default.uri()), + &config, + &policy(), + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await + .unwrap(); + assert_eq!(status, StatusCode::OK); + assert_eq!(selected.endpoint, default.uri()); + assert_eq!(selected.deployment, "default-model"); +} diff --git a/inference-router/tests/proxy_fake_upstream.rs b/inference-router/tests/proxy_fake_upstream.rs index 28c92f13..16c878a3 100644 --- a/inference-router/tests/proxy_fake_upstream.rs +++ b/inference-router/tests/proxy_fake_upstream.rs @@ -52,6 +52,18 @@ unsafe fn set_env(k: &str, v: &str) { unsafe { std::env::set_var(k, v) } } +fn azure_endpoint(base: &str) -> (String, reqwest::Client) { + let mut url = reqwest::Url::parse(base).unwrap(); + let address = url.socket_addrs(|| None).unwrap()[0]; + url.set_host(Some("test.openai.azure.com")).unwrap(); + let client = reqwest::Client::builder() + .no_proxy() + .resolve("test.openai.azure.com", address) + .build() + .unwrap(); + (url.to_string().trim_end_matches('/').into(), client) +} + #[tokio::test] async fn api_key_mode_proxies_chat_completion_with_filter_results() { let _g = ENV_LOCK.lock().unwrap(); @@ -68,15 +80,16 @@ async fn api_key_mode_proxies_chat_completion_with_filter_results() { let auth = WorkloadIdentityAuth::new(); assert!(auth.is_api_key_mode(), "expected API-key mode"); + let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { - endpoint: azure.base_url(), + endpoint, deployment: "gpt-4o".to_string(), sandbox_name: "test-sandbox".to_string(), provider: ProviderKind::AzureOpenAI, api_key: None, + provider_api_key: None, + authentication: Default::default(), }; - let client = reqwest::Client::new(); - let req_body = serde_json::json!({ "messages": [{"role": "user", "content": "hello fixture"}] }); @@ -148,14 +161,16 @@ async fn wi_mode_falls_back_to_imds_and_proxies_embeddings() { let auth = WorkloadIdentityAuth::new(); assert!(!auth.is_api_key_mode(), "expected WI mode"); + let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { - endpoint: azure.base_url(), + endpoint, deployment: "text-embedding-3-small".to_string(), sandbox_name: "test-sandbox-wi".to_string(), provider: ProviderKind::AzureOpenAI, api_key: None, + provider_api_key: None, + authentication: Default::default(), }; - let client = reqwest::Client::new(); let body = Bytes::from(r#"{"input":"hello"}"#.as_bytes().to_vec()); let (status, _headers, resp) = forward( @@ -219,14 +234,16 @@ async fn upstream_error_status_is_propagated() { .await; let auth = WorkloadIdentityAuth::new(); + let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { - endpoint: azure.base_url(), + endpoint, deployment: "gpt-4o".to_string(), sandbox_name: "test-sandbox-429".to_string(), provider: ProviderKind::AzureOpenAI, api_key: None, + provider_api_key: None, + authentication: Default::default(), }; - let client = reqwest::Client::new(); let (status, _headers, resp) = forward( &auth,