diff --git a/CHANGELOG.md b/CHANGELOG.md index a22c06412..b62f8e810 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,100 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] — `crd-well-oiled-machine` +### Multi-cloud LLM providers + pluggable guardrails (slice 1: Anthropic, Ollama, OpenAI Moderation) + +First slice of the "multi-cloud LLM providers + native guardrails" +roadmap theme. Credentials stay router-side (secret mount / env on the +sidecar) — the agent process keeps talking to the localhost proxy and +never sees a provider key. + +**`InferencePolicy` CRD (controller)** + +- `spec.provider` — new optional typed enum (`InferenceProvider`: + `azure-openai` | `anthropic` | `ollama` | `bedrock`), serialized with + the same kebab-case tags `ModelRef.provider` strings have always + documented, so existing YAML values stay valid. `bedrock` is + schema-accepted for forward-compat; the router answers 501 until the + Bedrock client lands (declared intent is never silently rerouted). +- `spec.guardrails[]` — new optional ordered guardrail pipeline + (`{provider: openai-moderation, applyTo: input|output|both}`, 1–8 + stages via CEL). Both new fields are mutually exclusive with + `spec.bundleRef` (CEL + reconciler defense-in-depth) and flow through + `compile_to_profile` into the compiled `inference-policy.json` + (`"provider"` / `"guardrails"` keys, null when absent). +- Helm CRD template regenerated (`crd-inferencepolicy.yaml`). +- Reconciler forwards router-only provider config to every sidecar when + present on the controller env: `ANTHROPIC_API_KEY`, + `ANTHROPIC_ENDPOINT`, `OLLAMA_ENDPOINT`, `OPENAI_MODERATION_API_KEY` + (falls back to `OPENAI_API_KEY`), `OPENAI_MODERATION_ENDPOINT`. + Endpoints (not secrets) join `CONFIG_HASH_INPUTS`. + +**Inference router — multi-provider upstreams** + +- New `provider` module: tag parsing + fail-closed resolution. + `spec.provider` is the sole routing selector; the pre-existing + `modelPreference.primary.provider` tag stays informational (it drove + no routing before this slice), so an unchanged CR that only set a + model preference keeps its Azure upstream. Unknown tags warn and fall + back to Azure; missing endpoint/credential → 503; unimplemented + `bedrock` → 501. Never a silent Azure fallback for a selected + non-Azure provider. +- `UpstreamConfig` carries `provider` + router-held `api_key`; + `proxy.rs` gains per-provider URL shapes (Anthropic: path verbatim; + Ollama: OpenAI-compat under `/v1/`) and auth schemes (Anthropic: + `x-api-key` + default `anthropic-version`; Ollama: no credential). + Agent-supplied `x-api-key` headers are stripped as before. +- `provider: anthropic` serves the Anthropic Messages surface natively + (`/anthropic/v1/messages`, `/v1/messages`): streaming, tool use and + multi-modal content pass through. OpenAI-shaped `/v1/chat/completions` + under an Anthropic policy returns an explicit 501 pointing at the + Messages surface (mirrors the GitHub-Models 501 precedent). +- `provider: ollama` serves `/v1/chat/completions` (buffered + SSE) + against `OLLAMA_ENDPOINT` with token metering and budget tracking. +- Routing + guardrails are wired on `/v1/chat/completions` and the + Anthropic Messages routes only. Every other inference-bearing route + (`/v1/completions`, `/v1/responses`, `/v1/embeddings`, image + generation, and the Foundry proxy families that serve model output — + `/agents*`, `/openai/responses*`, `/openai/conversations*`) doesn't + implement either and **fails closed** (501 on a non-Azure provider, + 403 when guardrails are declared) rather than silently bypass the + policy; a plain Azure policy uses them unchanged. The Foundry proxy + guard is default-deny: only an explicit exempt set of canonical + management/storage APIs (memory stores, files, vector stores, …) + bypasses it, and paths are canonicalized before routing — benign + empty segments (`//`, trailing `/`) are collapsed, while dot, + percent-encoded-dot, and encoded-slash/backslash segments are + rejected (400) so a request cannot normalize into a different + upstream route than the one classified. + +**Inference router — pluggable guardrail pipeline** + +- New `guardrails` module: `Guardrail` trait + OpenAI Moderation + backend (`POST {endpoint}/v1/moderations`, model + `omni-moderation-latest`, override via `OPENAI_MODERATION_MODEL`). +- Enforcement on the governed routes: request pre-flight (input + stages), buffered responses (including all Responses-API recovery + paths), and SSE streams via **hold-and-release** windows + (`GUARDRAIL_STREAM_SCAN_CHARS`, default 1000, clamped to the 16k scan + cap) — no assistant message text is delivered before a scan covers + it, including across partial-line chunk boundaries; flagged streams + are cut with a structured SSE error frame + `data: [DONE]`. Text over + 16k chars is scanned in successive windows, not truncated, so nothing + can be hidden past the cap. Tool-call arguments and provider + extended-"thinking" deltas are not yet scanned (roadmap). +- Fail-closed contract: declared-but-unbuildable stages reject the + request (503 `guardrail_misconfigured`); backend outages block (502 + `guardrail_unavailable`); unparseable bodies are scanned as raw text + rather than skipped. New `kars_guardrail_scans_total` metric + (provider / direction / outcome) and `x-kars-decision*` headers on + every block. + +Tests: compile/round-trip + enum wire-tag pins (controller), provider +resolution truth table, loader parsing, hold-and-release SSE guard +(clean / violation / split-event / scan-error), and wiremock +integration tests against fake Anthropic / Ollama / moderation +upstreams (`inference-router/tests/multi_provider_guardrails.rs`). + ## [0.1.26] — 2026-08-25 ### Security and dependency maintenance diff --git a/controller/src/config_hash.rs b/controller/src/config_hash.rs index 5a99ec205..b7ac69f1e 100644 --- a/controller/src/config_hash.rs +++ b/controller/src/config_hash.rs @@ -38,6 +38,11 @@ use std::sync::LazyLock; /// change and should be called out in the audit trail. pub const CONFIG_HASH_INPUTS: &[&str] = &[ "KARS_DISABLE_ENTRA_AUTH", + // Multi-provider endpoints (never secrets — matches the + // AZURE_OPENAI_API_KEY precedent). + "ANTHROPIC_ENDPOINT", + "OLLAMA_ENDPOINT", + "OPENAI_MODERATION_ENDPOINT", "AZURE_AUTHORITY_HOST", "AZURE_OPENAI_ENDPOINT", "AZURE_SUBSCRIPTION_ID", diff --git a/controller/src/crd_validations.rs b/controller/src/crd_validations.rs index bf121b5cd..7ffa74beb 100644 --- a/controller/src/crd_validations.rs +++ b/controller/src/crd_validations.rs @@ -277,8 +277,16 @@ pub fn inference_policy_validations() -> Vec { let severities = "['Safe','Low','Medium','High']"; vec![ ValidationRule { - rule: "!has(self.bundleRef) || (!has(self.tokenBudget) && !has(self.contentSafety) && !has(self.modelPreference) && !has(self.displayName))".into(), - message: Some("spec.bundleRef is mutually exclusive with spec.tokenBudget, spec.contentSafety, spec.modelPreference, and spec.displayName; the bundle carries those content fields".into()), + rule: "!has(self.bundleRef) || (!has(self.tokenBudget) && !has(self.contentSafety) && !has(self.modelPreference) && !has(self.provider) && !has(self.guardrails) && !has(self.displayName))".into(), + message: Some("spec.bundleRef is mutually exclusive with spec.tokenBudget, spec.contentSafety, spec.modelPreference, spec.provider, spec.guardrails, and spec.displayName; the bundle carries those content fields".into()), + reason: Some("FieldValueInvalid".into()), + ..ValidationRule::default() + }, + // An empty list is an authoring mistake (omit the field for + // no pipeline); the cap bounds per-request scan fan-out. + ValidationRule { + rule: "!has(self.guardrails) || (size(self.guardrails) >= 1 && size(self.guardrails) <= 8)".into(), + message: Some("spec.guardrails, when set, must contain 1-8 stages (omit the field for no pipeline)".into()), reason: Some("FieldValueInvalid".into()), ..ValidationRule::default() }, diff --git a/controller/src/inference_policy.rs b/controller/src/inference_policy.rs index cc0d35910..67a3ae866 100644 --- a/controller/src/inference_policy.rs +++ b/controller/src/inference_policy.rs @@ -54,6 +54,44 @@ use serde::{Deserialize, Serialize}; use crate::mcp_server::LocalObjectRef; +/// Inference provider selector. Explicit kebab-case renames (not +/// `rename_all`) keep the wire tags byte-identical to the free-form +/// `ModelRef.provider` strings this CRD has always documented, so a +/// CR can move to the typed field without a migration. `bedrock` is +/// schema-accepted for forward-compat; the router returns 501 until +/// its client lands. +#[derive(Serialize, Deserialize, Clone, Copy, Debug, Default, PartialEq, Eq, JsonSchema)] +pub enum InferenceProvider { + /// Azure OpenAI / Azure AI Foundry (default — Phase 1 substrate). + #[default] + #[serde(rename = "azure-openai")] + AzureOpenAI, + /// Anthropic Messages API (`api.anthropic.com` or a compatible + /// gateway). + #[serde(rename = "anthropic")] + Anthropic, + /// OpenAI-compatible Ollama server (in-cluster or local). + #[serde(rename = "ollama")] + Ollama, + /// AWS Bedrock — schema-level forward-compat; router support is a + /// follow-up slice. + #[serde(rename = "bedrock")] + AWSBedrock, +} + +impl InferenceProvider { + /// The kebab-case wire tag serde emits. + #[must_use] + pub fn as_tag(&self) -> &'static str { + match self { + Self::AzureOpenAI => "azure-openai", + Self::Anthropic => "anthropic", + Self::Ollama => "ollama", + Self::AWSBedrock => "bedrock", + } + } +} + /// `InferencePolicy.spec` — declares per-sandbox inference-time /// guardrails: token budgets, Content Safety severity floors, model /// preference + fallback chain. @@ -102,6 +140,20 @@ pub struct InferencePolicySpec { /// [`Self::bundle_ref`]. pub model_preference: Option, + /// Inference provider for the call sites this policy governs, and + /// the sole routing selector. Absent ⇒ the env-configured Azure + /// OpenAI / Foundry upstream. A non-Azure provider swaps the base + /// URL and auth scheme; credentials stay in the router sidecar, + /// never in the agent. Mutually exclusive with [`Self::bundle_ref`]. + pub provider: Option, + + /// Ordered guardrail pipeline stages the router runs around each + /// call (input pre-flight + response, buffered and streaming), in + /// declaration order — first flag blocks. Absent ⇒ only the + /// Phase 1 substrate (Foundry annotations + `contentSafety`). + /// Mutually exclusive with [`Self::bundle_ref`]. + pub guardrails: Option>, + /// Optional human-readable label. Mutually exclusive with /// [`Self::bundle_ref`] — when `bundleRef` is set, the label /// comes from the signed bundle. @@ -228,6 +280,45 @@ pub struct ModelRef { pub deployment: String, } +/// A single router-side guardrail stage. A stage whose backend isn't +/// configured (e.g. missing moderation key) fails the request closed, +/// never skips. +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, JsonSchema)] +#[serde(rename_all = "camelCase")] +pub struct GuardrailStage { + /// Guardrail backend. + pub provider: GuardrailProvider, + + /// Which direction(s) this stage scans. Absent ⇒ `both`. + pub apply_to: Option, +} + +/// Guardrail backend. Bedrock Guardrails / Model Armor are roadmap +/// follow-ups that extend this enum. +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, JsonSchema)] +pub enum GuardrailProvider { + /// OpenAI Moderation API; router-side key via + /// `OPENAI_MODERATION_API_KEY` (falls back to `OPENAI_API_KEY`). + #[serde(rename = "openai-moderation")] + OpenAIModeration, +} + +/// Scan direction for a [`GuardrailStage`]. +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, Default, JsonSchema)] +pub enum GuardrailApplyTo { + /// Scan only the request (prompt) text. + #[serde(rename = "input")] + Input, + /// Scan only the response (completion) text — buffered and + /// streaming. + #[serde(rename = "output")] + Output, + /// Scan both directions (default). + #[default] + #[serde(rename = "both")] + Both, +} + /// Status of an `InferencePolicy` reconcile. #[derive(Debug, Serialize, Deserialize, Default, Clone, JsonSchema)] #[serde(rename_all = "camelCase")] diff --git a/controller/src/inference_policy_compile.rs b/controller/src/inference_policy_compile.rs index 548f1def1..ae6c630c0 100644 --- a/controller/src/inference_policy_compile.rs +++ b/controller/src/inference_policy_compile.rs @@ -72,6 +72,8 @@ use crate::inference_policy::InferencePolicySpec; /// "tokenBudget": { "perRequestTokens": ..., "dailyTokens": ..., "monthlyTokens": ... } | null, /// "contentSafety": { "hate": ..., "selfHarm": ..., "sexual": ..., "violence": ..., "requirePromptShields": ... } | null, /// "modelPreference": { "primary": {provider, deployment}, "fallback": [...] } | null, +/// "provider": "azure-openai" | "anthropic" | "ollama" | "bedrock" | null, +/// "guardrails": [ { "provider": "openai-moderation", "applyTo": "input"|"output"|"both"|null }, ... ] | null, /// "displayName": "..." | null /// } /// ``` @@ -118,11 +120,29 @@ pub fn compile_to_profile(spec: &InferencePolicySpec) -> Value { }) }); + // Provider + guardrails emit the same kebab-case wire tags the + // CRD serde uses, so the router parses one vocabulary. + let provider = spec.provider.as_ref().map(|p| json!(p.as_tag())); + + let guardrails = spec.guardrails.as_ref().map(|stages| { + json!( + stages + .iter() + .map(|g| json!({ + "provider": g.provider, + "applyTo": g.apply_to, + })) + .collect::>() + ) + }); + json!({ "appliesTo": applies_to, "tokenBudget": token_budget, "contentSafety": content_safety, "modelPreference": model_preference, + "provider": provider, + "guardrails": guardrails, "displayName": spec.display_name, }) } @@ -195,7 +215,8 @@ pub fn inference_policy_digest(body: &[u8]) -> String { mod tests { use super::*; use crate::inference_policy::{ - ContentSafetyFloor, InferenceAppliesTo, InferencePolicySpec, ModelPreference, ModelRef, + ContentSafetyFloor, GuardrailApplyTo, GuardrailProvider, GuardrailStage, + InferenceAppliesTo, InferencePolicySpec, InferenceProvider, ModelPreference, ModelRef, TokenBudget, }; @@ -230,6 +251,17 @@ mod tests { deployment: "claude-3-5-sonnet".into(), }], }), + provider: Some(InferenceProvider::Anthropic), + guardrails: Some(vec![ + GuardrailStage { + provider: GuardrailProvider::OpenAIModeration, + apply_to: Some(GuardrailApplyTo::Output), + }, + GuardrailStage { + provider: GuardrailProvider::OpenAIModeration, + apply_to: None, + }, + ]), display_name: Some("Prod chat policy".into()), bundle_ref: None, } @@ -243,6 +275,8 @@ mod tests { assert!(profile.get("tokenBudget").unwrap().is_null()); assert!(profile.get("contentSafety").unwrap().is_null()); assert!(profile.get("modelPreference").unwrap().is_null()); + assert!(profile.get("provider").unwrap().is_null()); + assert!(profile.get("guardrails").unwrap().is_null()); assert!(profile.get("appliesTo").unwrap().is_object()); } @@ -271,9 +305,60 @@ mod tests { profile["modelPreference"]["fallback"][0]["provider"], "anthropic" ); + assert_eq!(profile["provider"], "anthropic"); + let stages = profile["guardrails"].as_array().unwrap(); + assert_eq!(stages.len(), 2); + assert_eq!(stages[0]["provider"], "openai-moderation"); + assert_eq!(stages[0]["applyTo"], "output"); + assert_eq!(stages[1]["provider"], "openai-moderation"); + assert!(stages[1]["applyTo"].is_null()); assert_eq!(profile["displayName"], "Prod chat policy"); } + #[test] + fn provider_enum_serializes_to_kebab_case_wire_tags() { + // Wire-contract pin: the typed `spec.provider` must emit the + // same kebab-case tags the free-form `ModelRef.provider` + // strings have always documented, so existing YAML values + // stay valid when operators migrate to the typed field. + for (variant, tag) in [ + (InferenceProvider::AzureOpenAI, "azure-openai"), + (InferenceProvider::Anthropic, "anthropic"), + (InferenceProvider::Ollama, "ollama"), + (InferenceProvider::AWSBedrock, "bedrock"), + ] { + assert_eq!(serde_json::to_value(variant).unwrap(), tag); + assert_eq!(variant.as_tag(), tag); + let parsed: InferenceProvider = serde_json::from_value(serde_json::json!(tag)).unwrap(); + assert_eq!(parsed, variant); + } + } + + #[test] + fn version_hash_changes_when_provider_changes() { + let mut a = full_spec(); + let b = full_spec(); + a.provider = Some(InferenceProvider::Ollama); + assert_ne!( + version_hash(&compile_to_profile(&a)), + version_hash(&compile_to_profile(&b)) + ); + } + + #[test] + fn version_hash_changes_when_guardrails_change() { + let mut a = full_spec(); + let b = full_spec(); + a.guardrails = Some(vec![GuardrailStage { + provider: GuardrailProvider::OpenAIModeration, + apply_to: Some(GuardrailApplyTo::Input), + }]); + assert_ne!( + version_hash(&compile_to_profile(&a)), + version_hash(&compile_to_profile(&b)) + ); + } + #[test] fn compile_is_deterministic() { let spec = full_spec(); diff --git a/controller/src/inference_policy_reconciler.rs b/controller/src/inference_policy_reconciler.rs index 30e64feda..d7a9eab67 100644 --- a/controller/src/inference_policy_reconciler.rs +++ b/controller/src/inference_policy_reconciler.rs @@ -375,6 +375,8 @@ async fn resolve_inference_source( let inline_any = spec.token_budget.is_some() || spec.content_safety.is_some() || spec.model_preference.is_some() + || spec.provider.is_some() + || spec.guardrails.is_some() || spec.display_name.is_some(); let bundle_set = spec.bundle_ref.is_some(); @@ -389,7 +391,8 @@ async fn resolve_inference_source( Some(( "InvalidSpec", "spec.bundleRef is mutually exclusive with spec.tokenBudget, \ - spec.contentSafety, spec.modelPreference, and spec.displayName" + spec.contentSafety, spec.modelPreference, spec.provider, \ + spec.guardrails, and spec.displayName" .into(), )), ); @@ -520,6 +523,11 @@ fn merge_bundle_with_selector( token_budget, content_safety, model_preference, + // The signed-bundle canonical format doesn't carry + // provider/guardrails yet; bundle-sourced policies keep the + // router defaults on these axes. + provider: None, + guardrails: None, display_name: verified.display_name.clone(), bundle_ref: None, } diff --git a/controller/src/reconciler/mod.rs b/controller/src/reconciler/mod.rs index 670225ba7..631050648 100644 --- a/controller/src/reconciler/mod.rs +++ b/controller/src/reconciler/mod.rs @@ -30,7 +30,7 @@ use serde_json::json; use std::sync::Arc; use tokio::time::Duration; -use crate::crd::{KarsSandbox, SandboxConfig}; +use crate::crd::KarsSandbox; use crate::fedcred::{FedCredConfig, FedCredManager}; pub(crate) mod byo_contract; @@ -41,220 +41,10 @@ pub(crate) mod trustgraph_mount; use mcp_egress::mcp_egress_rule; -/// Build pod security context, conditionally including SELinux options and -/// choosing between RuntimeDefault and Localhost seccomp profiles. -/// For Kata (confidential), we use RuntimeDefault since the VM provides isolation. -pub(crate) fn build_pod_security_context(cfg: &SandboxConfig) -> serde_json::Value { - // Standard and Confidential use RuntimeDefault seccomp: - // standard — basic container isolation, kernel-default syscall filter - // confidential — Kata VM boundary is the isolation layer - // Enhanced uses custom Localhost seccomp (kars-strict) for strict syscall allowlist - let seccomp = if cfg.isolation == "confidential" - || cfg.isolation == "standard" - || cfg.seccomp_profile == "RuntimeDefault" - || cfg.seccomp_profile.is_empty() - { - json!({ "type": "RuntimeDefault" }) - } else { - json!({ - "type": "Localhost", - "localhostProfile": format!("profiles/{}.json", cfg.seccomp_profile) - }) - }; - - let mut ctx = json!({ - "runAsNonRoot": cfg.run_as_non_root, - "runAsUser": 1000, - "runAsGroup": 1000, - "fsGroup": 1000, - "seccompProfile": seccomp - }); - - // Only set seLinuxOptions if a non-empty context is specified - if !cfg.selinux_context.is_empty() { - ctx.as_object_mut().unwrap().insert( - "seLinuxOptions".into(), - json!({ "type": cfg.selinux_context }), - ); - } - - ctx -} - -/// Returns (runtimeClassName, nodeSelector) based on the isolation level. -/// standard → runc on clawpool, no custom seccomp -/// enhanced → runc on clawpool + Localhost seccomp (kars-strict) -/// confidential → Kata VM isolation on katapool -pub(crate) fn isolation_scheduling(isolation: &str) -> (Option<&'static str>, &'static str) { - match isolation { - "confidential" => (Some("kata-vm-isolation"), "sandbox-kata"), - _ => (None, "sandbox"), // standard + enhanced both on clawpool - } -} - -/// Build the egress-guard init-container command. -/// -/// Standard sandboxes (every kind except SRE) get the full lockdown: -/// UID 1000 → loopback + DNS allowed, everything else dropped, with -/// :80/:443 NAT-redirected to the inference-router on :8444 for L7 -/// policy + audit. -/// -/// SRE-mode sandboxes (labelled `kars.azure.com/role=sre`) get ONE -/// extra rule inserted into the OUTPUT NAT chain BEFORE the generic -/// REDIRECT: apiserver-bound traffic (KUBERNETES_SERVICE_HOST : -/// KUBERNETES_SERVICE_PORT_HTTPS, both kubelet-auto-injected envs) -/// is RETURNed — i.e. NOT NAT'd to :8444 — so the SRE plugin's K8s -/// API client (sre_kube.py) can hit the apiserver directly with its -/// projected SA token. -/// -/// The K8s audit log is the audit surface for these apiserver calls -/// (the router's L7 audit doesn't capture them, but K8s audit is -/// stronger — every call carries the SA identity and the verb). -/// -/// Privilege-containment design: this capability is uniquely held by -/// the SRE sandbox per the proposal §7.8. Future Slice 3 will add -/// ValidatingAdmissionPolicies to gate WHO can apply the -/// `role=sre` label (only chart-installer SAs; see §7.8.10 design). -pub(crate) fn build_egress_guard_command(is_sre_sandbox: bool) -> String { - let mut cmd = String::with_capacity(1024); - // Filter chain (OUTPUT): UID 1000 → allow loopback + DNS + - // established, then DROP. Same for every sandbox kind. - cmd.push_str("iptables -A OUTPUT -m owner --uid-owner 1000 -o lo -j ACCEPT && "); - cmd.push_str("iptables -A OUTPUT -m owner --uid-owner 1000 -p udp --dport 53 -j ACCEPT && "); - cmd.push_str("iptables -A OUTPUT -m owner --uid-owner 1000 -p tcp --dport 53 -j ACCEPT && "); - cmd.push_str( - "iptables -A OUTPUT -m owner --uid-owner 1000 -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT && " - ); - - // SRE-mode-only: filter-chain ACCEPT for apiserver-bound traffic. - // The filter chain runs AFTER the NAT chain — the NAT-bypass RETURN - // below just decides "don't redirect", but the filter chain's DROP - // (next rule) would still kill the packet. We have to ACCEPT it - // here BEFORE the catch-all DROP. - if is_sre_sandbox { - cmd.push_str( - "iptables -A OUTPUT -m owner --uid-owner 1000 \ - -d \"${KUBERNETES_SERVICE_HOST}\" \ - -p tcp --dport \"${KUBERNETES_SERVICE_PORT_HTTPS:-443}\" \ - -j ACCEPT && ", - ); - } - - cmd.push_str("iptables -A OUTPUT -m owner --uid-owner 1000 -j DROP && "); - - // SRE-mode-only: NAT-chain apiserver bypass. Inserted BEFORE the - // generic :443 REDIRECT so apiserver traffic short-circuits to the - // real upstream rather than the router. KUBERNETES_SERVICE_HOST - // and KUBERNETES_SERVICE_PORT_HTTPS are auto-injected by the - // kubelet on every container (including init containers). - if is_sre_sandbox { - cmd.push_str( - "iptables -t nat -A OUTPUT -m owner --uid-owner 1000 \ - -d \"${KUBERNETES_SERVICE_HOST}\" \ - -p tcp --dport \"${KUBERNETES_SERVICE_PORT_HTTPS:-443}\" \ - -j RETURN && ", - ); - } - - // NAT chain (OUTPUT): :80/:443 → REDIRECT to :8444 (transparent - // proxy in the inference-router sidecar). Same for every sandbox. - cmd.push_str( - "iptables -t nat -A OUTPUT -m owner --uid-owner 1000 ! -o lo -p tcp --dport 80 -j REDIRECT --to-port 8444 && " - ); - cmd.push_str( - "iptables -t nat -A OUTPUT -m owner --uid-owner 1000 ! -o lo -p tcp --dport 443 -j REDIRECT --to-port 8444 && " - ); - - if is_sre_sandbox { - cmd.push_str( - "echo 'egress-guard: UID 1000 → transparent proxy on :8444 + apiserver bypass (SRE mode)'" - ); - } else { - cmd.push_str( - "echo 'egress-guard: UID 1000 → transparent proxy on :8444 (learn + enforce)'", - ); - } - - cmd -} - -#[cfg(test)] -#[allow(clippy::module_inception)] -mod egress_guard_tests { - use super::build_egress_guard_command; - - #[test] - fn standard_sandbox_has_no_apiserver_bypass() { - let cmd = build_egress_guard_command(false); - assert!(!cmd.contains("KUBERNETES_SERVICE_HOST")); - assert!(cmd.contains("REDIRECT --to-port 8444")); - assert!(cmd.contains("(learn + enforce)")); - assert!(!cmd.contains("apiserver bypass")); - } - - #[test] - fn sre_sandbox_inserts_apiserver_bypass_before_redirect() { - let cmd = build_egress_guard_command(true); - // The bypass MUST come before the :443 REDIRECT — otherwise - // the REDIRECT wins (iptables -A appends; rules evaluate in - // order) and the bypass is dead code. - let bypass_pos = cmd - .find("-t nat -A OUTPUT -m owner --uid-owner 1000 -d \"${KUBERNETES_SERVICE_HOST}\"") - .or_else(|| cmd.find("-t nat -A OUTPUT -m owner --uid-owner 1000 \t\t\t -d \"${KUBERNETES_SERVICE_HOST}\"")) - .or_else(|| { - // Match the NAT-chain bypass specifically (not the filter ACCEPT) - cmd.match_indices("-t nat -A OUTPUT") - .find(|(i, _)| cmd[*i..].contains("KUBERNETES_SERVICE_HOST")) - .map(|(i, _)| i) - }) - .expect("NAT-chain bypass rule missing"); - let redirect_pos = cmd - .find("--dport 443 -j REDIRECT") - .expect("redirect rule missing"); - assert!( - bypass_pos < redirect_pos, - "NAT bypass at {bypass_pos} must precede redirect at {redirect_pos}" - ); - assert!(cmd.contains("apiserver bypass (SRE mode)")); - - // ALSO check the filter-chain ACCEPT exists BEFORE the DROP — this - // was the bug we hit live: NAT bypass alone wasn't enough because - // the filter chain's DROP for UID 1000 killed the packet anyway. - let filter_accept = cmd - .find( - "-A OUTPUT -m owner --uid-owner 1000 -d \"${KUBERNETES_SERVICE_HOST}\"", - ) - .or_else(|| { - cmd.match_indices("-A OUTPUT -m owner --uid-owner 1000") - .find(|(i, _)| { - let tail = &cmd[*i..*i + 200.min(cmd.len() - *i)]; - tail.contains("KUBERNETES_SERVICE_HOST") && tail.contains("-j ACCEPT") - }) - .map(|(i, _)| i) - }) - .expect("filter-chain ACCEPT for apiserver missing"); - let filter_drop = cmd - .find("-A OUTPUT -m owner --uid-owner 1000 -j DROP") - .expect("filter DROP rule missing"); - assert!( - filter_accept < filter_drop, - "filter ACCEPT at {filter_accept} must precede DROP at {filter_drop}" - ); - } - - #[test] - fn both_modes_keep_the_filter_chain_lockdown() { - for is_sre in [false, true] { - let cmd = build_egress_guard_command(is_sre); - // The filter-chain DROP rule is the actual lockdown — must - // never be removed by either mode. - assert!( - cmd.contains("-A OUTPUT -m owner --uid-owner 1000 -j DROP"), - "filter-chain DROP missing for is_sre={is_sre}" - ); - } - } -} +mod pod_spec; +pub(crate) use pod_spec::{ + build_egress_guard_command, build_pod_security_context, isolation_scheduling, +}; #[derive(Debug, thiserror::Error)] enum ReconcileError { #[error("Kubernetes API error: {0}")] @@ -311,6 +101,24 @@ struct Context { dev_openai_api_key: String, dev_provider: String, dev_copilot_github_token: String, + /// Multi-provider inference credentials/endpoints (roadmap: + /// multi-cloud LLM providers + native guardrails). Forwarded to + /// every router sidecar when present so `InferencePolicy` + /// `spec.provider` / `spec.guardrails` can resolve. All sourced + /// from the controller's own env at startup (Helm + /// `controller.extraEnv`, typically referencing a Secret) — the + /// agent container NEVER receives these. + /// - `anthropic_api_key` / `anthropic_endpoint`: Anthropic + /// Messages API (`provider: anthropic`). + /// - `ollama_endpoint`: OpenAI-compatible Ollama server + /// (`provider: ollama`); no credential. + /// - `openai_moderation_api_key` / `openai_moderation_endpoint`: + /// OpenAI Moderation guardrail stage backend. + anthropic_api_key: String, + anthropic_endpoint: String, + ollama_endpoint: String, + openai_moderation_api_key: String, + openai_moderation_endpoint: String, /// `KARS_DEV_PROFILE=true` (set only in `kars dev`) — triggers /// relaxed sub-agent CRD defaults in the router spawn helper. dev_profile: bool, @@ -1930,6 +1738,22 @@ async fn reconcile(sandbox: Arc, ctx: Arc) -> Result Result<()> { dev_openai_api_key, dev_provider, dev_copilot_github_token, + anthropic_api_key: std::env::var("ANTHROPIC_API_KEY").unwrap_or_default(), + anthropic_endpoint: std::env::var("ANTHROPIC_ENDPOINT").unwrap_or_default(), + ollama_endpoint: std::env::var("OLLAMA_ENDPOINT").unwrap_or_default(), + openai_moderation_api_key: std::env::var("OPENAI_MODERATION_API_KEY") + .or_else(|_| std::env::var("OPENAI_API_KEY")) + .unwrap_or_default(), + openai_moderation_endpoint: std::env::var("OPENAI_MODERATION_ENDPOINT").unwrap_or_default(), dev_profile, cluster_name: std::env::var("CLUSTER_NAME") .ok() diff --git a/controller/src/reconciler/pod_spec.rs b/controller/src/reconciler/pod_spec.rs new file mode 100644 index 000000000..811bf1c39 --- /dev/null +++ b/controller/src/reconciler/pod_spec.rs @@ -0,0 +1,226 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Pure pod-spec helpers extracted from the reconciler: pod +//! security context, isolation-based scheduling, and the +//! egress-guard init-container command. Kept here to bound +//! `reconciler/mod.rs` size; all functions are side-effect-free. + +use serde_json::json; + +use crate::crd::SandboxConfig; + +/// Build pod security context, conditionally including SELinux options and +/// choosing between RuntimeDefault and Localhost seccomp profiles. +/// For Kata (confidential), we use RuntimeDefault since the VM provides isolation. +pub(crate) fn build_pod_security_context(cfg: &SandboxConfig) -> serde_json::Value { + // Standard and Confidential use RuntimeDefault seccomp: + // standard — basic container isolation, kernel-default syscall filter + // confidential — Kata VM boundary is the isolation layer + // Enhanced uses custom Localhost seccomp (kars-strict) for strict syscall allowlist + let seccomp = if cfg.isolation == "confidential" + || cfg.isolation == "standard" + || cfg.seccomp_profile == "RuntimeDefault" + || cfg.seccomp_profile.is_empty() + { + json!({ "type": "RuntimeDefault" }) + } else { + json!({ + "type": "Localhost", + "localhostProfile": format!("profiles/{}.json", cfg.seccomp_profile) + }) + }; + + let mut ctx = json!({ + "runAsNonRoot": cfg.run_as_non_root, + "runAsUser": 1000, + "runAsGroup": 1000, + "fsGroup": 1000, + "seccompProfile": seccomp + }); + + // Only set seLinuxOptions if a non-empty context is specified + if !cfg.selinux_context.is_empty() { + ctx.as_object_mut().unwrap().insert( + "seLinuxOptions".into(), + json!({ "type": cfg.selinux_context }), + ); + } + + ctx +} + +/// Returns (runtimeClassName, nodeSelector) based on the isolation level. +/// standard → runc on clawpool, no custom seccomp +/// enhanced → runc on clawpool + Localhost seccomp (kars-strict) +/// confidential → Kata VM isolation on katapool +pub(crate) fn isolation_scheduling(isolation: &str) -> (Option<&'static str>, &'static str) { + match isolation { + "confidential" => (Some("kata-vm-isolation"), "sandbox-kata"), + _ => (None, "sandbox"), // standard + enhanced both on clawpool + } +} + +/// Build the egress-guard init-container command. +/// +/// Standard sandboxes (every kind except SRE) get the full lockdown: +/// UID 1000 → loopback + DNS allowed, everything else dropped, with +/// :80/:443 NAT-redirected to the inference-router on :8444 for L7 +/// policy + audit. +/// +/// SRE-mode sandboxes (labelled `kars.azure.com/role=sre`) get ONE +/// extra rule inserted into the OUTPUT NAT chain BEFORE the generic +/// REDIRECT: apiserver-bound traffic (KUBERNETES_SERVICE_HOST : +/// KUBERNETES_SERVICE_PORT_HTTPS, both kubelet-auto-injected envs) +/// is RETURNed — i.e. NOT NAT'd to :8444 — so the SRE plugin's K8s +/// API client (sre_kube.py) can hit the apiserver directly with its +/// projected SA token. +/// +/// The K8s audit log is the audit surface for these apiserver calls +/// (the router's L7 audit doesn't capture them, but K8s audit is +/// stronger — every call carries the SA identity and the verb). +/// +/// Privilege-containment design: this capability is uniquely held by +/// the SRE sandbox per the proposal §7.8. Future Slice 3 will add +/// ValidatingAdmissionPolicies to gate WHO can apply the +/// `role=sre` label (only chart-installer SAs; see §7.8.10 design). +pub(crate) fn build_egress_guard_command(is_sre_sandbox: bool) -> String { + let mut cmd = String::with_capacity(1024); + // Filter chain (OUTPUT): UID 1000 → allow loopback + DNS + + // established, then DROP. Same for every sandbox kind. + cmd.push_str("iptables -A OUTPUT -m owner --uid-owner 1000 -o lo -j ACCEPT && "); + cmd.push_str("iptables -A OUTPUT -m owner --uid-owner 1000 -p udp --dport 53 -j ACCEPT && "); + cmd.push_str("iptables -A OUTPUT -m owner --uid-owner 1000 -p tcp --dport 53 -j ACCEPT && "); + cmd.push_str( + "iptables -A OUTPUT -m owner --uid-owner 1000 -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT && " + ); + + // SRE-mode-only: filter-chain ACCEPT for apiserver-bound traffic. + // The filter chain runs AFTER the NAT chain — the NAT-bypass RETURN + // below just decides "don't redirect", but the filter chain's DROP + // (next rule) would still kill the packet. We have to ACCEPT it + // here BEFORE the catch-all DROP. + if is_sre_sandbox { + cmd.push_str( + "iptables -A OUTPUT -m owner --uid-owner 1000 \ + -d \"${KUBERNETES_SERVICE_HOST}\" \ + -p tcp --dport \"${KUBERNETES_SERVICE_PORT_HTTPS:-443}\" \ + -j ACCEPT && ", + ); + } + + cmd.push_str("iptables -A OUTPUT -m owner --uid-owner 1000 -j DROP && "); + + // SRE-mode-only: NAT-chain apiserver bypass. Inserted BEFORE the + // generic :443 REDIRECT so apiserver traffic short-circuits to the + // real upstream rather than the router. KUBERNETES_SERVICE_HOST + // and KUBERNETES_SERVICE_PORT_HTTPS are auto-injected by the + // kubelet on every container (including init containers). + if is_sre_sandbox { + cmd.push_str( + "iptables -t nat -A OUTPUT -m owner --uid-owner 1000 \ + -d \"${KUBERNETES_SERVICE_HOST}\" \ + -p tcp --dport \"${KUBERNETES_SERVICE_PORT_HTTPS:-443}\" \ + -j RETURN && ", + ); + } + + // NAT chain (OUTPUT): :80/:443 → REDIRECT to :8444 (transparent + // proxy in the inference-router sidecar). Same for every sandbox. + cmd.push_str( + "iptables -t nat -A OUTPUT -m owner --uid-owner 1000 ! -o lo -p tcp --dport 80 -j REDIRECT --to-port 8444 && " + ); + cmd.push_str( + "iptables -t nat -A OUTPUT -m owner --uid-owner 1000 ! -o lo -p tcp --dport 443 -j REDIRECT --to-port 8444 && " + ); + + if is_sre_sandbox { + cmd.push_str( + "echo 'egress-guard: UID 1000 → transparent proxy on :8444 + apiserver bypass (SRE mode)'" + ); + } else { + cmd.push_str( + "echo 'egress-guard: UID 1000 → transparent proxy on :8444 (learn + enforce)'", + ); + } + + cmd +} + +#[cfg(test)] +#[allow(clippy::module_inception)] +mod egress_guard_tests { + use super::build_egress_guard_command; + + #[test] + fn standard_sandbox_has_no_apiserver_bypass() { + let cmd = build_egress_guard_command(false); + assert!(!cmd.contains("KUBERNETES_SERVICE_HOST")); + assert!(cmd.contains("REDIRECT --to-port 8444")); + assert!(cmd.contains("(learn + enforce)")); + assert!(!cmd.contains("apiserver bypass")); + } + + #[test] + fn sre_sandbox_inserts_apiserver_bypass_before_redirect() { + let cmd = build_egress_guard_command(true); + // The bypass MUST come before the :443 REDIRECT — otherwise + // the REDIRECT wins (iptables -A appends; rules evaluate in + // order) and the bypass is dead code. + let bypass_pos = cmd + .find("-t nat -A OUTPUT -m owner --uid-owner 1000 -d \"${KUBERNETES_SERVICE_HOST}\"") + .or_else(|| cmd.find("-t nat -A OUTPUT -m owner --uid-owner 1000 \t\t\t -d \"${KUBERNETES_SERVICE_HOST}\"")) + .or_else(|| { + // Match the NAT-chain bypass specifically (not the filter ACCEPT) + cmd.match_indices("-t nat -A OUTPUT") + .find(|(i, _)| cmd[*i..].contains("KUBERNETES_SERVICE_HOST")) + .map(|(i, _)| i) + }) + .expect("NAT-chain bypass rule missing"); + let redirect_pos = cmd + .find("--dport 443 -j REDIRECT") + .expect("redirect rule missing"); + assert!( + bypass_pos < redirect_pos, + "NAT bypass at {bypass_pos} must precede redirect at {redirect_pos}" + ); + assert!(cmd.contains("apiserver bypass (SRE mode)")); + + // ALSO check the filter-chain ACCEPT exists BEFORE the DROP — this + // was the bug we hit live: NAT bypass alone wasn't enough because + // the filter chain's DROP for UID 1000 killed the packet anyway. + let filter_accept = cmd + .find( + "-A OUTPUT -m owner --uid-owner 1000 -d \"${KUBERNETES_SERVICE_HOST}\"", + ) + .or_else(|| { + cmd.match_indices("-A OUTPUT -m owner --uid-owner 1000") + .find(|(i, _)| { + let tail = &cmd[*i..*i + 200.min(cmd.len() - *i)]; + tail.contains("KUBERNETES_SERVICE_HOST") && tail.contains("-j ACCEPT") + }) + .map(|(i, _)| i) + }) + .expect("filter-chain ACCEPT for apiserver missing"); + let filter_drop = cmd + .find("-A OUTPUT -m owner --uid-owner 1000 -j DROP") + .expect("filter DROP rule missing"); + assert!( + filter_accept < filter_drop, + "filter ACCEPT at {filter_accept} must precede DROP at {filter_drop}" + ); + } + + #[test] + fn both_modes_keep_the_filter_chain_lockdown() { + for is_sre in [false, true] { + let cmd = build_egress_guard_command(is_sre); + // The filter-chain DROP rule is the actual lockdown — must + // never be removed by either mode. + assert!( + cmd.contains("-A OUTPUT -m owner --uid-owner 1000 -j DROP"), + "filter-chain DROP missing for is_sre={is_sre}" + ); + } + } +} diff --git a/deploy/helm/kars/templates/crd-inferencepolicy.yaml b/deploy/helm/kars/templates/crd-inferencepolicy.yaml index fceb79c29..202ffc618 100644 --- a/deploy/helm/kars/templates/crd-inferencepolicy.yaml +++ b/deploy/helm/kars/templates/crd-inferencepolicy.yaml @@ -181,6 +181,37 @@ spec: comes from the signed bundle. nullable: true type: string + guardrails: + description: |- + Ordered guardrail pipeline stages the router runs around each + call (input pre-flight + response, buffered and streaming), in + declaration order — first flag blocks. Absent ⇒ only the + Phase 1 substrate (Foundry annotations + `contentSafety`). + Mutually exclusive with [`Self::bundle_ref`]. + items: + description: |- + A single router-side guardrail stage. A stage whose backend isn't + configured (e.g. missing moderation key) fails the request closed, + never skips. + properties: + applyTo: + description: Scan direction for a [`GuardrailStage`]. + enum: + - input + - output + - both + nullable: true + type: string + provider: + description: Guardrail backend. + enum: + - openai-moderation + type: string + required: + - provider + type: object + nullable: true + type: array modelPreference: description: |- Model preference + fallback chain. Optional. **Not** a router: @@ -226,6 +257,21 @@ spec: required: - primary type: object + provider: + description: |- + Inference provider selector. Explicit kebab-case renames (not + `rename_all`) keep the wire tags byte-identical to the free-form + `ModelRef.provider` strings this CRD has always documented, so a + CR can move to the typed field without a migration. `bedrock` is + schema-accepted for forward-compat; the router returns 501 until + its client lands. + enum: + - azure-openai + - anthropic + - ollama + - bedrock + nullable: true + type: string tokenBudget: description: |- Token-budget caps. Optional — absent ⇒ no budget enforcement. @@ -262,9 +308,12 @@ spec: - appliesTo type: object x-kubernetes-validations: - - message: spec.bundleRef is mutually exclusive with spec.tokenBudget, spec.contentSafety, spec.modelPreference, and spec.displayName; the bundle carries those content fields + - message: spec.bundleRef is mutually exclusive with spec.tokenBudget, spec.contentSafety, spec.modelPreference, spec.provider, spec.guardrails, and spec.displayName; the bundle carries those content fields + reason: FieldValueInvalid + rule: '!has(self.bundleRef) || (!has(self.tokenBudget) && !has(self.contentSafety) && !has(self.modelPreference) && !has(self.provider) && !has(self.guardrails) && !has(self.displayName))' + - message: spec.guardrails, when set, must contain 1-8 stages (omit the field for no pipeline) reason: FieldValueInvalid - rule: '!has(self.bundleRef) || (!has(self.tokenBudget) && !has(self.contentSafety) && !has(self.modelPreference) && !has(self.displayName))' + rule: '!has(self.guardrails) || (size(self.guardrails) >= 1 && size(self.guardrails) <= 8)' - message: spec.tokenBudget.monthlyTokens must be >= spec.tokenBudget.dailyTokens reason: FieldValueInvalid rule: '!has(self.tokenBudget) || !has(self.tokenBudget.monthlyTokens) || !has(self.tokenBudget.dailyTokens) || self.tokenBudget.monthlyTokens >= self.tokenBudget.dailyTokens' @@ -403,4 +452,3 @@ spec: storage: true subresources: status: {} - diff --git a/docs/api/crd-reference.md b/docs/api/crd-reference.md index 3f7c49029..f2009f861 100644 --- a/docs/api/crd-reference.md +++ b/docs/api/crd-reference.md @@ -458,6 +458,10 @@ spec: sandboxName: my-agent # exact match; empty = any in ns sandboxMatchLabels: {} # AND with sandboxName action: "*" # chat | responses | image | embeddings | * + provider: azure-openai # optional: azure-openai | anthropic | ollama | bedrock + guardrails: # optional: ordered pipeline, first flag blocks + - provider: openai-moderation + applyTo: both # input | output | both (default both) modelPreference: primary: provider: azure-openai # azure-openai | anthropic | gemini | bedrock | ollama @@ -480,13 +484,17 @@ spec: | Field | Notes | |---|---| | `spec.appliesTo` | Required selector — AND of `sandboxName`, `sandboxMatchLabels`, `action`. | -| `spec.modelPreference.primary` | `{provider, deployment}`. `provider` is one of `azure-openai`, `anthropic`, `gemini`, `bedrock`, `ollama`. | -| `spec.modelPreference.fallback[]` | Ordered fallback routes — first healthy wins, deterministically. No load-balancing. | +| `spec.provider` | Optional typed default provider (`azure-openai` \| `anthropic` \| `ollama` \| `bedrock`), and the **only** field that drives provider routing. Absent ⇒ the env-configured Azure OpenAI / Foundry upstream. `anthropic` serves the Anthropic Messages surface (`/anthropic/v1/messages`, streaming + tools pass-through) with the router-held `ANTHROPIC_API_KEY`; `ollama` serves OpenAI-compatible chat completions against `OLLAMA_ENDPOINT` (no credential). `bedrock` is schema-accepted but the router returns 501 until the Bedrock client lands. Routing applies to `/v1/chat/completions` and the Anthropic Messages routes only; the other inference routes (see note below) refuse a non-Azure provider with 501. Credentials/endpoints are router-sidecar config only — never visible to the agent. | +| `spec.guardrails[]` | Optional ordered guardrail pipeline (1–8 stages) the router runs on `/v1/chat/completions` and the Anthropic Messages routes. Stage: `{provider, applyTo}` with `provider: openai-moderation` (Bedrock Guardrails / Model Armor are roadmap follow-ups) and `applyTo: input \| output \| both`. Fail-closed: a declared stage that cannot run (missing key, backend outage) blocks the request. Streaming responses use hold-and-release windows — no assistant message text reaches the client before a scan covers it (tool-call arguments and provider "thinking" deltas are not yet scanned — see note). Text over 16k chars is scanned in successive windows, not truncated. | +| `spec.modelPreference.primary` | `{provider, deployment}`. `provider` is one of `azure-openai`, `anthropic`, `gemini`, `bedrock`, `ollama`. This tag stays **informational** — it does not select the upstream (use `spec.provider`); `modelPreference` drives deployment failover within the resolved provider. | +| `spec.modelPreference.fallback[]` | Ordered fallback routes — first healthy wins, deterministically. No load-balancing. Failover walks deployments on the resolved provider (cross-provider failover is a follow-up). | | `spec.tokenBudget.perRequestTokens` | Per-call hard cap. Inference calls exceeding this are refused **before** the upstream forward. | | `spec.tokenBudget.dailyTokens` / `monthlyTokens` | Accepted and surfaced in status; **aggregate enforcement is not yet wired** — see roadmap below. CEL enforces `monthlyTokens ≥ dailyTokens`. | | `spec.contentSafety` | Per-category severity floors (`Safe` \| `Low` \| `Medium` \| `High`). The router parses Foundry `prompt_filter_results` inline; there is **no** separate Content Safety call. | | `spec.contentSafety.requirePromptShields` | Fail-closed if Prompt Shields are advertised by the deployment but the response lacks the corresponding annotations. | -| `spec.bundleRef` | Signed OCI artifact alternative to inline `tokenBudget` / `contentSafety` / `modelPreference` / `displayName`. `appliesTo` always comes from the CR. | +| `spec.bundleRef` | Signed OCI artifact alternative to inline `tokenBudget` / `contentSafety` / `modelPreference` / `provider` / `guardrails` / `displayName`. `appliesTo` always comes from the CR. | + +> **Provider + guardrail enforcement scope today.** `spec.provider` routing and `spec.guardrails[]` scanning are wired on `/v1/chat/completions` and the Anthropic Messages routes (`/anthropic/v1/messages`, `/v1/messages`). Every other inference-bearing route, `/v1/completions`, `/v1/responses`, `/v1/embeddings`, image generation, and the Foundry proxy families that serve model output (`/agents*`, `/openai/responses*`, `/openai/conversations*`), does **not** implement either; rather than silently bypass the policy, they **fail closed** (501 for a non-Azure `spec.provider`, 403 when guardrails are declared). A plain Azure policy with no guardrails uses those routes exactly as before. The Foundry proxy guard is default-deny: only an explicit exempt set of canonical management/storage APIs (memory stores, knowledge bases, evaluations, files, vector stores, containers, and similar) bypasses it, because those surfaces are not general-purpose inference channels; every other or unknown Foundry path fails closed. Foundry proxy paths are canonicalized before routing: benign empty segments (`//`, trailing `/`) are collapsed, while dot segments (`.`/`..`, literal or percent-encoded), encoded slash/backslash forms, and malformed escapes are rejected outright (400 `invalid_path`) so a request cannot normalize into a different upstream route than the one classified. Guardrail output scanning covers assistant message text; tool-call arguments and provider extended-"thinking" deltas are not yet scanned. These gaps are on the roadmap. > **Budget enforcement scope today.** The router enforces `tokenBudget.perRequestTokens` on every model call. Aggregate counters across requests (`dailyTokens`, `monthlyTokens`) are **not yet persisted**; the fields are accepted and surfaced for forward compatibility but only the per-request limit fires denials today. Aggregate enforcement is on the roadmap — see [`docs/roadmap.md`](../roadmap.md#trust-topology-end-to-end). diff --git a/docs/architecture.md b/docs/architecture.md index 83e3fe531..4bd2cb3fa 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -145,7 +145,7 @@ In prose: > **Sub-agent inheritance.** When a parent spawns a sub-agent (`/sandbox/spawn`), the router propagates `OPENCLAW_MODEL`, `KARS_PROVIDER`, the upstream endpoint, and the auth credential (Copilot OAuth token or PAT) into the new container's environment. The child uses the same provider + model + credentials as its parent without per-spawn wiring. -> **More providers later.** Copilot, Foundry, and GitHub Models are the three backends wired in today. Adding more (direct Anthropic, Bedrock, AWS Q, third-party OpenAI-compatible gateways) is mostly a matter of an endpoint+auth recipe in `inference-router/src/proxy.rs::build_upstream_url` plus a CLI prompt branch. We're tracking provider-expansion through GitHub issues — please open a feature request describing the provider, auth model, and which Foundry-only features (if any) you'd want preserved. +> **More providers later.** Copilot, Foundry, GitHub Models, direct Anthropic (`InferencePolicy.spec.provider: anthropic` — native Messages pass-through, router-held `ANTHROPIC_API_KEY`), and OpenAI-compatible Ollama (`provider: ollama`, `OLLAMA_ENDPOINT`) are the backends wired in today. Adding more (Bedrock, Vertex, vLLM, third-party OpenAI-compatible gateways) is mostly a matter of an endpoint+auth recipe in `inference-router/src/provider.rs` + `proxy.rs::build_upstream_url` plus a CLI prompt branch. We're tracking provider-expansion through GitHub issues — please open a feature request describing the provider, auth model, and which Foundry-only features (if any) you'd want preserved. Every other external call (web fetch, MCP tool, sub-agent spawn, A2A peer message) goes through the same shape with a different policy module. The handler is `chat_completions_handler` in [`inference-router/src/routes/chat_completions.rs`](../inference-router/src/routes/chat_completions.rs). diff --git a/docs/security-audits/2026-08-25-multi-provider-guardrails.md b/docs/security-audits/2026-08-25-multi-provider-guardrails.md new file mode 100644 index 000000000..ebce1cd36 --- /dev/null +++ b/docs/security-audits/2026-08-25-multi-provider-guardrails.md @@ -0,0 +1,116 @@ +# Security Audit, Multi-provider LLM upstreams + pluggable guardrail pipeline (PR #488) + +Date: 2026-08-25 +Scope: `inference-router/src/routes/` (`chat_completions.rs`, `anthropic_messages.rs`, +`inference.rs`, `mod.rs`), `inference-router/src/guardrails/`, +`inference-router/src/provider.rs`, `inference-router/src/proxy.rs`, +`controller/src/inference_policy*.rs`, `controller/src/crd_validations.rs`, +`controller/src/reconciler/`. +Gated paths: `inference-router/src/routes/*`, `controller/src/reconciler/*`. + +## Summary + +Adds policy-driven routing to non-Azure LLM providers (Anthropic native, Ollama +OpenAI-compat; Bedrock stubbed → 501) via a new `InferencePolicy.spec.provider`, +and an ordered, fail-closed guardrail pipeline (`spec.guardrails[]`; first backend +OpenAI Moderation) that scans request input and model output for buffered and SSE +streaming responses. Provider credentials are held router-side only and never +reach the agent container. This audit also covers the review-driven hardening: +the raw-path/dot-segment guard on the Foundry proxy, and six functional +fail-closed/accounting fixes. + +## T1: New capability / attack surface? (YES) + +- **New egress targets.** The router can now originate requests to Anthropic and + Ollama endpoints. Targets are derived from router-side config + (`ANTHROPIC_ENDPOINT`, `OLLAMA_ENDPOINT`), not from request or CR content, so + the destinations are operator-controlled. Note the provider forward paths do + NOT currently invoke the blocklist (`is_blocked`); the blocklist/egress guard + applies to the Foundry proxy path, not these provider calls. `provider::resolve` + fails closed (501 unimplemented / 503 unconfigured) rather than defaulting to an + attacker-influenced target. +- **New outbound credential use.** Anthropic uses a router-held `x-api-key` + (`ANTHROPIC_API_KEY`); Ollama is unauthenticated. Inbound agent-supplied + `x-api-key` / `authorization` are stripped and replaced before forwarding + (verified by `multi_provider_guardrails` tests). Keys are forwarded only to the + router sidecar, never the agent container (`reconciler` skip-empty env block), + and only endpoints (never keys) enter the config hash. +- **New moderation backend call.** The guardrail pipeline calls the configured + OpenAI Moderation endpoint with a router-held key. A declared stage with no key + fails pipeline construction (503 `guardrail_misconfigured`). +- **New public request surface reachable pre-auth.** The guardrail/provider + enforcement runs on the unauthenticated public router. Enforcement scope is + bounded and documented (see T2). + +## T2: Security-control change? (YES, net strengthening, fail-closed) + +- **Guardrail enforcement is fail-closed throughout.** Unbuildable pipeline → + 503; backend outage → 502; violation → 403; streaming uses hold-and-release so + no model text reaches the client before a scan covers it, and a flagged scan + cuts the stream. +- **Route-coverage guard (default-deny).** Routes that cannot run the pipeline + refuse a policy that needs it rather than silently bypassing: `/v1/completions`, + `/v1/responses`, `/v1/embeddings`, image generation, and the Foundry proxy + inference families (`/agents*`, `/openai/responses*`, `/openai/conversations*`). + The Foundry proxy guard is **default-deny**: everything is guarded except an + explicit exempt list of canonical management/storage APIs. +- **Path-canonicalization guard (bypass fix).** `foundry_proxy` percent-decodes + each path segment, canonicalizes benign empty segments (`//`, trailing + `/`), and rejects (400 `invalid_path`) any dot segment (`.`/`..`, literal + or percent-encoded), encoded slash/backslash, or malformed escape before + classification or forwarding, closing a traversal + bypass where `/openai/files/../responses` normalized upstream into a guarded + inference route. Proven live (mock upstream received `/openai/responses` + pre-fix; 400 post-fix) and pinned by `foundry_route_guard` tests including a + `reqwest::Url` normalization premise test. +- **Malformed guardrail config fails closed.** A present-but-malformed + `guardrails` block no longer degrades to "no guardrails"; it poisons the policy + so every request refuses (previously a silent disable of both the scan and the + route-gap guard). +- **Buffered non-JSON output now scanned.** Output enforcement was hoisted out of + the JSON-parse block so a non-JSON/truncated upstream body is still scanned via + the raw-text fallback rather than returned verbatim. +- **UTF-8 integrity of scanned text.** SSE moderation now reconstructs multi-byte + code points split across chunk boundaries, so the moderator scans the same text + the client receives (no `U+FFFD` divergence). +- **All choices scanned.** Streaming extraction reads every `choices[]` entry, not + just `choices[0]`. +- **Credential handling unchanged elsewhere.** No change to Azure WI/IMDS/Copilot + auth; hop-by-hop headers stripped on Anthropic pass-through relays. + +## T3: Availability / fail-open risk? (NEUTRAL-to-INCREASED-strictness) + +- Fail-closed choices can turn backend outages into request denials (moderation + 502, provider 503). This is intended and documented; a plain Azure policy with + no guardrails is unaffected and uses every route exactly as before + (back-compat tests: `absent_provider_and_guardrails_keep_backcompat_defaults`). +- Streaming Anthropic now records token usage, closing a budget-accounting gap + (no new denial path; best-effort, recorded once at end-of-stream). +- SSE hold-and-release bounds retained scan context (`MAX_SCAN_CHARS`), so memory + stays bounded on long streams (regression-pinned). +- No new unbounded allocation, no new panics on request paths (path decoder and + UTF-8 carry are total; malformed input → 400 / lossy, never panic). + +## Verification + +- `cargo build --workspace`: clean. +- `cargo clippy --workspace --all-targets`: 0 warnings (includes the CI + `-D warnings` lints, `result_large_err` boxed, `collapsible_else_if` collapsed). +- `cargo fmt --all --check`: clean. +- `cargo test`: full router + controller suites pass, including new tests: + `foundry_route_guard` (traversal matrix), `multi_provider_guardrails`, guardrail + UTF-8 split / multi-choice / malformed-config / fail-closed, Anthropic streaming + usage. +- LOC gate: `guardrails.rs` decomposed into `guardrails/{mod,backend,stream,tests}.rs` + (all < 800); `reconciler/mod.rs` reduced to < 3700 via `reconciler/pod_spec.rs`. +- Live reproduction of the traversal bypass and its fix recorded against a mock + upstream. + +## Verdict + +Accept, net fail-closed strengthening of the inference plane; the new egress and +credential surface is router-side-only, bounded, and default-deny; the identified +bypass and functional gaps are fixed and regression-tested. + +Signed-off-by: John Seong +Signed-off-by: Pal Lakatos-Toth diff --git a/inference-router/src/config.rs b/inference-router/src/config.rs index 2707674a6..3438a106b 100644 --- a/inference-router/src/config.rs +++ b/inference-router/src/config.rs @@ -81,6 +81,49 @@ pub struct Config { /// Captured at config-load time so provider detection is a pure /// function on the `Config` struct (testable without env hacks). pub provider_override: Option, + + /// Anthropic Messages API endpoint used when an `InferencePolicy` + /// selects `provider: anthropic`. Default `https://api.anthropic.com`; + /// override with `ANTHROPIC_ENDPOINT` for gateways. + pub anthropic_endpoint: String, + + /// Anthropic API key — `ANTHROPIC_API_KEY` env or the + /// `anthropic-api-key` secret mount. Router-side only. `None` ⇒ + /// Anthropic policies fail closed. + pub anthropic_api_key: Option, + + /// OpenAI-compatible Ollama endpoint (e.g. + /// `http://ollama.ollama.svc:11434`) used when an + /// `InferencePolicy` selects `provider: ollama`. No default — the + /// operator must opt in via `OLLAMA_ENDPOINT`. + pub ollama_endpoint: Option, + + /// Base endpoint for the OpenAI Moderation guardrail stage. + /// Default `https://api.openai.com`; override with + /// `OPENAI_MODERATION_ENDPOINT`. + pub openai_moderation_endpoint: String, + + /// OpenAI Moderation key — `OPENAI_MODERATION_API_KEY` env (falls + /// back to `OPENAI_API_KEY`, then the `openai-moderation-api-key` + /// secret mount). `None` ⇒ `openai-moderation` stages fail closed. + pub openai_moderation_api_key: Option, + + /// Moderation model (`OPENAI_MODERATION_MODEL`, default + /// `omni-moderation-latest`). + pub openai_moderation_model: String, +} + +/// Read a credential from an env var, falling back to the standard +/// kars secret mounts (`/etc/kars/secrets/` then +/// `/run/secrets/`). Mirrors the admin-token lookup in +/// `routes::AppState::new`. Empty values are treated as unset. +fn secret_from_env_or_mount(env: &str, file: &str) -> Option { + std::env::var(env) + .ok() + .or_else(|| std::fs::read_to_string(format!("/etc/kars/secrets/{file}")).ok()) + .or_else(|| std::fs::read_to_string(format!("/run/secrets/{file}")).ok()) + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) } impl Config { @@ -145,6 +188,38 @@ impl Config { .ok() .filter(|s| !s.is_empty()) .map(|s| s.to_ascii_lowercase()), + + anthropic_endpoint: std::env::var("ANTHROPIC_ENDPOINT") + .ok() + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| "https://api.anthropic.com".into()), + + anthropic_api_key: secret_from_env_or_mount("ANTHROPIC_API_KEY", "anthropic-api-key"), + + ollama_endpoint: std::env::var("OLLAMA_ENDPOINT") + .ok() + .filter(|s| !s.is_empty()), + + openai_moderation_endpoint: std::env::var("OPENAI_MODERATION_ENDPOINT") + .ok() + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| "https://api.openai.com".into()), + + openai_moderation_api_key: secret_from_env_or_mount( + "OPENAI_MODERATION_API_KEY", + "openai-moderation-api-key", + ) + .or_else(|| { + std::env::var("OPENAI_API_KEY") + .ok() + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + }), + + openai_moderation_model: std::env::var("OPENAI_MODERATION_MODEL") + .ok() + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| "omni-moderation-latest".into()), }) } @@ -212,6 +287,12 @@ mod tests { registry_mode: RegistryMode::Local, registry_url: None, provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: None, + ollama_endpoint: None, + openai_moderation_endpoint: "https://api.openai.com".into(), + openai_moderation_api_key: None, + openai_moderation_model: "omni-moderation-latest".into(), } } diff --git a/inference-router/src/errors.rs b/inference-router/src/errors.rs index e1ce3b560..287729ba4 100644 --- a/inference-router/src/errors.rs +++ b/inference-router/src/errors.rs @@ -84,6 +84,32 @@ pub fn openai( ) } +/// OpenAI-shape error that additionally carries a machine-readable +/// `code`: `{"error": {"message", "type", "code"}}`. +/// +/// Use for provider/guardrail policy denials so a client can switch on +/// a single stable `error.code` across every inference route +/// (chat-completions, the sibling `/v1/*` routes, and the Anthropic +/// Messages route, which mirrors `code` into its own error object). +/// Additive to [`openai`], `type` keeps its per-route value. +pub fn openai_coded( + status: StatusCode, + msg: impl Into, + type_: &str, + code: &str, +) -> (StatusCode, Json) { + ( + status, + Json(json!({ + "error": { + "message": msg.into(), + "type": type_, + "code": code, + } + })), + ) +} + #[cfg(test)] mod tests { use super::*; diff --git a/inference-router/src/failover.rs b/inference-router/src/failover.rs index 91642637f..7fa2245c4 100644 --- a/inference-router/src/failover.rs +++ b/inference-router/src/failover.rs @@ -253,6 +253,8 @@ mod tests { endpoint: "https://example.openai.azure.com".into(), deployment: dep.to_string(), sandbox_name: "sbx".into(), + provider: crate::provider::ProviderKind::AzureOpenAI, + api_key: None, } } diff --git a/inference-router/src/guardrails/backend.rs b/inference-router/src/guardrails/backend.rs new file mode 100644 index 000000000..fe6591703 --- /dev/null +++ b/inference-router/src/guardrails/backend.rs @@ -0,0 +1,372 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! OpenAI Moderation backend, the ordered guardrail pipeline, and the +//! request/response text-extraction helpers. + +use std::sync::Arc; + +use async_trait::async_trait; +use reqwest::Client; + +use crate::config::Config; +use crate::metrics; + +use super::{ + ApplyTo, Direction, Guardrail, GuardrailError, GuardrailStageCfg, GuardrailVerdict, + GuardrailViolation, MAX_SCAN_CHARS, +}; + +// ─── OpenAI Moderation backend ─────────────────────────────────────────────── + +/// OpenAI Moderation API backend (`POST {endpoint}/v1/moderations`). +pub struct OpenAiModeration { + client: Client, + endpoint: String, + api_key: String, + model: String, +} + +impl OpenAiModeration { + #[must_use] + pub fn new(client: Client, endpoint: String, api_key: String, model: String) -> Self { + Self { + client, + endpoint, + api_key, + model, + } + } +} + +/// Parse a Moderation API response body into a verdict. Pure — unit +/// tested without I/O. Missing/malformed `results` is an error, not a +/// pass: an unparseable verdict must fail closed. +pub fn parse_moderation_response(body: &serde_json::Value) -> Result { + let result = body + .get("results") + .and_then(|r| r.as_array()) + .and_then(|r| r.first()) + .ok_or_else(|| "moderation response missing results[0]".to_string())?; + let flagged = result + .get("flagged") + .and_then(|f| f.as_bool()) + .ok_or_else(|| "moderation response missing results[0].flagged".to_string())?; + let categories = result + .get("categories") + .and_then(|c| c.as_object()) + .map(|c| { + c.iter() + .filter(|(_, v)| v.as_bool() == Some(true)) + .map(|(k, _)| k.clone()) + .collect() + }) + .unwrap_or_default(); + Ok(GuardrailVerdict { + flagged, + categories, + }) +} + +#[async_trait] +impl Guardrail for OpenAiModeration { + fn name(&self) -> &'static str { + "openai-moderation" + } + + async fn scan(&self, text: &str) -> Result { + let url = format!( + "{}/v1/moderations", + self.endpoint.trim_end_matches('/').trim_end_matches("/v1") + ); + let response = self + .client + .post(&url) + .bearer_auth(&self.api_key) + .json(&serde_json::json!({ "model": self.model, "input": text })) + .timeout(std::time::Duration::from_secs(10)) + .send() + .await + .map_err(|e| GuardrailError::Unavailable { + provider: "openai-moderation", + reason: format!("transport error: {e}"), + })?; + let status = response.status(); + if !status.is_success() { + let body = response.text().await.unwrap_or_default(); + let preview: String = body.chars().take(512).collect(); + return Err(GuardrailError::Unavailable { + provider: "openai-moderation", + reason: format!("upstream status {status}: {preview}"), + }); + } + let body: serde_json::Value = + response + .json() + .await + .map_err(|e| GuardrailError::Unavailable { + provider: "openai-moderation", + reason: format!("non-JSON response: {e}"), + })?; + parse_moderation_response(&body).map_err(|reason| GuardrailError::Unavailable { + provider: "openai-moderation", + reason, + }) + } +} + +// ─── Pipeline ──────────────────────────────────────────────────────────────── + +pub(crate) struct BuiltStage { + pub(crate) apply_to: ApplyTo, + pub(crate) guard: Arc, +} + +/// Ordered guardrail stages materialised from a policy snapshot. +/// Cheap to build per request (clones a shared `reqwest::Client`). +pub struct GuardrailPipeline { + pub(crate) stages: Vec, +} + +impl GuardrailPipeline { + /// Build from the compiled-policy stage list. A stage naming an + /// unknown backend, or one whose router-side credential/endpoint + /// is absent, is a construction error — see the module-level + /// fail-closed contract. + pub fn from_stages( + stages: &[GuardrailStageCfg], + config: &Config, + client: &Client, + ) -> Result { + let mut built = Vec::with_capacity(stages.len()); + for stage in stages { + match stage.provider.trim().to_ascii_lowercase().as_str() { + "openai-moderation" => { + let api_key = config.openai_moderation_api_key.clone().ok_or_else(|| { + GuardrailError::Config { + provider: stage.provider.clone(), + reason: "no API key configured (OPENAI_MODERATION_API_KEY, \ + OPENAI_API_KEY, or secret mount)" + .into(), + } + })?; + built.push(BuiltStage { + apply_to: stage.apply_to, + guard: Arc::new(OpenAiModeration::new( + client.clone(), + config.openai_moderation_endpoint.clone(), + api_key, + config.openai_moderation_model.clone(), + )), + }); + } + other => { + return Err(GuardrailError::Config { + provider: other.to_string(), + reason: "unknown guardrail backend".into(), + }); + } + } + } + Ok(Self { stages: built }) + } + + /// True when at least one stage covers `direction` — callers use + /// this to skip text extraction entirely on the hot path. + #[must_use] + pub fn covers(&self, direction: Direction) -> bool { + self.stages.iter().any(|s| s.apply_to.covers(direction)) + } + + /// Run every stage covering `direction` over `text`, first flag + /// wins. Text over [`MAX_SCAN_CHARS`] is scanned in successive + /// windows, not truncated, so content can't be hidden past the + /// cap. + pub async fn scan( + &self, + text: &str, + direction: Direction, + ) -> Result, GuardrailError> { + if text.is_empty() { + return Ok(None); + } + let windows = scan_windows(text); + for stage in self.stages.iter().filter(|s| s.apply_to.covers(direction)) { + for window in &windows { + match stage.guard.scan(window).await { + Ok(verdict) if verdict.flagged => { + metrics::GUARDRAIL_SCANS + .with_label_values(&[stage.guard.name(), direction.as_str(), "flagged"]) + .inc(); + return Ok(Some(GuardrailViolation { + provider: stage.guard.name(), + direction, + categories: verdict.categories, + })); + } + Ok(_) => { + metrics::GUARDRAIL_SCANS + .with_label_values(&[stage.guard.name(), direction.as_str(), "pass"]) + .inc(); + } + Err(e) => { + metrics::GUARDRAIL_SCANS + .with_label_values(&[stage.guard.name(), direction.as_str(), "error"]) + .inc(); + return Err(e); + } + } + } + } + Ok(None) + } +} + +/// Split `text` into consecutive windows of at most [`MAX_SCAN_CHARS`] +/// chars (never mid-char) so scanning all windows covers everything. +pub(crate) fn scan_windows(text: &str) -> Vec<&str> { + if text.len() <= MAX_SCAN_CHARS { + return vec![text]; + } + let mut windows = Vec::new(); + let mut start = 0; + let mut count = 0; + for (i, _) in text.char_indices() { + if count == MAX_SCAN_CHARS { + windows.push(&text[start..i]); + start = i; + count = 0; + } + count += 1; + } + windows.push(&text[start..]); + windows +} + +// ─── Request / response text extraction ────────────────────────────────────── + +/// Extract the human-visible text of an OpenAI chat-completions +/// request body: every `messages[].content` string, plus `text` +/// fields of array-shaped content parts. +#[must_use] +pub fn extract_openai_input_text(body: &serde_json::Value) -> String { + let mut out: Vec = Vec::new(); + if let Some(messages) = body.get("messages").and_then(|m| m.as_array()) { + for m in messages { + match m.get("content") { + Some(serde_json::Value::String(s)) if !s.is_empty() => out.push(s.clone()), + Some(serde_json::Value::Array(parts)) => { + for p in parts { + if let Some(t) = p.get("text").and_then(|t| t.as_str()) + && !t.is_empty() + { + out.push(t.to_string()); + } + } + } + _ => {} + } + } + } + out.join("\n") +} + +/// Extract the human-visible text of an Anthropic Messages request +/// body: `system` (string or parts) plus `messages[].content` text / +/// `tool_result` strings. +#[must_use] +pub fn extract_anthropic_input_text(body: &serde_json::Value) -> String { + let mut out: Vec = Vec::new(); + match body.get("system") { + Some(serde_json::Value::String(s)) if !s.is_empty() => out.push(s.clone()), + Some(serde_json::Value::Array(parts)) => { + for p in parts { + if let Some(t) = p.get("text").and_then(|t| t.as_str()) + && !t.is_empty() + { + out.push(t.to_string()); + } + } + } + _ => {} + } + if let Some(messages) = body.get("messages").and_then(|m| m.as_array()) { + for m in messages { + match m.get("content") { + Some(serde_json::Value::String(s)) if !s.is_empty() => out.push(s.clone()), + Some(serde_json::Value::Array(parts)) => { + for p in parts { + match p.get("type").and_then(|t| t.as_str()) { + Some("text") => { + if let Some(t) = p.get("text").and_then(|t| t.as_str()) + && !t.is_empty() + { + out.push(t.to_string()); + } + } + Some("tool_result") => { + if let Some(t) = p.get("content").and_then(|c| c.as_str()) + && !t.is_empty() + { + out.push(t.to_string()); + } + } + _ => {} + } + } + } + _ => {} + } + } + } + out.join("\n") +} + +/// Extract the assistant text of a buffered OpenAI chat-completions +/// response (`choices[*].message.content`). +#[must_use] +pub fn extract_openai_output_text(body: &serde_json::Value) -> String { + let mut out: Vec = Vec::new(); + if let Some(choices) = body.get("choices").and_then(|c| c.as_array()) { + for c in choices { + if let Some(t) = c + .get("message") + .and_then(|m| m.get("content")) + .and_then(|t| t.as_str()) + && !t.is_empty() + { + out.push(t.to_string()); + } + } + } + out.join("\n") +} + +/// Extract the assistant text of a buffered Anthropic Messages +/// response (`content[*].text`). +#[must_use] +pub fn extract_anthropic_output_text(body: &serde_json::Value) -> String { + let mut out: Vec = Vec::new(); + if let Some(content) = body.get("content").and_then(|c| c.as_array()) { + for block in content { + if let Some(t) = block.get("text").and_then(|t| t.as_str()) + && !t.is_empty() + { + out.push(t.to_string()); + } + } + } + out.join("\n") +} + +/// Extract scan text via `extract`, falling back to raw lossy-UTF-8 +/// bytes when the body isn't JSON, so a declared guardrail is never +/// skipped on a parse failure. (Parsed-but-empty extraction — e.g. a +/// tool-call-only response — is a deliberate pass.) +#[must_use] +pub fn scan_text_or_raw(body: &[u8], extract: impl FnOnce(&serde_json::Value) -> String) -> String { + match serde_json::from_slice::(body) { + Ok(v) => extract(&v), + Err(_) => String::from_utf8_lossy(body).into_owned(), + } +} diff --git a/inference-router/src/guardrails/mod.rs b/inference-router/src/guardrails/mod.rs new file mode 100644 index 000000000..ffe7feea4 --- /dev/null +++ b/inference-router/src/guardrails/mod.rs @@ -0,0 +1,225 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Pluggable guardrail pipeline for `InferencePolicy.spec.guardrails[]`. +//! +//! Stages run around each governed call — request text pre-flight and +//! response text (buffered + streaming). First backend is OpenAI +//! Moderation; the [`Guardrail`] trait is the extension point. +//! +//! Fail-closed: a declared stage that can't be built (unknown backend +//! / missing credential) or errors at runtime blocks the request +//! rather than passing unscanned content. +//! +//! Streaming uses hold-and-release: SSE chunks are withheld until the +//! accumulated text reaches [`STREAM_SCAN_THRESHOLD_CHARS`] (or the +//! stream ends) and a scan clears it, so no model text reaches the +//! client unscanned; a flagged scan cuts the stream with an error +//! frame. Text over [`MAX_SCAN_CHARS`] is scanned in successive +//! windows ([`scan_windows`]), never truncated. +//! +//! Split across `backend` (moderation + pipeline + extraction) and +//! `stream` (SSE hold-and-release) to bound file size; this module +//! keeps the shared config/verdict/error types and re-exports both. + +use async_trait::async_trait; + +mod backend; +mod stream; + +pub use backend::*; +pub use stream::*; + +/// Upper bound on characters submitted to a backend in one scan call. +pub const MAX_SCAN_CHARS: usize = 16_000; + +/// Default hold-and-release window for streaming output scans, in +/// characters of extracted delta text. Override with +/// `GUARDRAIL_STREAM_SCAN_CHARS`. +pub const STREAM_SCAN_THRESHOLD_CHARS: usize = 1_000; + +/// Env override for [`STREAM_SCAN_THRESHOLD_CHARS`]. +pub const STREAM_SCAN_THRESHOLD_ENV: &str = "GUARDRAIL_STREAM_SCAN_CHARS"; + +/// Scan direction, from the policy's `applyTo`. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum ApplyTo { + Input, + Output, + #[default] + Both, +} + +impl ApplyTo { + /// Liberal parse of the compiled-profile string. Unknown values + /// widen to `Both` — scanning more than asked is safe; scanning + /// less is not. + #[must_use] + pub fn parse(s: Option<&str>) -> Self { + match s.map(str::trim) { + Some(v) if v.eq_ignore_ascii_case("input") => Self::Input, + Some(v) if v.eq_ignore_ascii_case("output") => Self::Output, + Some(v) if v.eq_ignore_ascii_case("both") || v.is_empty() => Self::Both, + None => Self::Both, + Some(other) => { + tracing::warn!( + apply_to = other, + "guardrail applyTo not recognised — widening to 'both'" + ); + Self::Both + } + } + } + + #[must_use] + pub fn covers(&self, direction: Direction) -> bool { + matches!( + (self, direction), + (Self::Both, _) | (Self::Input, Direction::Input) | (Self::Output, Direction::Output) + ) + } +} + +/// Which side of the inference call a scan covers. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Direction { + Input, + Output, +} + +impl Direction { + #[must_use] + pub fn as_str(&self) -> &'static str { + match self { + Self::Input => "input", + Self::Output => "output", + } + } +} + +/// One compiled `guardrails[]` stage (`{provider, applyTo}`). Parsed +/// liberally; unknown-backend rejection happens at pipeline +/// construction, where a request exists to fail closed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GuardrailStageCfg { + pub provider: String, + pub apply_to: ApplyTo, +} + +impl GuardrailStageCfg { + /// Parse the compiled `guardrails` block. + /// + /// `null`/absent ⇒ `Ok(empty)` (legitimately no pipeline). A + /// present-but-malformed value, not an array, or an entry missing + /// a string `provider`, is `Err`: a declared-but-unbuildable + /// control must fail closed, never be silently dropped to "no + /// guardrails" (which would also disable the sibling route-gap + /// guard). The caller poisons the policy so every request refuses. + pub fn from_compiled_json(v: &serde_json::Value) -> Result, String> { + if v.is_null() { + return Ok(Vec::new()); + } + let Some(arr) = v.as_array() else { + return Err(format!( + "`guardrails` must be an array or null, got {}", + json_kind(v) + )); + }; + let mut out = Vec::with_capacity(arr.len()); + for (i, stage) in arr.iter().enumerate() { + let Some(provider) = stage.get("provider").and_then(|p| p.as_str()) else { + return Err(format!( + "guardrails[{i}] missing string `provider`: {stage}" + )); + }; + out.push(Self { + provider: provider.to_string(), + apply_to: ApplyTo::parse(stage.get("applyTo").and_then(|a| a.as_str())), + }); + } + Ok(out) + } +} + +/// One-word JSON kind for error messages. +fn json_kind(v: &serde_json::Value) -> &'static str { + match v { + serde_json::Value::Null => "null", + serde_json::Value::Bool(_) => "bool", + serde_json::Value::Number(_) => "number", + serde_json::Value::String(_) => "string", + serde_json::Value::Array(_) => "array", + serde_json::Value::Object(_) => "object", + } +} + +/// A guardrail verdict for one scanned text. +#[derive(Debug, Clone, Default)] +pub struct GuardrailVerdict { + pub flagged: bool, + /// Backend-specific category names that flagged (e.g. + /// `violence`, `hate/threatening`). + pub categories: Vec, +} + +/// A confirmed violation, carrying enough context for the audit log +/// and the client-facing error body. +#[derive(Debug, Clone)] +pub struct GuardrailViolation { + pub provider: &'static str, + pub direction: Direction, + pub categories: Vec, +} + +impl GuardrailViolation { + #[must_use] + pub fn message(&self) -> String { + format!( + "Blocked by guardrail '{}' ({}): flagged categories [{}]", + self.provider, + self.direction.as_str(), + self.categories.join(", ") + ) + } + + #[must_use] + pub fn code(&self) -> &'static str { + "guardrail_blocked" + } +} + +/// Errors from the pipeline. Both variants block the request +/// (fail-closed) but carry distinct codes so operators can tell a +/// config gap from a backend outage. +#[derive(Debug, thiserror::Error)] +pub enum GuardrailError { + #[error("guardrail stage '{provider}' cannot run: {reason}")] + Config { provider: String, reason: String }, + #[error("guardrail '{provider}' scan failed: {reason}")] + Unavailable { + provider: &'static str, + reason: String, + }, +} + +impl GuardrailError { + #[must_use] + pub fn code(&self) -> &'static str { + match self { + Self::Config { .. } => "guardrail_misconfigured", + Self::Unavailable { .. } => "guardrail_unavailable", + } + } +} + +/// One guardrail backend. `scan` returns the backend's verdict for a +/// single text; transport/parse failures are `Err` and block the +/// request at the pipeline layer. +#[async_trait] +pub trait Guardrail: Send + Sync { + fn name(&self) -> &'static str; + async fn scan(&self, text: &str) -> Result; +} + +#[cfg(test)] +mod tests; diff --git a/inference-router/src/guardrails/stream.rs b/inference-router/src/guardrails/stream.rs new file mode 100644 index 000000000..345b44aa0 --- /dev/null +++ b/inference-router/src/guardrails/stream.rs @@ -0,0 +1,406 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Streaming (SSE) hold-and-release output guard. + +use std::sync::Arc; + +use bytes::Bytes; +use futures::stream::BoxStream; +use futures::stream::StreamExt; + +use super::{ + Direction, GuardrailError, GuardrailPipeline, GuardrailViolation, MAX_SCAN_CHARS, + STREAM_SCAN_THRESHOLD_CHARS, STREAM_SCAN_THRESHOLD_ENV, +}; + +// ─── Streaming (SSE) guard ─────────────────────────────────────────────────── + +/// SSE wire dialect of the guarded stream — decides how delta text is +/// extracted from `data:` events. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum StreamDialect { + /// OpenAI chat-completions chunks: `choices[0].delta.content`. + OpenAiChat, + /// Anthropic Messages events: `content_block_delta` → + /// `delta.text`. + AnthropicMessages, +} + +/// Extract delta text from one complete SSE `data:` JSON payload. +#[must_use] +pub(crate) fn delta_text_from_event( + dialect: StreamDialect, + event: &serde_json::Value, +) -> Option { + match dialect { + StreamDialect::OpenAiChat => { + // Every choice, not just `choices[0]`, an OpenAI-compatible + // upstream can emit multiple choices in one frame, and a + // later choice's text must not slip through unscanned. + let texts: Vec<&str> = event + .get("choices") + .and_then(|c| c.as_array()) + .map(|choices| { + choices + .iter() + .filter_map(|c| { + c.get("delta") + .and_then(|d| d.get("content")) + .and_then(|t| t.as_str()) + .filter(|t| !t.is_empty()) + }) + .collect() + }) + .unwrap_or_default(); + if texts.is_empty() { + None + } else { + Some(texts.join("\n")) + } + } + StreamDialect::AnthropicMessages => { + if event.get("type").and_then(|t| t.as_str()) == Some("content_block_delta") { + event + .get("delta") + .and_then(|d| d.get("text")) + .and_then(|t| t.as_str()) + .filter(|t| !t.is_empty()) + .map(str::to_string) + } else { + None + } + } + } +} + +/// Client-facing SSE error frame for a guardrail cut. The OpenAI-shape +/// error object works for both dialects' SDK error paths. +#[must_use] +pub fn violation_sse_frame(violation: &GuardrailViolation) -> Bytes { + Bytes::from(format!( + "data: {}\n\ndata: [DONE]\n\n", + serde_json::json!({ + "error": { + "message": violation.message(), + "type": "content_policy_violation", + "code": violation.code() + } + }) + )) +} + +/// SSE frame for a guardrail that could not run (config gap or +/// backend outage) — carries the error's own `type`/`code` so a +/// fail-closed cut is never mislabelled as a content violation. +#[must_use] +pub fn error_sse_frame(err: &GuardrailError) -> Bytes { + Bytes::from(format!( + "data: {}\n\ndata: [DONE]\n\n", + serde_json::json!({ + "error": { + "message": err.to_string(), + "type": "guardrail_error", + "code": err.code() + } + }) + )) +} + +/// Effective hold-and-release window size, clamped to +/// [`MAX_SCAN_CHARS`] so a window can never accumulate more text than +/// one scan covers. +fn stream_scan_threshold() -> usize { + std::env::var(STREAM_SCAN_THRESHOLD_ENV) + .ok() + .and_then(|v| v.parse().ok()) + .filter(|v: &usize| *v > 0) + .unwrap_or(STREAM_SCAN_THRESHOLD_CHARS) + .min(MAX_SCAN_CHARS) +} + +/// Hold-and-release state machine for one guarded SSE stream. Kept +/// separate from the stream adaptor so the release/hold/block logic +/// is unit-testable with a fake [`Guardrail`]. +pub(crate) struct SseGuardState { + pipeline: Arc, + dialect: StreamDialect, + threshold: usize, + /// Raw chunks held back until the text they carry has been + /// covered by a scan. + held: Vec, + /// Carry buffer for a trailing incomplete UTF-8 sequence split + /// across chunk boundaries, decoded once its continuation arrives + /// so moderation scans the same code points the client receives. + byte_carry: Vec, + /// Carry buffer for `data:` lines split across chunk boundaries. + line_carry: String, + /// All delta text accumulated so far (scan context). + pub(crate) accumulated: String, + /// Chars of `accumulated` not yet covered by a scan. + unscanned: usize, +} + +/// What the state machine wants the adaptor to emit next. +pub(crate) enum SseGuardStep { + /// Forward these bytes (possibly empty ⇒ nothing to emit yet). + Release(Vec), + /// Emit this terminal frame and drop the upstream stream. + Cut(Bytes), +} + +impl SseGuardState { + pub(crate) fn new( + pipeline: Arc, + dialect: StreamDialect, + threshold: usize, + ) -> Self { + Self { + pipeline, + dialect, + threshold, + held: Vec::new(), + byte_carry: Vec::new(), + line_carry: String::new(), + accumulated: String::new(), + unscanned: 0, + } + } + + fn ingest_text(&mut self, chunk: &[u8]) { + // Prepend any trailing incomplete UTF-8 bytes from the previous + // chunk, then decode only the complete-code-point prefix. A + // genuine incomplete sequence at the end is carried for the + // next chunk; a real mid-stream invalid byte is decoded + // lossily (matches the pre-existing behaviour for that + // pathological case) rather than carried forever. + let mut bytes = std::mem::take(&mut self.byte_carry); + bytes.extend_from_slice(chunk); + let decodable = match std::str::from_utf8(&bytes) { + Ok(_) => bytes.len(), + Err(e) if e.error_len().is_none() => e.valid_up_to(), + Err(_) => bytes.len(), + }; + self.byte_carry = bytes.split_off(decodable); + self.line_carry.push_str(&String::from_utf8_lossy(&bytes)); + let (complete, rest) = match self.line_carry.rfind('\n') { + Some(idx) => { + let (c, r) = self.line_carry.split_at(idx + 1); + (c.to_string(), r.to_string()) + } + None => (String::new(), std::mem::take(&mut self.line_carry)), + }; + self.line_carry = rest; + for line in complete.lines() { + self.ingest_line(line); + } + } + + fn ingest_line(&mut self, line: &str) { + // SSE permits `data:` with or without a following space. + let Some(payload) = line.trim().strip_prefix("data:") else { + return; + }; + let payload = payload.trim_start(); + if payload.is_empty() || payload == "[DONE]" { + return; + } + match serde_json::from_str::(payload) { + Ok(event) => { + if let Some(text) = delta_text_from_event(self.dialect, &event) { + self.unscanned += text.chars().count(); + self.accumulated.push_str(&text); + } + // Valid JSON with no delta text is a structural frame + // (ping / role-only / stop) — nothing to scan. + } + // Unrecognised (non-JSON) frame: scan the raw payload so + // it can't bypass the scan. + Err(_) => { + self.unscanned += payload.chars().count(); + self.accumulated.push_str(payload); + } + } + } + + pub(crate) async fn on_chunk(&mut self, chunk: Bytes) -> SseGuardStep { + self.ingest_text(&chunk); + self.held.push(chunk); + // A pending partial line, or a partial UTF-8 sequence, has + // bytes in `held` whose text is not yet counted; never release + // until it completes. + if !self.line_carry.is_empty() || !self.byte_carry.is_empty() { + return SseGuardStep::Release(Vec::new()); + } + if self.unscanned < self.threshold { + if self.unscanned == 0 { + return SseGuardStep::Release(std::mem::take(&mut self.held)); + } + return SseGuardStep::Release(Vec::new()); + } + self.scan_and_release().await + } + + async fn on_end(&mut self) -> SseGuardStep { + // Flush any leftover incomplete UTF-8 bytes (lossily, the + // stream ended mid-sequence) into the line carry first… + if !self.byte_carry.is_empty() { + let tail = std::mem::take(&mut self.byte_carry); + self.line_carry.push_str(&String::from_utf8_lossy(&tail)); + } + // …then flush a trailing unterminated line so it's scanned too. + if !self.line_carry.is_empty() { + let line = std::mem::take(&mut self.line_carry); + self.ingest_line(&line); + } + if self.unscanned == 0 { + return SseGuardStep::Release(std::mem::take(&mut self.held)); + } + self.scan_and_release().await + } + + async fn scan_and_release(&mut self) -> SseGuardStep { + match self + .pipeline + .scan(&self.accumulated, Direction::Output) + .await + { + Ok(None) => { + self.unscanned = 0; + self.trim_scan_context(); + SseGuardStep::Release(std::mem::take(&mut self.held)) + } + Ok(Some(violation)) => SseGuardStep::Cut(violation_sse_frame(&violation)), + Err(e) => SseGuardStep::Cut(error_sse_frame(&e)), + } + } + + /// After a clean scan, retain `MAX_SCAN_CHARS - threshold` chars + /// of context: bounds memory and keeps the next scan in one window + /// while overlapping for cross-boundary detection. + fn trim_scan_context(&mut self) { + let keep = MAX_SCAN_CHARS.saturating_sub(self.threshold).max(1); + if self.accumulated.chars().count() <= keep { + return; + } + let start = self + .accumulated + .char_indices() + .rev() + .nth(keep - 1) + .map_or(0, |(i, _)| i); + self.accumulated = self.accumulated.split_off(start); + } +} + +/// Wrap an SSE byte stream with the hold-and-release output guard. +/// `sandbox` and `policy_digest` feed the audit log line on a cut. +/// +/// No-op-cheap when the pipeline has no output stages — callers +/// should check [`GuardrailPipeline::covers`] and skip the wrap. +pub fn guard_sse_stream( + stream: BoxStream<'static, Result>, + pipeline: Arc, + dialect: StreamDialect, + sandbox: String, + policy_digest: String, +) -> BoxStream<'static, Result> +where + E: Send + 'static, +{ + let state = SseGuardState::new(pipeline, dialect, stream_scan_threshold()); + + struct Ctx { + inner: BoxStream<'static, Result>, + state: SseGuardState, + sandbox: String, + policy_digest: String, + /// Terminal frame queued for emission; stream ends after. + pending_cut: Option, + finished: bool, + } + + let ctx = Ctx { + inner: stream, + state, + sandbox, + policy_digest, + pending_cut: None, + finished: false, + }; + + futures::stream::unfold(ctx, |mut ctx| async move { + if let Some(frame) = ctx.pending_cut.take() { + ctx.finished = true; + return Some((Ok(frame), ctx)); + } + if ctx.finished { + return None; + } + loop { + match ctx.inner.next().await { + Some(Ok(chunk)) => match ctx.state.on_chunk(chunk).await { + SseGuardStep::Release(chunks) if chunks.is_empty() => continue, + SseGuardStep::Release(chunks) => { + let merged = merge_chunks(chunks); + return Some((Ok(merged), ctx)); + } + SseGuardStep::Cut(frame) => { + tracing::warn!( + target: "inference.audit", + sandbox = %ctx.sandbox, + inference_policy_digest = %ctx.policy_digest, + decision = "deny", + gate = "guardrail_stream", + "guardrail pipeline cut SSE stream" + ); + ctx.finished = true; + return Some((Ok(frame), ctx)); + } + }, + Some(Err(e)) => { + // Upstream transport error: surface it verbatim. + // Held chunks are dropped — their text was never + // scanned, so releasing them would violate the + // scanned-before-delivery contract. + ctx.finished = true; + return Some((Err(e), ctx)); + } + None => match ctx.state.on_end().await { + SseGuardStep::Release(chunks) => { + ctx.finished = true; + if chunks.is_empty() { + return None; + } + return Some((Ok(merge_chunks(chunks)), ctx)); + } + SseGuardStep::Cut(frame) => { + tracing::warn!( + target: "inference.audit", + sandbox = %ctx.sandbox, + inference_policy_digest = %ctx.policy_digest, + decision = "deny", + gate = "guardrail_stream", + "guardrail pipeline cut SSE stream at end-of-stream" + ); + ctx.finished = true; + return Some((Ok(frame), ctx)); + } + }, + } + } + }) + .boxed() +} + +fn merge_chunks(chunks: Vec) -> Bytes { + if chunks.len() == 1 { + return chunks.into_iter().next().expect("len checked"); + } + let total: usize = chunks.iter().map(Bytes::len).sum(); + let mut merged = Vec::with_capacity(total); + for c in chunks { + merged.extend_from_slice(&c); + } + Bytes::from(merged) +} diff --git a/inference-router/src/guardrails/tests.rs b/inference-router/src/guardrails/tests.rs new file mode 100644 index 000000000..633e5fc0b --- /dev/null +++ b/inference-router/src/guardrails/tests.rs @@ -0,0 +1,698 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Guardrail unit tests (config, moderation, extraction, pipeline, SSE). + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use bytes::Bytes; +use futures::stream::{BoxStream, StreamExt}; + +use super::backend::{BuiltStage, scan_windows}; +use super::stream::{SseGuardState, SseGuardStep, delta_text_from_event}; +use super::*; + +// ---- config parsing ---- + +#[test] +fn apply_to_parses_liberally_and_widens_unknowns() { + assert_eq!(ApplyTo::parse(Some("input")), ApplyTo::Input); + assert_eq!(ApplyTo::parse(Some("OUTPUT")), ApplyTo::Output); + assert_eq!(ApplyTo::parse(Some("both")), ApplyTo::Both); + assert_eq!(ApplyTo::parse(None), ApplyTo::Both); + assert_eq!(ApplyTo::parse(Some("sideways")), ApplyTo::Both); +} + +#[test] +fn apply_to_covers_directions() { + assert!(ApplyTo::Both.covers(Direction::Input)); + assert!(ApplyTo::Both.covers(Direction::Output)); + assert!(ApplyTo::Input.covers(Direction::Input)); + assert!(!ApplyTo::Input.covers(Direction::Output)); + assert!(ApplyTo::Output.covers(Direction::Output)); + assert!(!ApplyTo::Output.covers(Direction::Input)); +} + +#[test] +fn stage_cfg_parses_compiled_json() { + let v = serde_json::json!([ + { "provider": "openai-moderation", "applyTo": "output" }, + { "provider": "openai-moderation", "applyTo": null }, + ]); + let stages = GuardrailStageCfg::from_compiled_json(&v).expect("valid stages"); + assert_eq!(stages.len(), 2); + assert_eq!(stages[0].provider, "openai-moderation"); + assert_eq!(stages[0].apply_to, ApplyTo::Output); + assert_eq!(stages[1].apply_to, ApplyTo::Both); +} + +#[test] +fn stage_cfg_null_is_no_pipeline() { + assert!( + GuardrailStageCfg::from_compiled_json(&serde_json::Value::Null) + .expect("null ⇒ ok") + .is_empty() + ); +} + +#[test] +fn stage_cfg_malformed_fails_closed() { + // A declared-but-malformed block must be an error, never a + // silent drop to "no guardrails". + // Not an array: + assert!(GuardrailStageCfg::from_compiled_json(&serde_json::json!({})).is_err()); + assert!( + GuardrailStageCfg::from_compiled_json(&serde_json::json!("openai-moderation")).is_err() + ); + // An entry missing a string provider: + let missing = serde_json::json!([ + { "provider": "openai-moderation", "applyTo": "output" }, + { "applyTo": "input" } + ]); + assert!(GuardrailStageCfg::from_compiled_json(&missing).is_err()); +} + +// ---- moderation response parsing ---- + +#[test] +fn moderation_parse_flags_and_categories() { + let body = serde_json::json!({ + "results": [{ + "flagged": true, + "categories": { "violence": true, "hate": false, "self-harm": true } + }] + }); + let v = parse_moderation_response(&body).unwrap(); + assert!(v.flagged); + let mut cats = v.categories.clone(); + cats.sort(); + assert_eq!(cats, vec!["self-harm", "violence"]); +} + +#[test] +fn moderation_parse_pass() { + let body = serde_json::json!({ "results": [{ "flagged": false, "categories": {} }] }); + let v = parse_moderation_response(&body).unwrap(); + assert!(!v.flagged); + assert!(v.categories.is_empty()); +} + +#[test] +fn moderation_parse_fails_closed_on_malformed() { + assert!(parse_moderation_response(&serde_json::json!({})).is_err()); + assert!(parse_moderation_response(&serde_json::json!({ "results": [] })).is_err()); + assert!( + parse_moderation_response(&serde_json::json!({ "results": [{ "categories": {} }] })) + .is_err() + ); +} + +// ---- text extraction ---- + +#[test] +fn scan_text_or_raw_extracts_json_and_falls_back_to_raw() { + let json = br#"{"choices":[{"message":{"content":"answer"}}]}"#; + assert_eq!(scan_text_or_raw(json, extract_openai_output_text), "answer"); + let not_json = b"plain text that failed to parse"; + assert_eq!( + scan_text_or_raw(not_json, extract_openai_output_text), + "plain text that failed to parse" + ); +} + +#[test] +fn openai_input_text_handles_string_and_parts() { + let body = serde_json::json!({ + "messages": [ + { "role": "system", "content": "be nice" }, + { "role": "user", "content": [ { "type": "text", "text": "hello" }, + { "type": "image_url", "image_url": {} } ] } + ] + }); + assert_eq!(extract_openai_input_text(&body), "be nice\nhello"); +} + +#[test] +fn anthropic_input_text_handles_system_and_tool_results() { + let body = serde_json::json!({ + "system": "be nice", + "messages": [ + { "role": "user", "content": [ + { "type": "text", "text": "hello" }, + { "type": "tool_result", "content": "result text" } + ]}, + { "role": "assistant", "content": "earlier reply" } + ] + }); + assert_eq!( + extract_anthropic_input_text(&body), + "be nice\nhello\nresult text\nearlier reply" + ); +} + +#[test] +fn output_text_extractors() { + let openai = serde_json::json!({ + "choices": [ { "message": { "content": "answer" } } ] + }); + assert_eq!(extract_openai_output_text(&openai), "answer"); + let anthropic = serde_json::json!({ + "content": [ { "type": "text", "text": "answer" } ] + }); + assert_eq!(extract_anthropic_output_text(&anthropic), "answer"); +} + +#[test] +fn delta_extraction_per_dialect() { + let openai = serde_json::json!({ + "choices": [ { "delta": { "content": "hi" } } ] + }); + assert_eq!( + delta_text_from_event(StreamDialect::OpenAiChat, &openai), + Some("hi".to_string()) + ); + let anthropic = serde_json::json!({ + "type": "content_block_delta", + "delta": { "type": "text_delta", "text": "hi" } + }); + assert_eq!( + delta_text_from_event(StreamDialect::AnthropicMessages, &anthropic), + Some("hi".to_string()) + ); + let other = serde_json::json!({ "type": "message_start" }); + assert_eq!( + delta_text_from_event(StreamDialect::AnthropicMessages, &other), + None + ); +} + +struct MarkerGuard { + marker: &'static str, + scans: Arc, + fail: bool, +} + +#[async_trait] +impl Guardrail for MarkerGuard { + fn name(&self) -> &'static str { + "marker-test" + } + async fn scan(&self, text: &str) -> Result { + self.scans.fetch_add(1, Ordering::SeqCst); + if self.fail { + return Err(GuardrailError::Unavailable { + provider: "marker-test", + reason: "boom".into(), + }); + } + Ok(GuardrailVerdict { + flagged: text.contains(self.marker), + categories: vec!["marker".into()], + }) + } +} + +fn pipeline_with( + marker: &'static str, + apply_to: ApplyTo, + scans: Arc, + fail: bool, +) -> GuardrailPipeline { + GuardrailPipeline { + stages: vec![BuiltStage { + apply_to, + guard: Arc::new(MarkerGuard { + marker, + scans, + fail, + }), + }], + } +} + +#[tokio::test] +async fn pipeline_scan_flags_and_passes() { + let scans = Arc::new(AtomicUsize::new(0)); + let p = pipeline_with("BAD", ApplyTo::Both, scans.clone(), false); + assert!( + p.scan("all good", Direction::Input) + .await + .unwrap() + .is_none() + ); + let v = p + .scan("some BAD text", Direction::Output) + .await + .unwrap() + .expect("flagged"); + assert_eq!(v.provider, "marker-test"); + assert_eq!(v.direction, Direction::Output); + assert_eq!(scans.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn pipeline_skips_direction_not_covered() { + let scans = Arc::new(AtomicUsize::new(0)); + let p = pipeline_with("BAD", ApplyTo::Output, scans.clone(), false); + assert!(!p.covers(Direction::Input)); + assert!( + p.scan("BAD input", Direction::Input) + .await + .unwrap() + .is_none() + ); + assert_eq!(scans.load(Ordering::SeqCst), 0, "input stage must not run"); +} + +#[tokio::test] +async fn pipeline_scan_error_fails_closed() { + let scans = Arc::new(AtomicUsize::new(0)); + let p = pipeline_with("BAD", ApplyTo::Both, scans, true); + let err = p.scan("anything", Direction::Output).await.unwrap_err(); + assert_eq!(err.code(), "guardrail_unavailable"); +} + +#[test] +fn pipeline_from_stages_rejects_unknown_backend() { + let cfg = crate::config::Config::from_env().expect("env config"); + let stages = vec![GuardrailStageCfg { + provider: "not-a-backend".into(), + apply_to: ApplyTo::Both, + }]; + let err = GuardrailPipeline::from_stages(&stages, &cfg, &reqwest::Client::new()) + .err() + .expect("must fail"); + assert_eq!(err.code(), "guardrail_misconfigured"); +} + +fn sse_chunk(text: &str) -> Bytes { + Bytes::from(format!( + "data: {}\n\n", + serde_json::json!({ "choices": [ { "delta": { "content": text } } ] }) + )) +} + +fn collect_stream( + stream: BoxStream<'static, Result>, +) -> impl std::future::Future { + use futures::TryStreamExt; + async move { + let all: Vec = stream.try_collect().await.expect("stream ok"); + all.iter() + .map(|b| String::from_utf8_lossy(b).into_owned()) + .collect() + } +} + +#[tokio::test] +async fn sse_guard_releases_clean_stream_intact() { + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with("BAD", ApplyTo::Output, scans.clone(), false)); + let chunks: Vec> = vec![ + Ok(sse_chunk("hello ")), + Ok(sse_chunk("world")), + Ok(Bytes::from("data: [DONE]\n\n")), + ]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!(out.contains("hello ")); + assert!(out.contains("world")); + assert!(out.contains("[DONE]")); + // Under-threshold text ⇒ exactly one end-of-stream scan. + assert_eq!(scans.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +async fn sse_guard_cuts_stream_on_violation_and_withholds_text() { + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with("BAD", ApplyTo::Output, scans, false)); + let chunks: Vec> = vec![ + Ok(sse_chunk("this is BAD content")), + Ok(sse_chunk("more text that must never be seen")), + ]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!( + !out.contains("BAD content"), + "flagged text must never reach the client: {out}" + ); + assert!(out.contains("guardrail_blocked")); + assert!(out.contains("data: [DONE]")); +} + +#[tokio::test] +async fn sse_guard_holds_text_until_scanned_across_threshold() { + // Force a tiny threshold via a long first chunk: text length + // over the default threshold triggers a mid-stream scan. + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with("BAD", ApplyTo::Output, scans.clone(), false)); + let big = "x".repeat(STREAM_SCAN_THRESHOLD_CHARS + 10); + let chunks: Vec> = + vec![Ok(sse_chunk(&big)), Ok(sse_chunk("tail"))]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!(out.contains(&big)); + assert!(out.contains("tail")); + // One mid-stream scan (threshold) + one at end-of-stream for + // the tail. + assert_eq!(scans.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn sse_guard_cuts_on_scan_error() { + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with("BAD", ApplyTo::Output, scans, true)); + let chunks: Vec> = vec![Ok(sse_chunk("hello"))]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!(!out.contains("hello"), "unscanned text must be withheld"); + assert!(out.contains("guardrail_unavailable")); +} + +#[tokio::test] +async fn sse_guard_passes_non_text_frames_through_untouched() { + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with("BAD", ApplyTo::Output, scans.clone(), false)); + let chunks: Vec> = vec![ + Ok(Bytes::from(": keepalive\n\n")), + Ok(Bytes::from("data: [DONE]\n\n")), + ]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!(out.contains(": keepalive")); + assert!(out.contains("[DONE]")); + assert_eq!( + scans.load(Ordering::SeqCst), + 0, + "no text ⇒ no scan round-trips" + ); +} + +#[tokio::test] +async fn sse_guard_catches_data_prefix_without_space() { + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with("FORBIDDEN", ApplyTo::Output, scans, false)); + let event = serde_json::json!({ "choices": [ { "delta": { "content": "FORBIDDEN text" } } ] }); + let chunks: Vec> = + vec![Ok(Bytes::from(format!("data:{event}\n\n")))]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!( + !out.contains("FORBIDDEN"), + "spaceless data: events must still be scanned: {out}" + ); + assert!(out.contains("guardrail_blocked")); +} + +#[tokio::test] +async fn sse_guard_state_bounds_scan_context_on_long_streams() { + // Regression: `accumulated` must not grow without bound on + // long-lived streams — after every clean scan the retained + // context is trimmed to MAX_SCAN_CHARS (the most a future + // scan can consume anyway). + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with("BAD", ApplyTo::Output, scans.clone(), false)); + let mut state = SseGuardState::new(p, StreamDialect::OpenAiChat, 10); + for _ in 0..40 { + let step = state.on_chunk(sse_chunk(&"y".repeat(1000))).await; + assert!(matches!(step, SseGuardStep::Release(_))); + } + assert!( + scans.load(Ordering::SeqCst) >= 40, + "every chunk over threshold scans" + ); + assert!( + state.accumulated.chars().count() <= MAX_SCAN_CHARS, + "scan context must stay bounded, got {}", + state.accumulated.chars().count() + ); +} + +#[tokio::test] +async fn sse_guard_handles_events_split_across_chunks() { + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with("FORBIDDEN", ApplyTo::Output, scans, false)); + let full = sse_chunk("this is FORBIDDEN text"); + let (a, b) = full.split_at(20); + let chunks: Vec> = + vec![Ok(Bytes::copy_from_slice(a)), Ok(Bytes::copy_from_slice(b))]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!( + !out.contains("FORBIDDEN"), + "split-event text must still be caught: {out}" + ); + assert!(out.contains("guardrail_blocked")); +} + +#[tokio::test] +async fn sse_guard_scans_non_json_data_frames() { + // Upstream drift: a `data:` line whose payload isn't valid + // JSON must still be scanned, not fast-released unscanned. + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with( + "FORBIDDEN", + ApplyTo::Output, + scans.clone(), + false, + )); + let chunks: Vec> = + vec![Ok(Bytes::from("data: this is FORBIDDEN not-json\n\n"))]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!( + !out.contains("FORBIDDEN not-json"), + "non-JSON data frame must be scanned, not leaked: {out}" + ); + assert!(out.contains("guardrail_blocked")); + assert!( + scans.load(Ordering::SeqCst) >= 1, + "raw frame must be scanned" + ); +} + +#[tokio::test] +async fn sse_guard_passes_json_structural_frames_without_scanning() { + // Valid-JSON frames with no delta text (ping / role-only / + // stop) carry no model text and are released without a scan + // round-trip. + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with( + "FORBIDDEN", + ApplyTo::Output, + scans.clone(), + false, + )); + let chunks: Vec> = vec![ + Ok(Bytes::from( + "data: {\"choices\":[{\"delta\":{\"role\":\"assistant\"}}]}\n\n", + )), + Ok(Bytes::from("data: [DONE]\n\n")), + ]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!(out.contains("\"role\":\"assistant\"")); + assert!(out.contains("[DONE]")); + assert_eq!( + scans.load(Ordering::SeqCst), + 0, + "structural frames don't scan" + ); +} + +#[tokio::test] +async fn sse_guard_withholds_bytes_when_split_inside_content_string() { + // Regression (B1): a chunk boundary *inside* the JSON content + // string leaves the delta text uncounted in line_carry while + // its bytes sit in `held`. The guard must not fast-release + // those bytes, or flagged model text ships unscanned. + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with("FORBIDDEN", ApplyTo::Output, scans, false)); + let full = sse_chunk("this is FORBIDDEN text"); + let s = String::from_utf8(full.to_vec()).unwrap(); + let cut = s.find("FORB").unwrap() + 2; // split mid-word, mid-string + let (a, b) = s.split_at(cut); + let chunks: Vec> = + vec![Ok(Bytes::from(a.to_owned())), Ok(Bytes::from(b.to_owned()))]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!( + !out.contains("FORB"), + "no partial content bytes may reach the client unscanned: {out}" + ); + assert!(out.contains("guardrail_blocked")); +} + +#[test] +fn scan_windows_covers_every_char_without_truncation() { + let short = "hello"; + assert_eq!(scan_windows(short), vec!["hello"]); + + let long: String = "a".repeat(MAX_SCAN_CHARS) + &"b".repeat(500); + let windows = scan_windows(&long); + assert_eq!(windows.len(), 2); + assert_eq!(windows[0].chars().count(), MAX_SCAN_CHARS); + assert_eq!(windows[1].chars().count(), 500); + assert_eq!(windows.concat(), long, "no char dropped across windows"); +} + +#[tokio::test] +async fn pipeline_scans_past_the_cap_via_windows() { + // Content hidden past MAX_SCAN_CHARS must still be caught — + // padding can't push it out of a truncated window anymore. + let scans = Arc::new(AtomicUsize::new(0)); + let p = pipeline_with("NEEDLE", ApplyTo::Input, scans.clone(), false); + let text = "x".repeat(MAX_SCAN_CHARS + 100) + " NEEDLE"; + let v = p + .scan(&text, Direction::Input) + .await + .unwrap() + .expect("needle past the cap is still flagged"); + assert_eq!(v.provider, "marker-test"); + assert!(scans.load(Ordering::SeqCst) >= 2, "must scan >1 window"); +} + +#[test] +fn delta_text_scans_every_choice_not_just_first() { + // A multi-choice frame must surface all choices' text, or a + // later choice ships unscanned. + let event = serde_json::json!({ + "choices": [ + { "delta": { "content": "safe intro " } }, + { "delta": { "content": "FORBIDDEN payload" } } + ] + }); + let text = delta_text_from_event(StreamDialect::OpenAiChat, &event) + .expect("multi-choice frame has text"); + assert!(text.contains("safe intro"), "first choice: {text}"); + assert!(text.contains("FORBIDDEN payload"), "second choice: {text}"); +} + +#[tokio::test] +async fn sse_guard_scans_utf8_split_across_chunks() { + // Regression: a multi-byte code point split across chunk + // boundaries must be reconstructed before scanning. With the + // old per-chunk lossy decode, the split char became U+FFFD, so + // a flagged multi-byte phrase evaded moderation while the + // client received the intact bytes. The marker is CJK so the + // scan only flags if the code points survived the split. + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with("爆弾", ApplyTo::Output, scans, false)); + let full = sse_chunk("plan: 爆弾 now"); + let bytes = full.to_vec(); + let s = std::str::from_utf8(&bytes).unwrap(); + // Split one byte into the 3-byte '爆' code point. + let split = s.find('爆').unwrap() + 1; + let (a, b) = bytes.split_at(split); + let chunks: Vec> = vec![ + Ok(Bytes::copy_from_slice(a)), + Ok(Bytes::copy_from_slice(b)), + Ok(Bytes::from("data: [DONE]\n\n")), + ]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!( + out.contains("guardrail_blocked"), + "reconstructed multi-byte text must be flagged, got: {out}" + ); + assert!( + !out.contains('\u{FFFD}'), + "no replacement char should appear from a split code point: {out}" + ); +} + +#[tokio::test] +async fn sse_guard_passes_clean_utf8_split_intact() { + // The benign counterpart: a clean stream whose multi-byte char + // is split must still be delivered byte-for-byte (no U+FFFD, + // no drops) once reconstructed. + let scans = Arc::new(AtomicUsize::new(0)); + let p = Arc::new(pipeline_with("NOPE", ApplyTo::Output, scans, false)); + let full = sse_chunk("café ☕ ready"); + let bytes = full.to_vec(); + let s = std::str::from_utf8(&bytes).unwrap(); + let split = s.find('☕').unwrap() + 1; // mid 3-byte code point + let (a, b) = bytes.split_at(split); + let chunks: Vec> = vec![ + Ok(Bytes::copy_from_slice(a)), + Ok(Bytes::copy_from_slice(b)), + Ok(Bytes::from("data: [DONE]\n\n")), + ]; + let guarded = guard_sse_stream( + futures::stream::iter(chunks).boxed(), + p, + StreamDialect::OpenAiChat, + "sbx".into(), + "sha256:t".into(), + ); + let out = collect_stream(guarded).await; + assert!(out.contains("café ☕ ready"), "intact delivery: {out}"); + assert!(!out.contains('\u{FFFD}'), "no corruption: {out}"); +} diff --git a/inference-router/src/inference_policy_loader.rs b/inference-router/src/inference_policy_loader.rs index 30defe06c..bf7b1efbb 100644 --- a/inference-router/src/inference_policy_loader.rs +++ b/inference-router/src/inference_policy_loader.rs @@ -165,6 +165,14 @@ pub struct LoadedInferencePolicy { /// back to the env-driven default deployment (back-compat). pub model_preference: Option, + /// `spec.provider` — raw kebab-case tag, resolved per request by + /// [`crate::provider::resolve`]. `None` ⇒ Azure (back-compat). + pub provider: Option, + + /// `spec.guardrails[]` — pipeline stages; validity is checked at + /// build time in [`crate::guardrails::GuardrailPipeline::from_stages`]. + pub guardrails: Vec, + /// Whole profile JSON, kept so subsequent sub-slices can pick up /// other axes without a new loader. pub raw: serde_json::Value, @@ -327,6 +335,39 @@ pub fn load_inference_policy_from_dir( .unwrap_or(&serde_json::Value::Null), ); + // Multi-provider slice: raw provider tag + guardrail stages. + // Both parse liberally here (defence-in-depth: never crash the + // data plane on schema drift); enforcement-relevant strictness + // lives at the per-request consumption sites. + let provider = parsed + .get("provider") + .and_then(|p| p.as_str()) + .filter(|p| !p.trim().is_empty()) + .map(str::to_string); + // `null`/absent ⇒ no pipeline. A present-but-malformed block does + // NOT degrade to "no guardrails" (that would silently disable both + // the scan and the sibling route-gap guard). Instead poison the + // policy with a sentinel stage: it is non-empty (so the route-gap + // guard fails closed) and names an unbuildable backend (so + // `GuardrailPipeline::from_stages` returns a config error and every + // request fails closed with 503 guardrail_misconfigured). + let guardrails = match crate::guardrails::GuardrailStageCfg::from_compiled_json( + parsed.get("guardrails").unwrap_or(&serde_json::Value::Null), + ) { + Ok(stages) => stages, + Err(reason) => { + tracing::error!( + file = %file.display(), + %reason, + "InferencePolicy guardrails malformed, failing closed for every request" + ); + vec![crate::guardrails::GuardrailStageCfg { + provider: format!("malformed-guardrail-config ({reason})"), + apply_to: crate::guardrails::ApplyTo::Both, + }] + } + }; + // Digest layout matches controller `inference_policy_digest`: // length-prefixed (name, body) hashed with sha256. let canonical = canonical_bytes_for_digest(INFERENCE_POLICY_FILENAME, &body); @@ -342,6 +383,8 @@ pub fn load_inference_policy_from_dir( daily_tokens = ?daily_tokens, monthly_tokens = ?monthly_tokens, content_safety_active = content_safety.is_active(), + provider = ?provider, + guardrail_stages = guardrails.len(), primary_deployment = ?model_preference.as_ref().map(|m| m.primary.deployment.as_str()), fallback_count = model_preference.as_ref().map(|m| m.fallback.len()).unwrap_or(0), // Surface the actual fallback chain (not just the count) so ops @@ -383,6 +426,8 @@ pub fn load_inference_policy_from_dir( monthly_tokens, content_safety, model_preference, + provider, + guardrails, raw: parsed, }) } @@ -507,6 +552,13 @@ pub struct InferencePolicySnapshot { /// `primary.deployment`; the `fallback` chain is captured for /// Slice 2d.2's health-aware failover. pub model_preference: Option, + /// `spec.provider` — raw tag consumed by + /// [`crate::routes::apply_provider_resolution`] per request. + pub provider: Option, + /// `spec.guardrails[]` — consumed by + /// [`crate::guardrails::GuardrailPipeline::from_stages`] per + /// request. Empty ⇒ no pipeline. + pub guardrails: Vec, } /// Take a single read-lock snapshot of the currently-loaded policy. @@ -529,6 +581,8 @@ pub async fn current_snapshot(handle: &LoadedInferencePolicyHandle) -> Inference monthly_tokens: p.monthly_tokens, content_safety: p.content_safety.clone(), model_preference: p.model_preference.clone(), + provider: p.provider.clone(), + guardrails: p.guardrails.clone(), }) .unwrap_or_default() } @@ -743,6 +797,11 @@ mod tests { deployment: "gpt-5.4-us".into(), }], }), + provider: Some("anthropic".into()), + guardrails: vec![crate::guardrails::GuardrailStageCfg { + provider: "openai-moderation".into(), + apply_to: crate::guardrails::ApplyTo::Both, + }], raw: serde_json::Value::Null, }; let handle: LoadedInferencePolicyHandle = @@ -761,6 +820,73 @@ mod tests { assert_eq!(mp.primary.deployment, "gpt-5.4-eu"); assert_eq!(mp.fallback.len(), 1); assert_eq!(mp.fallback[0].deployment, "gpt-5.4-us"); + assert_eq!(snap.provider.as_deref(), Some("anthropic")); + assert_eq!(snap.guardrails.len(), 1); + assert_eq!(snap.guardrails[0].provider, "openai-moderation"); + } + + #[test] + fn loads_provider_and_guardrails_when_present() { + // Multi-provider slice: compiled JSON carries `provider` + + // `guardrails`; the loader lifts both onto + // `LoadedInferencePolicy` verbatim (interpretation is + // per-request). + let tmp = TempDir::new().unwrap(); + let profile = serde_json::json!({ + "appliesTo": { "sandboxName": "agent-x", "sandboxMatchLabels": {}, "action": null }, + "tokenBudget": null, + "contentSafety": null, + "modelPreference": null, + "provider": "ollama", + "guardrails": [ + { "provider": "openai-moderation", "applyTo": "input" }, + { "provider": "openai-moderation", "applyTo": null } + ], + "displayName": null + }); + write_profile(tmp.path(), INFERENCE_POLICY_FILENAME, &profile); + + let reg = registry(); + let outcome = load_inference_policy_from_dir(tmp.path().to_str().unwrap(), ®); + let loaded = match outcome { + LoadOutcome::Loaded(p) => p, + other => panic!("expected Loaded, got {other:?}"), + }; + assert_eq!(loaded.provider.as_deref(), Some("ollama")); + assert_eq!(loaded.guardrails.len(), 2); + assert_eq!(loaded.guardrails[0].provider, "openai-moderation"); + assert_eq!( + loaded.guardrails[0].apply_to, + crate::guardrails::ApplyTo::Input + ); + assert_eq!( + loaded.guardrails[1].apply_to, + crate::guardrails::ApplyTo::Both + ); + } + + #[test] + fn absent_provider_and_guardrails_keep_backcompat_defaults() { + // Pre-slice compiled profiles (no `provider` / `guardrails` + // keys at all) must load exactly as before. + let tmp = TempDir::new().unwrap(); + let profile = serde_json::json!({ + "appliesTo": { "sandboxName": null, "sandboxMatchLabels": {}, "action": null }, + "tokenBudget": { "perRequestTokens": 1024 }, + "contentSafety": null, + "modelPreference": null, + "displayName": null + }); + write_profile(tmp.path(), INFERENCE_POLICY_FILENAME, &profile); + + let reg = registry(); + let outcome = load_inference_policy_from_dir(tmp.path().to_str().unwrap(), ®); + let loaded = match outcome { + LoadOutcome::Loaded(p) => p, + other => panic!("expected Loaded, got {other:?}"), + }; + assert!(loaded.provider.is_none()); + assert!(loaded.guardrails.is_empty()); } #[test] diff --git a/inference-router/src/lib.rs b/inference-router/src/lib.rs index ca4ac9cd9..a8cc38a83 100644 --- a/inference-router/src/lib.rs +++ b/inference-router/src/lib.rs @@ -32,6 +32,7 @@ pub mod errors; pub mod failover; pub mod forward_proxy; pub mod governance; +pub mod guardrails; pub mod handoff; pub mod inference_policy_loader; pub mod mcp; @@ -40,6 +41,7 @@ pub mod mesh; pub mod metrics; pub mod policy_envelope; pub mod policy_status; +pub mod provider; pub mod providers; pub mod proxy; pub mod rate_limiter; diff --git a/inference-router/src/metrics.rs b/inference-router/src/metrics.rs index 3ce7895c2..faeb00f53 100644 --- a/inference-router/src/metrics.rs +++ b/inference-router/src/metrics.rs @@ -38,6 +38,19 @@ pub static TOKENS_USED: LazyLock = LazyLock::new(|| { .unwrap() }); +/// Guardrail pipeline scans by backend, direction, and outcome +/// (`pass` | `flagged` | `error`). +pub static GUARDRAIL_SCANS: LazyLock = LazyLock::new(|| { + register_int_counter_vec!( + opts!( + "kars_guardrail_scans_total", + "Guardrail pipeline scans by backend, direction, and outcome" + ), + &["provider", "direction", "outcome"] + ) + .unwrap() +}); + // ── AGT Governance metrics ────────────────────────────────────────────────── /// Total AGT policy evaluations by decision (allow, deny, requires_approval, rate_limited). diff --git a/inference-router/src/provider.rs b/inference-router/src/provider.rs new file mode 100644 index 000000000..06ea713df --- /dev/null +++ b/inference-router/src/provider.rs @@ -0,0 +1,300 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Multi-provider upstream resolution. +//! +//! Maps `InferencePolicy.spec.provider` onto a concrete upstream +//! target (base URL + auth scheme). Credentials/endpoints come only +//! from the router's env / secret mounts — never from the agent. +//! +//! `spec.provider` is the sole routing selector; the pre-existing +//! `modelPreference.primary.provider` tag stays informational, so +//! adding routing doesn't reroute CRs that only set a model +//! preference. Absent/empty or unrecognised ⇒ Azure. `bedrock` +//! (unimplemented) and a provider with missing config both fail the +//! request closed rather than silently falling back to Azure. + +use crate::config::Config; + +/// Providers the router can actually forward to today. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum ProviderKind { + /// Azure OpenAI / Foundry (also GitHub Models + Copilot via + /// endpoint detection) — the Phase 1 substrate. Default. + #[default] + AzureOpenAI, + /// Anthropic Messages API — native pass-through on + /// `/v1/messages`; auth via `x-api-key` from the router-side + /// secret. + Anthropic, + /// OpenAI-compatible Ollama server — pass-through on + /// `/v1/chat/completions`; no auth. + Ollama, +} + +impl ProviderKind { + /// Kebab-case wire tag, matching the controller-side + /// `InferenceProvider::as_tag`. + #[must_use] + pub fn as_tag(&self) -> &'static str { + match self { + Self::AzureOpenAI => "azure-openai", + Self::Anthropic => "anthropic", + Self::Ollama => "ollama", + } + } +} + +/// Why a provider tag could not be turned into a forwardable upstream. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum ProviderError { + /// Tag is recognised by the CRD schema but the router has no + /// client for it yet (`bedrock`). HTTP mapping: 501. + #[error("provider '{tag}' is not implemented by this router build")] + Unimplemented { tag: String }, + /// Provider needs an endpoint the router was not configured with. + /// HTTP mapping: 503 (operator config gap, not a caller bug). + #[error( + "provider '{provider}' selected by InferencePolicy but {env} is not configured on the router" + )] + MissingEndpoint { + provider: &'static str, + env: &'static str, + }, + /// Provider needs a credential the router was not configured + /// with. HTTP mapping: 503. + #[error( + "provider '{provider}' selected by InferencePolicy but no credential is configured ({env} or secret mount)" + )] + MissingCredential { + provider: &'static str, + env: &'static str, + }, +} + +/// Parse a policy provider tag. `Ok(None)` means "no opinion" (empty +/// or unknown tag — logged by the caller, keeps the pre-slice +/// informational-only behaviour for tags like `gemini`). +/// `Err(Unimplemented)` is reserved for tags the CRD schema accepts +/// but the router cannot serve, so declared intent fails loudly. +pub fn parse_tag(tag: &str) -> Result, ProviderError> { + match tag.trim().to_ascii_lowercase().as_str() { + "azure-openai" => Ok(Some(ProviderKind::AzureOpenAI)), + "anthropic" => Ok(Some(ProviderKind::Anthropic)), + "ollama" => Ok(Some(ProviderKind::Ollama)), + "bedrock" => Err(ProviderError::Unimplemented { + tag: "bedrock".into(), + }), + _ => Ok(None), + } +} + +/// Concrete upstream target after resolution. For non-Azure providers +/// this carries the endpoint (and credential) the proxy layer needs; +/// `AzureOpenAI` keeps the env-driven endpoint already present on +/// `UpstreamConfig`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ProviderTarget { + AzureOpenAI, + Anthropic { endpoint: String, api_key: String }, + Ollama { endpoint: String }, +} + +impl ProviderTarget { + #[must_use] + pub fn kind(&self) -> ProviderKind { + match self { + Self::AzureOpenAI => ProviderKind::AzureOpenAI, + Self::Anthropic { .. } => ProviderKind::Anthropic, + Self::Ollama { .. } => ProviderKind::Ollama, + } + } +} + +/// Resolve the upstream target from `spec.provider` (the sole routing +/// selector; see module docs). +pub fn resolve(policy_tag: Option<&str>, config: &Config) -> Result { + let kind = effective_kind(policy_tag)?; + target_for(kind, config) +} + +fn effective_kind(policy_tag: Option<&str>) -> Result { + if let Some(tag) = policy_tag.filter(|t| !t.trim().is_empty()) { + match parse_tag(tag)? { + Some(kind) => return Ok(kind), + None => { + tracing::warn!( + tag, + "InferencePolicy spec.provider tag not recognised — using azure-openai" + ); + } + } + } + Ok(ProviderKind::AzureOpenAI) +} + +fn target_for(kind: ProviderKind, config: &Config) -> Result { + match kind { + ProviderKind::AzureOpenAI => Ok(ProviderTarget::AzureOpenAI), + ProviderKind::Anthropic => { + let api_key = + config + .anthropic_api_key + .clone() + .ok_or(ProviderError::MissingCredential { + provider: "anthropic", + env: "ANTHROPIC_API_KEY", + })?; + Ok(ProviderTarget::Anthropic { + endpoint: config.anthropic_endpoint.clone(), + api_key, + }) + } + ProviderKind::Ollama => { + let endpoint = + config + .ollama_endpoint + .clone() + .ok_or(ProviderError::MissingEndpoint { + provider: "ollama", + env: "OLLAMA_ENDPOINT", + })?; + Ok(ProviderTarget::Ollama { endpoint }) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{Config, RegistryMode}; + + fn cfg(anthropic_key: Option<&str>, ollama: Option<&str>) -> Config { + Config { + port: 8443, + foundry_endpoint: None, + foundry_project_endpoint: None, + azure_openai_endpoint: Some("https://contoso.openai.azure.com".into()), + default_model: "gpt-4o-mini".into(), + content_safety_enabled: false, + prompt_shields_enabled: false, + content_safety_endpoint: None, + token_budget_daily: 0, + token_budget_per_request: 0, + registry_mode: RegistryMode::Local, + registry_url: None, + provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: anthropic_key.map(String::from), + ollama_endpoint: ollama.map(String::from), + openai_moderation_endpoint: "https://api.openai.com".into(), + openai_moderation_api_key: None, + openai_moderation_model: "omni-moderation-latest".into(), + } + } + + #[test] + fn parse_recognises_supported_tags_case_insensitively() { + assert_eq!( + parse_tag("azure-openai").unwrap(), + Some(ProviderKind::AzureOpenAI) + ); + assert_eq!( + parse_tag("Anthropic").unwrap(), + Some(ProviderKind::Anthropic) + ); + assert_eq!(parse_tag(" ollama ").unwrap(), Some(ProviderKind::Ollama)); + } + + #[test] + fn parse_returns_none_for_unknown_tags() { + assert_eq!(parse_tag("gemini").unwrap(), None); + assert_eq!(parse_tag("").unwrap(), None); + assert_eq!(parse_tag("Foundry").unwrap(), None); + } + + #[test] + fn parse_rejects_bedrock_as_unimplemented() { + assert!(matches!( + parse_tag("bedrock"), + Err(ProviderError::Unimplemented { .. }) + )); + } + + #[test] + fn absent_provider_resolves_to_azure() { + assert_eq!( + resolve(None, &cfg(None, None)).unwrap(), + ProviderTarget::AzureOpenAI + ); + assert_eq!( + resolve(Some(""), &cfg(None, None)).unwrap(), + ProviderTarget::AzureOpenAI + ); + } + + #[test] + fn unknown_provider_tag_falls_back_to_azure() { + // A documented-but-unrouted tag like gemini must not error — + // it stays informational, request goes to Azure. + assert_eq!( + resolve(Some("gemini"), &cfg(Some("sk-x"), None)).unwrap(), + ProviderTarget::AzureOpenAI + ); + } + + #[test] + fn anthropic_provider_resolves_with_key() { + assert_eq!( + resolve(Some("anthropic"), &cfg(Some("sk-x"), None)).unwrap(), + ProviderTarget::Anthropic { + endpoint: "https://api.anthropic.com".into(), + api_key: "sk-x".into(), + } + ); + } + + #[test] + fn anthropic_without_key_fails_closed() { + assert!(matches!( + resolve(Some("anthropic"), &cfg(None, None)), + Err(ProviderError::MissingCredential { + provider: "anthropic", + .. + }) + )); + } + + #[test] + fn ollama_without_endpoint_fails_closed() { + assert!(matches!( + resolve(Some("ollama"), &cfg(None, None)), + Err(ProviderError::MissingEndpoint { + provider: "ollama", + .. + }) + )); + } + + #[test] + fn ollama_with_endpoint_resolves() { + assert_eq!( + resolve( + Some("ollama"), + &cfg(None, Some("http://ollama.ollama.svc:11434")) + ) + .unwrap(), + ProviderTarget::Ollama { + endpoint: "http://ollama.ollama.svc:11434".into() + } + ); + } + + #[test] + fn bedrock_is_unimplemented_not_silent() { + assert!(matches!( + resolve(Some("bedrock"), &cfg(None, None)), + Err(ProviderError::Unimplemented { .. }) + )); + } +} diff --git a/inference-router/src/proxy.rs b/inference-router/src/proxy.rs index 87da56894..5d79b145e 100644 --- a/inference-router/src/proxy.rs +++ b/inference-router/src/proxy.rs @@ -18,6 +18,7 @@ use crate::copilot_auth::{ self, COPILOT_INTEGRATION_ID, CopilotTokenCache, EDITOR_PLUGIN_VERSION, EDITOR_VERSION, }; use crate::metrics; +use crate::provider::ProviderKind; use std::sync::Arc; /// Upstream configuration for a single request. @@ -26,6 +27,39 @@ pub struct UpstreamConfig { pub endpoint: String, pub deployment: String, pub sandbox_name: String, + /// Which provider family `endpoint` belongs to — drives URL shape + /// and auth scheme. `AzureOpenAI` preserves the historic + /// behaviour (incl. GitHub Models / Copilot endpoint detection). + pub provider: ProviderKind, + /// Static API key for providers that use one (`Anthropic`). + /// Filled by `provider::resolve` from router-side config only — + /// never from the inbound request. + pub api_key: Option, +} + +impl UpstreamConfig { + /// The historic constructor shape: an Azure OpenAI / Foundry + /// upstream authenticated via Workload Identity / API-key mode. + #[must_use] + pub fn azure(endpoint: String, deployment: String, sandbox_name: String) -> Self { + Self { + endpoint, + deployment, + sandbox_name, + provider: ProviderKind::AzureOpenAI, + api_key: None, + } + } +} + +/// Credential material resolved for one upstream request. Which +/// header it lands in is provider-specific (`Authorization: Bearer` +/// vs `x-api-key` vs nothing at all for unauthenticated in-cluster +/// Ollama). +pub enum UpstreamCredential { + Bearer(String), + AnthropicApiKey(String), + None, } /// Determine the correct token audience for the upstream endpoint. @@ -51,7 +85,7 @@ fn token_audience(endpoint: &str) -> &'static str { fn build_upstream_headers( request_headers: &HeaderMap, _auth: &WorkloadIdentityAuth, - token: &str, + credential: &UpstreamCredential, endpoint: &str, ) -> Result { let mut headers = HeaderMap::new(); @@ -68,13 +102,32 @@ fn build_upstream_headers( } } - // Both API-key and Entra modes use Authorization: Bearer for the unified - // /openai/v1/ endpoint format. Azure OpenAI accepts API keys as Bearer tokens. - // Copilot also uses Bearer (with the exchanged Copilot JWT). - headers.insert( - "authorization", - HeaderValue::from_str(&format!("Bearer {token}")).context("Invalid token")?, - ); + match credential { + // Both API-key and Entra modes use Authorization: Bearer for the unified + // /openai/v1/ endpoint format. Azure OpenAI accepts API keys as Bearer tokens. + // Copilot also uses Bearer (with the exchanged Copilot JWT). + UpstreamCredential::Bearer(token) => { + headers.insert( + "authorization", + HeaderValue::from_str(&format!("Bearer {token}")).context("Invalid token")?, + ); + } + // Anthropic's Messages API authenticates with `x-api-key` and + // requires an `anthropic-version` header. The inbound SDK + // value (when present) was already copied through above — + // only the default is filled in here. + UpstreamCredential::AnthropicApiKey(key) => { + headers.insert( + "x-api-key", + HeaderValue::from_str(key).context("Invalid Anthropic API key")?, + ); + headers + .entry("anthropic-version") + .or_insert(HeaderValue::from_static("2023-06-01")); + } + // Unauthenticated upstream (in-cluster Ollama). + UpstreamCredential::None => {} + } headers .entry("content-type") .or_insert(HeaderValue::from_static("application/json")); @@ -132,6 +185,38 @@ pub async fn token_for_endpoint( } } +/// 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), + } +} + /// Record Prometheus metrics from a completed request. fn record_metrics( upstream: &UpstreamConfig, @@ -194,20 +279,26 @@ pub async fn forward( let (upstream_url, body) = build_upstream_url(auth, upstream, path, request_body)?; - let mode = if is_copilot_endpoint(&upstream.endpoint) { - "copilot" - } else if auth.is_api_key_mode() { - "dev" - } else { - "foundry" + let mode = match upstream.provider { + ProviderKind::Anthropic => "anthropic", + ProviderKind::Ollama => "ollama", + ProviderKind::AzureOpenAI => { + if is_copilot_endpoint(&upstream.endpoint) { + "copilot" + } else if auth.is_api_key_mode() { + "dev" + } else { + "foundry" + } + } }; tracing::info!(sandbox = %upstream.sandbox_name, model = %upstream.deployment, mode = %mode, "Forwarding inference"); - let token = token_for_endpoint(auth, copilot, &upstream.endpoint) + let credential = credential_for_upstream(auth, copilot, upstream) .await .context("Failed to acquire auth token")?; - let headers = build_upstream_headers(request_headers, auth, &token, &upstream.endpoint)?; + let headers = build_upstream_headers(request_headers, auth, &credential, &upstream.endpoint)?; tracing::info!(sandbox = %upstream.sandbox_name, url = %upstream_url, body_len = body.len(), "Sending upstream request"); @@ -430,10 +521,10 @@ pub async fn forward_stream( tracing::info!(sandbox = %upstream.sandbox_name, model = %upstream.deployment, mode = "stream", "Forwarding SSE stream"); - let token = token_for_endpoint(&auth, copilot.as_deref(), &upstream.endpoint) + 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, &token, &upstream.endpoint)?; + let headers = build_upstream_headers(&request_headers, &auth, &credential, &upstream.endpoint)?; let start = Instant::now(); @@ -591,6 +682,10 @@ fn is_github_models_endpoint(endpoint: &str) -> bool { /// Uses the unified /openai/v1/ format — works with both API-key and Entra auth. /// /// Routing rules: +/// - Anthropic: no path rewrite — callers pass Messages-API paths +/// (`v1/messages`) verbatim. +/// - Ollama: OpenAI-compat lives under `/v1/` — `chat/completions` +/// becomes `{endpoint}/v1/chat/completions`. /// - GitHub Copilot (`api.githubcopilot.com`): no path rewrite; OpenClaw /// sends OpenAI-shape to `/chat/completions` and Anthropic-shape to /// `/v1/messages`. We forward those paths unchanged. @@ -602,20 +697,34 @@ fn build_upstream_url( path: &str, request_body: Bytes, ) -> Result<(String, Bytes)> { - let url = if is_github_models_endpoint(&upstream.endpoint) - || is_copilot_endpoint(&upstream.endpoint) - { - format!( + let url = match upstream.provider { + ProviderKind::Anthropic => format!( "{}/{}", upstream.endpoint.trim_end_matches('/'), path.trim_start_matches('/'), - ) - } else { - format!( - "{}/openai/v1/{}", + ), + ProviderKind::Ollama => format!( + "{}/v1/{}", upstream.endpoint.trim_end_matches('/'), - path.trim_start_matches('/'), - ) + path.trim_start_matches('/').trim_start_matches("v1/"), + ), + ProviderKind::AzureOpenAI => { + if is_github_models_endpoint(&upstream.endpoint) + || is_copilot_endpoint(&upstream.endpoint) + { + format!( + "{}/{}", + upstream.endpoint.trim_end_matches('/'), + path.trim_start_matches('/'), + ) + } else { + format!( + "{}/openai/v1/{}", + upstream.endpoint.trim_end_matches('/'), + path.trim_start_matches('/'), + ) + } + } }; let body = if let Ok(mut body_json) = serde_json::from_slice::(&request_body) { @@ -642,6 +751,7 @@ fn build_upstream_url( // requesting reasoning encryption on a stripped input is a no-op // 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) && let Some(obj) = body_json.as_object_mut() diff --git a/inference-router/src/routes/anthropic_messages.rs b/inference-router/src/routes/anthropic_messages.rs index 4bbc1416b..0e939f24c 100644 --- a/inference-router/src/routes/anthropic_messages.rs +++ b/inference-router/src/routes/anthropic_messages.rs @@ -25,20 +25,100 @@ use futures::stream::StreamExt; use serde_json::{Value, json}; use super::AppState; +use crate::guardrails::{self, Direction, GuardrailError, GuardrailPipeline}; +use crate::provider::{ProviderError, ProviderKind}; use crate::proxy; +use std::sync::Arc; + +/// Framing / hop-by-hop headers that must not be copied from an +/// upstream response onto a rebuilt one — hyper re-frames the body +/// itself, and a stale `transfer-encoding: chunked` (Anthropic over +/// HTTP/1.1) makes it abort the connection without a response. +fn is_hop_by_hop(name: &str) -> bool { + matches!( + name, + "connection" + | "content-length" + | "keep-alive" + | "proxy-authenticate" + | "proxy-authorization" + | "te" + | "trailer" + | "transfer-encoding" + | "upgrade" + ) +} fn deny_response(status: StatusCode, message: &str, code: &str) -> axum::response::Response { - ( + // Anthropic wire shape (`error.type`) plus an explicit `error.code` + // mirroring the OpenAI routes, so a client can switch on one stable + // `error.code` across every inference route. For non-policy errors + // (bad request, rate limit) `type` and `code` coincide. + deny_coded(status, message, code, code, false) +} + +/// Policy denial (provider/guardrail): keeps the Anthropic-native +/// `error.type` but carries the stable machine code in `error.code` +/// and attaches the `x-kars-decision*` headers, matching the +/// chat-completions contract. Distinct `error_type` and `code` are the +/// point: a client switches on `error.code` (e.g. `guardrail_blocked` +/// vs `guardrail_misconfigured` vs `guardrail_unavailable`) while +/// `error.type` stays the Anthropic-shaped value clients already read. +fn deny_policy( + status: StatusCode, + message: &str, + error_type: &str, + code: &str, +) -> axum::response::Response { + deny_coded(status, message, error_type, code, true) +} + +fn deny_coded( + status: StatusCode, + message: &str, + error_type: &str, + code: &str, + decision_headers: bool, +) -> axum::response::Response { + let mut resp = ( status, Json(json!({ "type": "error", "error": { - "type": code, + "type": error_type, + "code": code, "message": message, } })), ) - .into_response() + .into_response(); + if decision_headers { + super::chat_completions::insert_decision_headers( + &mut resp, + "blocked", + "InferencePolicy", + message, + ); + } + resp +} + +/// HTTP status for a guardrail pipeline error: 503 misconfigured +/// (declared stage cannot be built), 502 backend outage. +fn guardrail_error_status(e: &GuardrailError) -> StatusCode { + match e { + GuardrailError::Config { .. } => StatusCode::SERVICE_UNAVAILABLE, + GuardrailError::Unavailable { .. } => StatusCode::BAD_GATEWAY, + } +} + +/// Stable machine code for a provider-resolution failure, matching +/// `chat_completions::provider_error_response`. +fn provider_error_code(e: &ProviderError) -> &'static str { + match e { + ProviderError::Unimplemented { .. } => "provider_unimplemented", + _ => "provider_unconfigured", + } } /// Convert Anthropic Messages-shaped JSON to OpenAI chat-completions-shaped JSON. @@ -264,11 +344,104 @@ pub(super) async fn anthropic_messages( // Slice 2d.1: honour `InferencePolicy.modelPreference.primary.deployment`. crate::routes::apply_model_preference_override(&mut upstream, &policy); - // Copilot exposes a native Anthropic Messages endpoint at /v1/messages. - // Skip translation entirely and forward the body as-is, preserving the - // streaming + tool_use + multi-modal contracts of the Anthropic SDK. - if proxy::is_copilot_endpoint(&upstream.endpoint) { - return forward_anthropic_passthrough(state, sandbox_name, headers, body, upstream).await; + // Retarget at the policy-selected provider (fails closed). + if let Err(e) = crate::routes::apply_provider_resolution(&state, &mut upstream, &policy) { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "provider_resolution", + error = %e, + "InferencePolicy provider could not be resolved (anthropic route)" + ); + let status = match e { + ProviderError::Unimplemented { .. } => StatusCode::NOT_IMPLEMENTED, + _ => StatusCode::SERVICE_UNAVAILABLE, + }; + return deny_policy(status, &e.to_string(), "api_error", provider_error_code(&e)); + } + + // Guardrail pipeline; a declared-but-unbuildable stage blocks. + let guardrail_pipeline = + match super::chat_completions::build_guardrail_pipeline(&state, &policy) { + Ok(p) => p, + Err(e) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_config", + error = %e, + "guardrail pipeline could not be built (anthropic route) — failing closed" + ); + return deny_policy( + guardrail_error_status(&e), + &e.to_string(), + "api_error", + e.code(), + ); + } + }; + if let Some(ref p) = guardrail_pipeline + && p.covers(Direction::Input) + { + let input_text = guardrails::extract_anthropic_input_text(&req_json); + match p.scan(&input_text, Direction::Input).await { + Ok(None) => {} + Ok(Some(v)) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_input", + categories = ?v.categories, + "guardrail pipeline blocked request (anthropic route)" + ); + return deny_policy( + StatusCode::FORBIDDEN, + &v.message(), + "content_policy_violation", + v.code(), + ); + } + Err(e) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_input", + error = %e, + "guardrail pipeline unavailable (anthropic route) — failing closed" + ); + return deny_policy( + guardrail_error_status(&e), + &e.to_string(), + "api_error", + e.code(), + ); + } + } + } + + // Native Messages pass-through (provider: anthropic, or Copilot's + // native /v1/messages) — no translation. + if upstream.provider == ProviderKind::Anthropic + || proxy::is_copilot_endpoint(&upstream.endpoint) + { + return forward_anthropic_passthrough( + state, + sandbox_name, + headers, + body, + upstream, + guardrail_pipeline, + policy.digest.clone(), + ) + .await; } // Translate Anthropic -> OpenAI chat completions request shape. @@ -337,6 +510,52 @@ pub(super) async fn anthropic_messages( } }; let anthropic_resp = openai_to_anthropic(&openai_resp, &requested_model); + + // Guardrail output scan (buffered, translated path). + if let Some(p) = guardrail_pipeline + .as_ref() + .filter(|p| p.covers(Direction::Output)) + { + let text = guardrails::extract_anthropic_output_text(&anthropic_resp); + match p.scan(&text, Direction::Output).await { + Ok(None) => {} + Ok(Some(v)) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_output", + categories = ?v.categories, + "guardrail pipeline blocked translated response (anthropic route)" + ); + return deny_policy( + StatusCode::FORBIDDEN, + &v.message(), + "content_policy_violation", + v.code(), + ); + } + Err(e) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_output", + error = %e, + "guardrail pipeline unavailable (anthropic route) — failing closed" + ); + return deny_policy( + guardrail_error_status(&e), + &e.to_string(), + "api_error", + e.code(), + ); + } + } + } + (StatusCode::OK, Json(anthropic_resp)).into_response() } Err(e) => { @@ -350,6 +569,183 @@ pub(super) async fn anthropic_messages( } } +/// Update running token counts from one Anthropic SSE event. +/// `message_start` carries `message.usage.input_tokens` (and an initial +/// `output_tokens`); each `message_delta` carries the cumulative +/// `usage.output_tokens`. +fn update_anthropic_usage(ev: &Value, input: &mut u64, output: &mut u64) { + match ev.get("type").and_then(|t| t.as_str()) { + Some("message_start") => { + if let Some(usage) = ev.get("message").and_then(|m| m.get("usage")) { + if let Some(i) = usage.get("input_tokens").and_then(|v| v.as_u64()) { + *input = i; + } + if let Some(o) = usage.get("output_tokens").and_then(|v| v.as_u64()) { + *output = o; + } + } + } + Some("message_delta") => { + if let Some(o) = ev + .get("usage") + .and_then(|u| u.get("output_tokens")) + .and_then(|v| v.as_u64()) + { + *output = o; + } + } + _ => {} + } +} + +/// Shared streaming-usage accumulator. The observer (inside the guard) +/// writes it as `usage` events arrive; the finalizer (outside the +/// guard) reads it once at the terminal. +#[derive(Default)] +struct StreamUsage { + input: u64, + output: u64, +} + +/// Observe token usage on a streamed Anthropic passthrough. Wraps the +/// RAW upstream stream (INSIDE the guard) and updates the shared +/// [`StreamUsage`] as it sees `message_start` / `message_delta` events. +/// Every byte passes through unchanged; nothing is recorded here. +/// +/// Recording is deliberately split out to [`finalize_stream_usage`] +/// (which wraps OUTSIDE the guard) because on a mid-stream guardrail +/// cut the guard stops polling its inner stream — so an inner tap's +/// terminal branch would never run and usage would be lost. The outer +/// finalizer still sees the guard's terminal and records what the +/// observer accumulated up to the cut. +fn observe_stream_usage( + stream: futures::stream::BoxStream<'static, Result>, + usage: Arc>, +) -> futures::stream::BoxStream<'static, Result> +where + E: Send + 'static, +{ + struct Ctx { + inner: futures::stream::BoxStream<'static, Result>, + usage: Arc>, + // Trailing incomplete UTF-8 bytes carried across a chunk + // boundary, mirroring the guardrail SSE state so a split + // multi-byte code point is not corrupted before parsing. + byte_carry: Vec, + line_carry: String, + } + fn scan_lines(ctx: &mut Ctx, text: &str) { + for line in text.lines() { + let Some(payload) = line.trim().strip_prefix("data:") else { + continue; + }; + let payload = payload.trim(); + if payload.is_empty() || payload == "[DONE]" { + continue; + } + if let Ok(ev) = serde_json::from_str::(payload) { + let mut u = ctx.usage.lock().unwrap(); + let StreamUsage { input, output } = &mut *u; + update_anthropic_usage(&ev, input, output); + } + } + } + let ctx = Ctx { + inner: stream, + usage, + byte_carry: Vec::new(), + line_carry: String::new(), + }; + futures::stream::unfold(ctx, |mut ctx| async move { + match ctx.inner.next().await { + Some(Ok(chunk)) => { + let mut bytes = std::mem::take(&mut ctx.byte_carry); + bytes.extend_from_slice(&chunk); + let decodable = match std::str::from_utf8(&bytes) { + Ok(_) => bytes.len(), + Err(e) if e.error_len().is_none() => e.valid_up_to(), + Err(_) => bytes.len(), + }; + ctx.byte_carry = bytes.split_off(decodable); + ctx.line_carry.push_str(&String::from_utf8_lossy(&bytes)); + if let Some(idx) = ctx.line_carry.rfind('\n') { + let complete = ctx.line_carry[..=idx].to_string(); + ctx.line_carry = ctx.line_carry[idx + 1..].to_string(); + scan_lines(&mut ctx, &complete); + } + Some((Ok(chunk), ctx)) + } + Some(Err(e)) => Some((Err(e), ctx)), + None => { + if !ctx.byte_carry.is_empty() { + let tail = std::mem::take(&mut ctx.byte_carry); + ctx.line_carry.push_str(&String::from_utf8_lossy(&tail)); + } + let tail = std::mem::take(&mut ctx.line_carry); + scan_lines(&mut ctx, &tail); + None + } + } + }) + .boxed() +} + +/// Record the observed streaming usage to the budget exactly once, at +/// the terminal of the (possibly guard-cut) stream. Wraps OUTSIDE the +/// guard so a mid-stream cut, an upstream error, or a clean end all +/// funnel through here. Every byte passes through unchanged. +fn finalize_stream_usage( + stream: futures::stream::BoxStream<'static, Result>, + usage: Arc>, + budget: crate::budget::TokenBudgetTracker, + sandbox: String, +) -> futures::stream::BoxStream<'static, Result> +where + E: Send + 'static, +{ + struct Ctx { + inner: futures::stream::BoxStream<'static, Result>, + usage: Arc>, + budget: crate::budget::TokenBudgetTracker, + sandbox: String, + recorded: bool, + } + async fn record(ctx: &mut Ctx) { + if ctx.recorded { + return; + } + ctx.recorded = true; + let total = { + let u = ctx.usage.lock().unwrap(); + u.input + u.output + }; + if total > 0 { + ctx.budget.record_usage(&ctx.sandbox, total).await; + } + } + let ctx = Ctx { + inner: stream, + usage, + budget, + sandbox, + recorded: false, + }; + futures::stream::unfold(ctx, |mut ctx| async move { + match ctx.inner.next().await { + Some(Ok(chunk)) => Some((Ok(chunk), ctx)), + Some(Err(e)) => { + record(&mut ctx).await; + Some((Err(e), ctx)) + } + None => { + record(&mut ctx).await; + None + } + } + }) + .boxed() +} + /// Native passthrough for Copilot's Anthropic Messages API. /// /// No translation: forwards body verbatim to `{copilot_endpoint}/v1/messages`, @@ -362,6 +758,8 @@ async fn forward_anthropic_passthrough( headers: HeaderMap, body: Bytes, upstream: crate::proxy::UpstreamConfig, + guardrail_pipeline: Option>, + policy_digest: String, ) -> axum::response::Response { let is_stream = serde_json::from_slice::(&body) .ok() @@ -393,10 +791,38 @@ async fn forward_anthropic_passthrough( .await { Ok((status, resp_headers, stream)) => { - let body = Body::from_stream(stream.map(|c| c.map_err(std::io::Error::other))); + // Observe usage INSIDE the guard (sees raw upstream + // events), then record OUTSIDE the guard so a mid-stream + // cut still bills the tokens consumed up to the cut. + let usage = Arc::new(std::sync::Mutex::new(StreamUsage::default())); + let observed = observe_stream_usage(stream, usage.clone()); + // Streaming output scan (Anthropic event dialect). + let guarded = match guardrail_pipeline + .as_ref() + .filter(|p| p.covers(Direction::Output)) + { + Some(p) => guardrails::guard_sse_stream( + observed, + p.clone(), + guardrails::StreamDialect::AnthropicMessages, + sandbox_name.to_string(), + policy_digest.clone(), + ), + None => observed, + }; + let finalized = finalize_stream_usage( + guarded, + usage, + state.budget.clone(), + sandbox_name.to_string(), + ); + let body = Body::from_stream(finalized.map(|c| c.map_err(std::io::Error::other))); let mut resp = axum::response::Response::builder().status(status); if let Some(h) = resp.headers_mut() { for (n, v) in resp_headers.iter() { + if is_hop_by_hop(n.as_str()) { + continue; + } h.insert(n.clone(), v.clone()); } h.insert( @@ -454,9 +880,61 @@ async fn forward_anthropic_passthrough( } } + // Guardrail output scan (buffered, Anthropic shape). + if status.is_success() + && let Some(p) = guardrail_pipeline + .as_ref() + .filter(|p| p.covers(Direction::Output)) + { + let text = guardrails::scan_text_or_raw( + &resp_body, + guardrails::extract_anthropic_output_text, + ); + match p.scan(&text, Direction::Output).await { + Ok(None) => {} + Ok(Some(v)) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy_digest, + decision = "deny", + gate = "guardrail_output", + categories = ?v.categories, + "guardrail pipeline blocked buffered response (anthropic route)" + ); + return deny_policy( + StatusCode::FORBIDDEN, + &v.message(), + "content_policy_violation", + v.code(), + ); + } + Err(e) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy_digest, + decision = "deny", + gate = "guardrail_output", + error = %e, + "guardrail pipeline unavailable (anthropic route) — failing closed" + ); + return deny_policy( + guardrail_error_status(&e), + &e.to_string(), + "api_error", + e.code(), + ); + } + } + } + let mut resp = axum::response::Response::builder().status(status); if let Some(h) = resp.headers_mut() { for (n, v) in resp_headers.iter() { + if is_hop_by_hop(n.as_str()) { + continue; + } h.insert(n.clone(), v.clone()); } if !h.contains_key(axum::http::header::CONTENT_TYPE) { @@ -568,4 +1046,173 @@ mod tests { let openai = anthropic_to_openai(&req); assert_eq!(openai["stop"], json!(["END", "STOP"])); } + + #[test] + fn usage_tap_reads_message_start_and_latest_delta() { + let mut input = 0; + let mut output = 0; + update_anthropic_usage( + &json!({"type": "message_start", "message": {"usage": {"input_tokens": 42, "output_tokens": 1}}}), + &mut input, + &mut output, + ); + update_anthropic_usage( + &json!({"type": "message_delta", "usage": {"output_tokens": 5}}), + &mut input, + &mut output, + ); + update_anthropic_usage( + &json!({"type": "message_delta", "usage": {"output_tokens": 17}}), + &mut input, + &mut output, + ); + assert_eq!(input, 42, "input from message_start"); + assert_eq!(output, 17, "output is the latest cumulative delta"); + } + + #[tokio::test] + async fn streaming_usage_is_recorded_to_budget() { + use bytes::Bytes; + use futures::stream::StreamExt as _; + let budget = crate::budget::TokenBudgetTracker::new(1_000_000, 0); + // Anthropic SSE split so a delta lands across a chunk boundary. + let frames = [ + "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":30,\"output_tokens\":1}}}\n\n", + "event: message_delta\ndata: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":", + "12}}\n\n", + "data: [DONE]\n\n", + ]; + let chunks: Vec> = + frames.iter().map(|f| Ok(Bytes::from(*f))).collect(); + // observer (no guard here) -> finalizer, mirroring the no-output- + // guardrail passthrough path. + let usage = Arc::new(std::sync::Mutex::new(StreamUsage::default())); + let observed = observe_stream_usage(futures::stream::iter(chunks).boxed(), usage.clone()); + let finalized = finalize_stream_usage(observed, usage, budget.clone(), "sbx".into()); + let _drained: Vec<_> = finalized.collect().await; + let (used, _) = budget.get_usage("sbx").await; + assert_eq!(used, 42, "30 input + 12 output recorded once at stream end"); + } + + #[tokio::test] + async fn buffered_policy_denial_carries_code_and_decision_headers() { + // Anthropic-native `error.type` is preserved, the stable machine + // code lands in `error.code`, and the decision headers are set. + let resp = deny_policy( + StatusCode::FORBIDDEN, + "blocked by guardrail", + "content_policy_violation", + "guardrail_blocked", + ); + assert_eq!(resp.status(), StatusCode::FORBIDDEN); + assert_eq!( + resp.headers() + .get("x-kars-decision") + .and_then(|v| v.to_str().ok()), + Some("blocked") + ); + assert!(resp.headers().get("x-kars-decision-by").is_some()); + let bytes = axum::body::to_bytes(resp.into_body(), 65536).await.unwrap(); + let v: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(v["error"]["type"], "content_policy_violation"); + assert_eq!(v["error"]["code"], "guardrail_blocked"); + } + + #[tokio::test] + async fn non_policy_denial_has_no_decision_headers() { + // A plain client error is not a policy decision: type == code and + // no decision headers. + let resp = deny_response( + StatusCode::TOO_MANY_REQUESTS, + "slow down", + "rate_limit_error", + ); + assert!(resp.headers().get("x-kars-decision").is_none()); + let bytes = axum::body::to_bytes(resp.into_body(), 65536).await.unwrap(); + let v: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(v["error"]["type"], "rate_limit_error"); + assert_eq!(v["error"]["code"], "rate_limit_error"); + } + + /// Regression: a mid-stream guardrail cut must still bill the tokens + /// consumed up to the cut. The observer sits inside the guard; the + /// finalizer outside records on the guard's terminal even though the + /// guard stopped polling its inner stream at the cut. + #[tokio::test] + async fn streaming_usage_recorded_when_guard_cuts_midstream() { + use bytes::Bytes; + use futures::stream::StreamExt as _; + use wiremock::matchers::{method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + // Moderation that flags any text containing "BLOCKED". + let moderation = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/moderations")) + .respond_with(move |req: &wiremock::Request| { + let flagged = String::from_utf8_lossy(&req.body).contains("BLOCKED"); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "results": [{ "flagged": flagged, "categories": { "illicit": flagged } }] + })) + }) + .mount(&moderation) + .await; + + let mut config = crate::config::Config::from_env().expect("config"); + config.openai_moderation_endpoint = moderation.uri(); + config.openai_moderation_api_key = Some("sk-mod-test".into()); + let pipeline = Arc::new( + GuardrailPipeline::from_stages( + &[crate::guardrails::GuardrailStageCfg { + provider: "openai-moderation".into(), + apply_to: crate::guardrails::ApplyTo::Output, + }], + &config, + &reqwest::Client::new(), + ) + .expect("pipeline builds"), + ); + + let budget = crate::budget::TokenBudgetTracker::new(1_000_000, 0); + + // message_start carries input=30; a single content_block_delta + // over the scan threshold forces a mid-stream scan that flags and + // cuts BEFORE any [DONE]. The trailing message_delta never arrives. + let big = "BLOCKED ".repeat(200); // > STREAM_SCAN_THRESHOLD_CHARS + let frames = vec![ + Ok(Bytes::from( + "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":30,\"output_tokens\":1}}}\n\n".to_string(), + )), + Ok(Bytes::from(format!( + "event: content_block_delta\ndata: {{\"type\":\"content_block_delta\",\"delta\":{{\"type\":\"text_delta\",\"text\":\"{big}\"}}}}\n\n" + ))), + ]; + let stream: futures::stream::BoxStream<'static, Result> = + futures::stream::iter(frames).boxed(); + + let usage = Arc::new(std::sync::Mutex::new(StreamUsage::default())); + let observed = observe_stream_usage(stream, usage.clone()); + let guarded = guardrails::guard_sse_stream( + observed, + pipeline, + guardrails::StreamDialect::AnthropicMessages, + "sbx".into(), + "sha256:t".into(), + ); + let finalized = finalize_stream_usage(guarded, usage, budget.clone(), "sbx".into()); + let out: String = finalized + .collect::>() + .await + .into_iter() + .filter_map(|c| c.ok()) + .map(|b| String::from_utf8_lossy(&b).into_owned()) + .collect(); + + assert!(out.contains("guardrail_blocked"), "guard must cut: {out}"); + let (used, _) = budget.get_usage("sbx").await; + assert_eq!( + used, 31, + "input(30)+message_start output(1) billed once despite the mid-stream cut, got {used}" + ); + } } diff --git a/inference-router/src/routes/chat_completions.rs b/inference-router/src/routes/chat_completions.rs index 2b53a7e20..62d539a9a 100644 --- a/inference-router/src/routes/chat_completions.rs +++ b/inference-router/src/routes/chat_completions.rs @@ -20,15 +20,18 @@ use futures::stream::StreamExt; use super::AppState; use super::inference_translate::{chat_to_responses_body, responses_to_chat_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; /// Inject the canonical `x-kars-decision*` triplet onto a response /// so downstream tooling (conformance-runner, observability pipelines, /// HTTP clients) can read the policy decision without parsing the /// upstream-specific body wording. CR/LF are sanitized from the reason /// per HTTP/1.1 header rules. -fn insert_decision_headers( +pub(super) fn insert_decision_headers( response: &mut axum::response::Response, decision: &'static str, by_kind: &'static str, @@ -46,6 +49,233 @@ fn insert_decision_headers( } } +/// Map a provider-resolution failure onto an OpenAI-shaped error +/// response: 501 for schema-valid-but-unimplemented providers +/// (`bedrock`), 503 for router-side config gaps. Never falls back to +/// Azure — see `routes::apply_provider_resolution`. +pub(super) fn provider_error_response(e: &ProviderError) -> axum::response::Response { + let (status, code) = match e { + ProviderError::Unimplemented { .. } => { + (StatusCode::NOT_IMPLEMENTED, "provider_unimplemented") + } + ProviderError::MissingEndpoint { .. } | ProviderError::MissingCredential { .. } => { + (StatusCode::SERVICE_UNAVAILABLE, "provider_unconfigured") + } + }; + let mut resp = ( + status, + Json(serde_json::json!({ + "error": { "message": e.to_string(), "type": "provider_error", "code": code } + })), + ) + .into_response(); + insert_decision_headers(&mut resp, "blocked", "InferencePolicy", &e.to_string()); + resp +} + +/// 403 response for a confirmed guardrail violation, with the +/// canonical `x-kars-decision*` triplet attached. +pub(super) fn guardrail_violation_response(v: &GuardrailViolation) -> axum::response::Response { + let mut resp = ( + StatusCode::FORBIDDEN, + Json(serde_json::json!({ + "error": { "message": v.message(), "type": "content_policy_violation", "code": v.code() } + })), + ) + .into_response(); + insert_decision_headers(&mut resp, "blocked", "InferencePolicy", &v.message()); + resp +} + +/// Fail-closed response for a guardrail that could not run: 503 for a +/// misconfigured stage, 502 for a backend outage. +pub(super) fn guardrail_error_response(e: &GuardrailError) -> axum::response::Response { + let status = match e { + GuardrailError::Config { .. } => StatusCode::SERVICE_UNAVAILABLE, + GuardrailError::Unavailable { .. } => StatusCode::BAD_GATEWAY, + }; + let mut resp = ( + status, + Json(serde_json::json!({ + "error": { "message": e.to_string(), "type": "guardrail_error", "code": e.code() } + })), + ) + .into_response(); + insert_decision_headers(&mut resp, "blocked", "InferencePolicy", &e.to_string()); + resp +} + +/// Build the policy's guardrail pipeline, or `None` when the policy +/// declares no stages. A declared-but-unbuildable pipeline is a +/// request-blocking error (fail closed). +pub(super) fn build_guardrail_pipeline( + state: &AppState, + policy: &crate::inference_policy_loader::InferencePolicySnapshot, +) -> Result>, GuardrailError> { + if policy.guardrails.is_empty() { + return Ok(None); + } + GuardrailPipeline::from_stages(&policy.guardrails, &state.config, &state.client) + .map(|p| Some(Arc::new(p))) +} + +/// A blocked output scan, kept as data so each transport picks its +/// wire shape (buffered → HTTP 403/502/503, SSE → error frame). +pub(super) enum OutputGuardrailBlock { + Violation(GuardrailViolation), + Error(GuardrailError), +} + +impl OutputGuardrailBlock { + /// SSE frame carrying this block's own type/code, so a backend + /// outage isn't mislabelled a content violation. + pub(super) fn sse_frame(&self) -> bytes::Bytes { + match self { + Self::Violation(v) => guardrails::violation_sse_frame(v), + Self::Error(e) => guardrails::error_sse_frame(e), + } + } +} + +/// Run the output-direction guardrail stages over a buffered +/// OpenAI-shaped response body. `None` when there is nothing to do or +/// the scan passes; `Some(block)` when the response must not reach +/// the client. Emits the audit log line on every block. +pub(super) async fn scan_openai_output_guardrails( + pipeline: Option<&Arc>, + resp_body: &[u8], + sandbox_name: &str, + policy_digest: &str, +) -> Option { + let p = pipeline?; + if !p.covers(Direction::Output) { + return None; + } + let text = guardrails::scan_text_or_raw(resp_body, guardrails::extract_openai_output_text); + match p.scan(&text, Direction::Output).await { + Ok(None) => None, + Ok(Some(v)) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy_digest, + decision = "deny", + gate = "guardrail_output", + categories = ?v.categories, + "guardrail pipeline blocked buffered response" + ); + Some(OutputGuardrailBlock::Violation(v)) + } + Err(e) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy_digest, + decision = "deny", + gate = "guardrail_output", + error = %e, + "guardrail pipeline unavailable — failing closed" + ); + Some(OutputGuardrailBlock::Error(e)) + } + } +} + +/// HTTP-response wrapper over [`scan_openai_output_guardrails`] for +/// the buffered branches: `Err(response)` carries the ready-made +/// block response (403 violation / 502 unavailable / 503 config). +/// +/// The error is boxed: an `axum::Response` is large, and an unboxed +/// `Result<(), Response>` trips `clippy::result_large_err` under +/// `-D warnings`. +pub(super) async fn enforce_openai_output_guardrails( + pipeline: Option<&Arc>, + resp_body: &[u8], + sandbox_name: &str, + policy_digest: &str, +) -> Result<(), Box> { + match scan_openai_output_guardrails(pipeline, resp_body, sandbox_name, policy_digest).await { + None => Ok(()), + Some(OutputGuardrailBlock::Violation(v)) => Err(Box::new(guardrail_violation_response(&v))), + Some(OutputGuardrailBlock::Error(e)) => Err(Box::new(guardrail_error_response(&e))), + } +} + +/// Why a route without provider/guardrail enforcement must refuse the +/// active policy. `Ok` ⇒ the route may proceed (Azure upstream, no +/// guardrails — the historic behaviour). +#[derive(Debug, PartialEq, Eq)] +pub(super) enum RouteGap { + NonAzureProvider, + Guardrails, +} + +/// Pure classifier for [`guard_unenforced_route`]: does this policy +/// need enforcement a non-chat route can't provide? Kept `AppState`- +/// free so the truth table is unit-testable. +pub(super) fn classify_route_gap( + policy: &crate::inference_policy_loader::InferencePolicySnapshot, + config: &crate::config::Config, +) -> Result<(), RouteGap> { + if !matches!( + crate::provider::resolve(policy.provider.as_deref(), config), + Ok(crate::provider::ProviderTarget::AzureOpenAI) + ) { + return Err(RouteGap::NonAzureProvider); + } + if !policy.guardrails.is_empty() { + return Err(RouteGap::Guardrails); + } + Ok(()) +} + +/// Fail-closed guard for inference/generation routes that do NOT +/// implement provider routing or the guardrail pipeline +/// (`completions`, `responses`, `embeddings`, image generation). +/// Such a route would silently bypass a policy that selects a +/// non-Azure provider or declares guardrail stages, so it refuses +/// instead. `None` ⇒ proceed. +pub(super) async fn guard_unenforced_route( + state: &AppState, + sandbox: &str, + route: &'static str, +) -> Option { + let policy = crate::inference_policy_loader::current_snapshot(&state.inference_policy).await; + let gap = classify_route_gap(&policy, &state.config).err()?; + let (status, code, msg) = match gap { + RouteGap::NonAzureProvider => ( + StatusCode::NOT_IMPLEMENTED, + "provider_unimplemented", + format!( + "InferencePolicy selects a non-Azure provider, unsupported on {route} in this \ + release — use /v1/chat/completions or /anthropic/v1/messages" + ), + ), + RouteGap::Guardrails => ( + StatusCode::FORBIDDEN, + "guardrail_route_unsupported", + format!( + "InferencePolicy declares a guardrail pipeline, not enforced on {route} in this \ + release — use /v1/chat/completions or /anthropic/v1/messages" + ), + ), + }; + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "route_enforcement_gap", + route, + "{msg}" + ); + // Code-carrying shape: `type` == `code` here (historic value), plus + // an explicit `error.code` so clients switch on one field everywhere. + let mut resp = errors::openai_coded(status, &msg, code, code).into_response(); + insert_decision_headers(&mut resp, "blocked", "InferencePolicy", &msg); + Some(resp) +} + /// POST /v1/chat/completions — the primary inference endpoint. pub(super) async fn chat_completions( State(state): State, @@ -197,6 +427,90 @@ pub(super) async fn chat_completions( // Slice 2d.1: honour `InferencePolicy.modelPreference.primary.deployment`. crate::routes::apply_model_preference_override(&mut upstream, &policy); + // Retarget at the policy-selected provider; fails closed. + if let Err(e) = crate::routes::apply_provider_resolution(&state, &mut upstream, &policy) { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "provider_resolution", + error = %e, + "InferencePolicy provider could not be resolved" + ); + return provider_error_response(&e); + } + + // Anthropic serves the Messages API — an OpenAI-shaped request + // here gets an explicit 501 pointing at /anthropic/v1/messages. + // Carries the coded error + decision headers like every other + // provider/guardrail denial on this route. + if upstream.provider == ProviderKind::Anthropic { + let msg = "InferencePolicy selects provider 'anthropic', which serves the Anthropic \ + Messages API — send Anthropic-shaped requests to /anthropic/v1/messages \ + (chat-completions translation for Anthropic is not implemented)"; + let mut resp = errors::openai_coded( + StatusCode::NOT_IMPLEMENTED, + msg, + "provider_error", + "provider_unimplemented", + ) + .into_response(); + insert_decision_headers(&mut resp, "blocked", "InferencePolicy", msg); + return resp; + } + + // Guardrail pipeline; a declared-but-unbuildable stage blocks. + let guardrail_pipeline = match build_guardrail_pipeline(&state, &policy) { + Ok(p) => p, + Err(e) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_config", + error = %e, + "guardrail pipeline could not be built — failing closed" + ); + return guardrail_error_response(&e); + } + }; + + // Input-direction scan, before any upstream forward. + if let Some(ref p) = guardrail_pipeline + && p.covers(Direction::Input) + { + let input_text = guardrails::scan_text_or_raw(&body, guardrails::extract_openai_input_text); + match p.scan(&input_text, Direction::Input).await { + Ok(None) => {} + Ok(Some(v)) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_input", + categories = ?v.categories, + "guardrail pipeline blocked request" + ); + return guardrail_violation_response(&v); + } + Err(e) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_input", + error = %e, + "guardrail pipeline unavailable — failing closed" + ); + return guardrail_error_response(&e); + } + } + } + // Defence-in-depth tool-schema filter: even when the upstream // runtime (e.g. a raw OpenAI SDK client outside OpenClaw) sends // `tools[]` schemas the AGT plugin would normally never have @@ -254,6 +568,8 @@ pub(super) async fn chat_completions( let headers = headers.clone(); let budget = state.budget.clone(); let sandbox_owned = sandbox_name.to_string(); + let guardrail_for_task = guardrail_pipeline.clone(); + let digest_for_task = policy.digest.clone(); tokio::spawn(async move { // Send keepalive comments every 5 seconds while waiting @@ -293,6 +609,20 @@ pub(super) async fn chat_completions( { budget.record_usage(&sandbox_owned, total).await; } + // Output scan on the Responses-API recovery + // path; block as an SSE frame (already committed + // to text/event-stream). + if let Some(block) = scan_openai_output_guardrails( + guardrail_for_task.as_ref(), + &chat_body, + &sandbox_owned, + &digest_for_task, + ) + .await + { + let _ = tx.send(Ok(block.sse_frame())).await; + return; + } let sse_data = format!( "data: {}\n\ndata: [DONE]\n\n", String::from_utf8_lossy(&chat_body) @@ -347,6 +677,16 @@ pub(super) async fn chat_completions( { state.budget.record_usage(sandbox_name, total).await; } + if let Err(block) = enforce_openai_output_guardrails( + guardrail_pipeline.as_ref(), + &chat_body, + sandbox_name, + &policy.digest, + ) + .await + { + return *block; + } let mut response = (resp_status, Body::from(chat_body)).into_response(); if let Some(ct) = resp_hdrs.get("content-type") { response.headers_mut().insert("content-type", ct.clone()); @@ -448,6 +788,16 @@ pub(super) async fn chat_completions( { state.budget.record_usage(sandbox_name, total).await; } + if let Err(block) = enforce_openai_output_guardrails( + guardrail_pipeline.as_ref(), + &chat_body, + sandbox_name, + &policy.digest, + ) + .await + { + return *block; + } // Wrap as SSE so the streaming client can parse it let sse = format!( "data: {}\n\ndata: [DONE]\n\n", @@ -493,6 +843,7 @@ pub(super) async fn chat_completions( let stream_blocked = std::sync::Arc::new(std::sync::atomic::AtomicBool::new(false)); let floor_for_stream = stream_floor.clone(); let digest_for_stream = stream_policy_digest.clone(); + let sandbox_for_guard = sandbox_owned.clone(); let wrapped = stream.map(move |chunk| { use std::sync::atomic::Ordering; if stream_blocked.load(Ordering::Relaxed) { @@ -581,7 +932,22 @@ pub(super) async fn chat_completions( } chunk }); - let body = Body::from_stream(wrapped); + // Streaming output scan; skipped when no output stages. + let guarded: futures::stream::BoxStream<'static, Result> = + match guardrail_pipeline + .as_ref() + .filter(|p| p.covers(Direction::Output)) + { + Some(p) => guardrails::guard_sse_stream( + wrapped.boxed(), + p.clone(), + guardrails::StreamDialect::OpenAiChat, + sandbox_for_guard, + policy.digest.clone(), + ), + None => wrapped.boxed(), + }; + let body = Body::from_stream(guarded); let mut response = (status, body).into_response(); if let Some(ct) = resp_headers.get("content-type") { response.headers_mut().insert("content-type", ct.clone()); @@ -665,6 +1031,16 @@ pub(super) async fn chat_completions( { state.budget.record_usage(sandbox_name, total).await; } + if let Err(block) = enforce_openai_output_guardrails( + guardrail_pipeline.as_ref(), + &chat_body, + sandbox_name, + &policy.digest, + ) + .await + { + return *block; + } let mut response = (resp_status, Body::from(chat_body)).into_response(); if let Some(ct) = resp_hdrs.get("content-type") { response.headers_mut().insert("content-type", ct.clone()); @@ -704,6 +1080,23 @@ pub(super) async fn chat_completions( } } + // Buffered output guardrail scan. Runs on EVERY buffered + // response, hoisted OUT of the JSON-parse block below so a + // non-JSON / truncated upstream body can't skip a declared + // output guardrail (`scan_text_or_raw` falls back to raw + // bytes). Fails closed: 403 violation / 502 unavailable / + // 503 misconfigured. + if let Err(block) = enforce_openai_output_guardrails( + guardrail_pipeline.as_ref(), + &resp_body, + sandbox_name, + &policy.digest, + ) + .await + { + return *block; + } + // Parse Foundry guardrail annotations and report to AGT governance. // On 200: prompt_filter_results at top level // On 400: error.innererror.content_filter_result @@ -1099,6 +1492,64 @@ async fn filter_disallowed_tools( #[cfg(test)] mod tests { use super::*; + use crate::inference_policy_loader::InferencePolicySnapshot; + + fn azure_cfg() -> crate::config::Config { + // No provider endpoints/keys configured — a bare Azure box. + let mut c = crate::config::Config::from_env().expect("config"); + c.anthropic_api_key = None; + c.ollama_endpoint = None; + c + } + + fn snapshot(provider: Option<&str>, guardrails: bool) -> InferencePolicySnapshot { + InferencePolicySnapshot { + provider: provider.map(String::from), + guardrails: if guardrails { + vec![crate::guardrails::GuardrailStageCfg { + provider: "openai-moderation".into(), + apply_to: crate::guardrails::ApplyTo::Both, + }] + } else { + vec![] + }, + ..InferencePolicySnapshot::default() + } + } + + #[test] + fn route_gap_allows_plain_azure_policy() { + assert!(classify_route_gap(&snapshot(None, false), &azure_cfg()).is_ok()); + assert!(classify_route_gap(&snapshot(Some("azure-openai"), false), &azure_cfg()).is_ok()); + // An unknown/informational tag routes to Azure ⇒ still allowed. + assert!(classify_route_gap(&snapshot(Some("gemini"), false), &azure_cfg()).is_ok()); + } + + #[test] + fn route_gap_refuses_non_azure_provider() { + // ollama/anthropic/bedrock all fail closed on these routes, + // whether or not their config is present. + assert_eq!( + classify_route_gap(&snapshot(Some("ollama"), false), &azure_cfg()), + Err(RouteGap::NonAzureProvider) + ); + assert_eq!( + classify_route_gap(&snapshot(Some("anthropic"), false), &azure_cfg()), + Err(RouteGap::NonAzureProvider) + ); + assert_eq!( + classify_route_gap(&snapshot(Some("bedrock"), false), &azure_cfg()), + Err(RouteGap::NonAzureProvider) + ); + } + + #[test] + fn route_gap_refuses_declared_guardrails() { + assert_eq!( + classify_route_gap(&snapshot(None, true), &azure_cfg()), + Err(RouteGap::Guardrails) + ); + } #[test] fn gate_allow_when_no_policy_cap() { diff --git a/inference-router/src/routes/inference.rs b/inference-router/src/routes/inference.rs index ae833e8ca..d0979b725 100644 --- a/inference-router/src/routes/inference.rs +++ b/inference-router/src/routes/inference.rs @@ -250,6 +250,12 @@ async fn completions( ) -> impl IntoResponse { let sandbox_name = resolve_sandbox_name(&headers); + if let Some(resp) = + super::chat_completions::guard_unenforced_route(&state, sandbox_name, "completions").await + { + return resp; + } + let upstream = state.upstream_config(sandbox_name); match proxy::forward( &state.auth, @@ -281,6 +287,12 @@ async fn responses( let sandbox_name_owned = resolve_sandbox_name(&headers).to_string(); let sandbox_name = sandbox_name_owned.as_str(); + if let Some(resp) = + super::chat_completions::guard_unenforced_route(&state, sandbox_name, "responses").await + { + return resp; + } + // Slice 2 DoD #7 — snapshot policy early so every audit log // emitted from this handler can carry `inference_policy_digest`. let policy = crate::inference_policy_loader::current_snapshot(&state.inference_policy).await; @@ -416,6 +428,12 @@ async fn embeddings( ) -> impl IntoResponse { let sandbox_name = resolve_sandbox_name(&headers); + if let Some(resp) = + super::chat_completions::guard_unenforced_route(&state, sandbox_name, "embeddings").await + { + return resp; + } + // Embeddings need a different deployment than chat — extract model from request body let mut upstream = state.upstream_config(sandbox_name); if let Ok(body_json) = serde_json::from_slice::(&body) { @@ -459,6 +477,13 @@ async fn images_generations( ) -> impl IntoResponse { let sandbox_name = resolve_sandbox_name(&headers); + if let Some(resp) = + super::chat_completions::guard_unenforced_route(&state, sandbox_name, "images/generations") + .await + { + return resp; + } + // AGT policy check — image generation is a tool invocation { let action = format!("image_generation:{deployment}"); @@ -666,6 +691,103 @@ async fn list_deployments(State(state): State) -> impl IntoResponse { } } +/// Strictly percent-decode one path segment. `None` on a malformed +/// escape (`%` not followed by two hex digits) or invalid UTF-8 , +/// callers reject the request rather than guess. +fn percent_decode_segment(seg: &str) -> Option { + let bytes = seg.as_bytes(); + let mut out = Vec::with_capacity(bytes.len()); + let mut i = 0; + while i < bytes.len() { + if bytes[i] == b'%' { + let hi = (*bytes.get(i + 1)? as char).to_digit(16)?; + let lo = (*bytes.get(i + 2)? as char).to_digit(16)?; + out.push((hi * 16 + lo) as u8); + i += 3; + } else { + out.push(bytes[i]); + i += 1; + } + } + String::from_utf8(out).ok() +} + +/// Split a raw request path into percent-decoded segments, rejecting +/// (`None`) any form that could reach a different upstream route than +/// the one classified: dot segments (`.` / `..`, literal or +/// percent-encoded) and decoded slashes/backslashes smuggled inside a +/// segment (`..%2f`, `..%5c`). [`foundry_proxy`] concatenates the raw +/// path into the upstream URL, where URL parsing resolves dot segments, +/// so classification and forwarding would otherwise disagree. +/// +/// Benign empty segments (`//`, trailing `/`) are COLLAPSED, not +/// rejected: they normalize away without creating a traversal, and +/// management paths like `/openai/files/` have always been proxied. +/// Rejecting them would 400 live customer traffic even with no policy +/// configured, so we canonicalize instead. +fn decoded_path_segments(path: &str) -> Option> { + let mut segments = Vec::new(); + for raw in path.trim_start_matches('/').split('/') { + let seg = percent_decode_segment(raw)?; + if seg.is_empty() { + continue; // collapse `//` and trailing `/` + } + if seg == "." || seg == ".." || seg.contains('/') || seg.contains('\\') { + return None; + } + segments.push(seg); + } + Some(segments) +} + +/// Default-deny guard classification for [`foundry_proxy`], over +/// segments already canonicalized by [`decoded_path_segments`]. +/// +/// The proxy implements neither provider routing nor the guardrail +/// pipeline, so any route that can serve model-generated inference +/// output must fail closed when the active `InferencePolicy` needs +/// either, otherwise an agent blocked on `/v1/chat/completions` +/// could rerun the same inference here and receive an unscanned +/// response. Rather than enumerating the inference-bearing families +/// (a miss re-opens the bypass), everything is guarded EXCEPT an +/// explicit exempt set of canonical management/storage APIs that do +/// not proxy an inference channel. Returns the route label for +/// `guard_unenforced_route`, or `None` when exempt, enforcement +/// scope is documented in docs/api/crd-reference.md. +fn guarded_foundry_route(segments: &[String]) -> Option<&'static str> { + const EXEMPT: &[&[&str]] = &[ + &["memory_stores"], + &["knowledgebases"], + &["evaluations"], + &["evaluators"], + &["evaluationrules"], + &["evaluationtaxonomies"], + &["indexes"], + &["connections"], + &["deployments"], + &["datasets"], + &["insights"], + &["schedules"], + &["redTeams"], + &["openai", "evals"], + &["openai", "vector_stores"], + &["openai", "files"], + &["openai", "containers"], + &["openai", "fine-tuning"], + ]; + if EXEMPT.iter().any(|family| { + segments.len() >= family.len() && family.iter().zip(segments).all(|(f, s)| s == f) + }) { + return None; + } + Some(match segments { + [first, ..] if first == "agents" => "agents", + [a, b, ..] if a == "openai" && b == "responses" => "openai/responses", + [a, b, ..] if a == "openai" && b == "conversations" => "openai/conversations", + _ => "foundry-proxy", + }) +} + /// Generic Foundry project-level API proxy. /// Forwards requests to the Foundry project endpoint with IMDS auth (ai.azure.com audience). /// Handles: /agents/*, /memory-stores/*, /knowledgebases/*, /evaluations/* @@ -678,6 +800,40 @@ async fn foundry_proxy( ) -> impl IntoResponse { let sandbox_name = resolve_sandbox_name(&headers); + // The raw path is concatenated into the upstream URL below, and + // URL parsing resolves `.`/`..` segments (including + // percent-encoded `%2e` forms), a crafted path could therefore + // reach a different upstream route than the one classified here. + // Benign empty segments are canonicalized; only forms that could + // normalize away into a different route are rejected. + let Some(segments) = decoded_path_segments(uri.path()) else { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + path = %uri.path(), + decision = "deny", + gate = "path_canonicalization", + "Foundry proxy path contains a dot segment, encoded slash/backslash, or malformed escape" + ); + return errors::openai_coded( + StatusCode::BAD_REQUEST, + "path contains a dot segment, encoded slash/backslash, or malformed escape", + "invalid_path", + "invalid_path", + ) + .into_response(); + }; + + // Default-deny fail-closed guard: every Foundry proxy route is + // guarded unless it is on the canonical management/storage exempt + // list (see `guarded_foundry_route` / `guard_unenforced_route`). + if let Some(route) = guarded_foundry_route(&segments) + && let Some(resp) = + super::chat_completions::guard_unenforced_route(&state, sandbox_name, route).await + { + return resp; + } + // Use project endpoint for agent/standalone APIs, fall back to foundry/openai endpoint let endpoint = state .config @@ -914,7 +1070,110 @@ async fn foundry_proxy( #[cfg(test)] mod tests { - use super::strip_project_prefix; + use super::{decoded_path_segments, guarded_foundry_route, strip_project_prefix}; + + fn classify(path: &str) -> Option<&'static str> { + guarded_foundry_route(&decoded_path_segments(path).expect("canonical path")) + } + + #[test] + fn classifies_inference_bearing_foundry_paths() { + assert_eq!(classify("/agents"), Some("agents")); + assert_eq!(classify("/agents/a1/runs"), Some("agents")); + assert_eq!(classify("/openai/responses"), Some("openai/responses")); + assert_eq!( + classify("/openai/responses/resp_123"), + Some("openai/responses") + ); + assert_eq!( + classify("/openai/conversations"), + Some("openai/conversations") + ); + assert_eq!( + classify("/openai/conversations/c1/items"), + Some("openai/conversations") + ); + } + + #[test] + fn unknown_and_lookalike_families_fail_closed() { + // Default-deny: anything not on the exempt list is guarded, + // including lookalike names and families this router has + // never heard of. + for path in ["/agentsmith", "/openai/responsesx", "/openai/threads"] { + assert_eq!(classify(path), Some("foundry-proxy"), "path: {path}"); + } + } + + #[test] + fn exempts_management_and_storage_families() { + for path in [ + "/memory_stores", + "/knowledgebases/kb1/queries", + "/openai/files", + "/openai/vector_stores/vs1", + "/openai/containers/c1/files/f1/content", + "/openai/fine-tuning/jobs", + "/evaluations", + "/redTeams/runs", + "/schedules/s1", + ] { + assert_eq!(classify(path), None, "path: {path}"); + } + } + + #[test] + fn rejects_paths_that_normalize_differently() { + for path in [ + "/openai/files/../responses", + "/openai/files/%2e%2e/responses", + "/openai/files/%2E%2E/responses", + "/openai/files/./responses", + "/openai/files/%2e/responses", + "/memory_stores/../openai/responses", + "/openai/files/..%2fresponses", // decoded slash inside a segment + "/openai/files/..%5cresponses", // decoded backslash + "/openai/files/%zz", // malformed escape + "/openai/files/%2", // truncated escape + ] { + assert_eq!(decoded_path_segments(path), None, "path: {path}"); + } + } + + #[test] + fn collapses_benign_empty_segments_instead_of_rejecting() { + // Trailing and double slashes are canonicalized, not rejected, + // so historically proxied management paths keep working with no + // policy configured. Classification is unchanged by the collapse. + assert_eq!( + decoded_path_segments("/openai/files/"), + Some(vec!["openai".to_string(), "files".to_string()]) + ); + assert_eq!( + decoded_path_segments("/memory_stores/"), + Some(vec!["memory_stores".to_string()]) + ); + assert_eq!( + decoded_path_segments("/openai//files"), + Some(vec!["openai".to_string(), "files".to_string()]) + ); + // Exempt families stay exempt; guarded families stay guarded. + assert_eq!(classify("/openai/files/"), None); + assert_eq!(classify("/memory_stores/"), None); + assert_eq!(classify("/openai//responses"), Some("openai/responses")); + assert_eq!(classify("/agents/"), Some("agents")); + } + + #[test] + fn decodes_benign_percent_encoding_for_classification() { + // Percent-encoding that hides a guarded family name must not + // dodge classification: the raw path is forwarded verbatim + // and most upstream servers decode it back. + assert_eq!(classify("/openai/%72esponses"), Some("openai/responses")); + assert_eq!(classify("/%61gents/a1"), Some("agents")); + // Encoded exempt families stay exempt. + assert_eq!(classify("/openai/%66iles/f1"), None); + } #[test] fn strips_foundry_project_prefix() { diff --git a/inference-router/src/routes/mod.rs b/inference-router/src/routes/mod.rs index 38428bd5f..bcc92eca1 100644 --- a/inference-router/src/routes/mod.rs +++ b/inference-router/src/routes/mod.rs @@ -181,23 +181,21 @@ impl AppState { .unwrap_or_else(|_| "/var/lib/kars/token-budgets.json".into()); let budget = if persist_path.is_empty() { TokenBudgetTracker::new(config.token_budget_daily, config.token_budget_per_request) + } else if let Some(parent) = std::path::Path::new(&persist_path).parent() + && let Err(e) = std::fs::create_dir_all(parent) + { + tracing::warn!( + path = %parent.display(), + error = %e, + "Could not create token-budget persistence dir — falling back to in-memory" + ); + TokenBudgetTracker::new(config.token_budget_daily, config.token_budget_per_request) } else { - if let Some(parent) = std::path::Path::new(&persist_path).parent() - && let Err(e) = std::fs::create_dir_all(parent) - { - tracing::warn!( - path = %parent.display(), - error = %e, - "Could not create token-budget persistence dir — falling back to in-memory" - ); - TokenBudgetTracker::new(config.token_budget_daily, config.token_budget_per_request) - } else { - TokenBudgetTracker::with_persistence( - config.token_budget_daily, - config.token_budget_per_request, - &persist_path, - ) - } + TokenBudgetTracker::with_persistence( + config.token_budget_daily, + config.token_budget_per_request, + &persist_path, + ) }; let sandbox_name = std::env::var("SANDBOX_NAME").unwrap_or_else(|_| "unknown".into()); @@ -366,12 +364,45 @@ impl AppState { .and_then(|g| g.clone()) .unwrap_or_else(|| self.config.default_model.clone()); - UpstreamConfig { - endpoint, - deployment, - sandbox_name: sandbox_name.to_string(), + UpstreamConfig::azure(endpoint, deployment, sandbox_name.to_string()) + } +} + +/// Retarget `upstream` at the policy-selected provider (no-op for +/// Azure / no policy). Fails closed on `bedrock` or missing config — +/// the handler must surface the error, not fall back to Azure. +pub(crate) fn apply_provider_resolution( + state: &AppState, + upstream: &mut UpstreamConfig, + policy: &crate::inference_policy_loader::InferencePolicySnapshot, +) -> Result<(), crate::provider::ProviderError> { + let target = crate::provider::resolve(policy.provider.as_deref(), &state.config)?; + match target { + crate::provider::ProviderTarget::AzureOpenAI => {} + crate::provider::ProviderTarget::Anthropic { endpoint, api_key } => { + tracing::info!( + sandbox = %upstream.sandbox_name, + endpoint = %endpoint, + digest = %policy.digest, + "InferencePolicy provider: routing to Anthropic" + ); + upstream.endpoint = endpoint; + upstream.provider = crate::provider::ProviderKind::Anthropic; + upstream.api_key = Some(api_key); + } + crate::provider::ProviderTarget::Ollama { endpoint } => { + tracing::info!( + sandbox = %upstream.sandbox_name, + endpoint = %endpoint, + digest = %policy.digest, + "InferencePolicy provider: routing to Ollama" + ); + upstream.endpoint = endpoint; + upstream.provider = crate::provider::ProviderKind::Ollama; + upstream.api_key = None; } } + Ok(()) } /// Slice 2d.1 — apply `modelPreference.primary.deployment` from a diff --git a/inference-router/tests/agt_governance_integration.rs b/inference-router/tests/agt_governance_integration.rs index d397f43e6..190a1dde1 100644 --- a/inference-router/tests/agt_governance_integration.rs +++ b/inference-router/tests/agt_governance_integration.rs @@ -55,6 +55,12 @@ fn test_state(sandbox: &str, admin_token: Option<&str>) -> AppState { registry_mode: RegistryMode::Local, registry_url: None, provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: None, + ollama_endpoint: None, + openai_moderation_endpoint: "https://api.openai.com".into(), + openai_moderation_api_key: None, + openai_moderation_model: "omni-moderation-latest".into(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), policy_provider: Arc::clone(&governance) as Arc, diff --git a/inference-router/tests/anthropic_buffered_guardrail.rs b/inference-router/tests/anthropic_buffered_guardrail.rs new file mode 100644 index 000000000..d23733e3e --- /dev/null +++ b/inference-router/tests/anthropic_buffered_guardrail.rs @@ -0,0 +1,265 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Route-level tests for the NATIVE buffered Anthropic passthrough +//! guardrail contract (Pal's round-4 HIGH). These drive the real +//! `/v1/messages` handler end to end, not `deny_policy` directly: +//! an `InferencePolicy` selecting `provider: anthropic` with an +//! `openai-moderation` output stage, a wiremock Anthropic upstream, and +//! a wiremock moderation backend. +//! +//! The identical request must return the same stable `error.code` and +//! `x-kars-decision*` headers whether it is buffered or streamed. Here +//! we pin the buffered (`stream: false`) native path: a flagged output +//! gives 403 `guardrail_blocked` and a moderation outage gives 502 +//! `guardrail_unavailable`, both with the decision headers attached. + +use std::sync::Arc; + +use axum::{ + Router, + body::Body, + http::{Request, StatusCode}, +}; +use serde_json::Value; +use tower::ServiceExt; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +use kars_inference_router::auth::WorkloadIdentityAuth; +use kars_inference_router::blocklist::Blocklist; +use kars_inference_router::budget::TokenBudgetTracker; +use kars_inference_router::config::{Config, RegistryMode}; +use kars_inference_router::egress_blocked::BlockedBuffer; +use kars_inference_router::governance::Governance; +use kars_inference_router::guardrails::{ApplyTo, GuardrailStageCfg}; +use kars_inference_router::handoff::{ + DrainState, HandoffSession, HandoffTokenStore, PendingHandoffStore, +}; +use kars_inference_router::inference_policy_loader::LoadedInferencePolicy; +use kars_inference_router::mesh::{MeshInbox, MeshMetrics}; +use kars_inference_router::policy_status::PolicyStatusRegistry; +use kars_inference_router::providers::{AuditSink, PolicyDecisionProvider, SigningProvider}; +use kars_inference_router::routes::{AppState, inference_routes}; + +fn test_state(anthropic_endpoint: String, moderation_endpoint: String) -> AppState { + let policy_status = Arc::new(PolicyStatusRegistry::new()); + let governance = Arc::new(Governance::new_with_status( + "sb-test", + policy_status.clone(), + )); + AppState { + auth: Arc::new(WorkloadIdentityAuth::new()), + copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), + client: reqwest::Client::new(), + config: Arc::new(Config { + port: 0, + foundry_endpoint: None, + foundry_project_endpoint: None, + azure_openai_endpoint: None, + default_model: "claude-x".into(), + content_safety_enabled: false, + prompt_shields_enabled: false, + content_safety_endpoint: None, + token_budget_daily: 1_000_000_000, + token_budget_per_request: 1_000_000_000, + registry_mode: RegistryMode::Local, + registry_url: None, + provider_override: None, + anthropic_endpoint, + anthropic_api_key: Some("sk-ant-router-held".into()), + ollama_endpoint: None, + openai_moderation_endpoint: moderation_endpoint, + openai_moderation_api_key: Some("sk-mod-test".into()), + openai_moderation_model: "omni-moderation-latest".into(), + }), + budget: TokenBudgetTracker::new(1_000_000_000, 1_000_000_000), + policy_provider: Arc::clone(&governance) as Arc, + audit_sink: Arc::clone(&governance) as Arc, + signing_provider: Arc::clone(&governance) as Arc, + governance, + blocklist: Blocklist::disabled(), + blocked_egress: Arc::new(BlockedBuffer::with_defaults()), + sandbox_name: Arc::new("sb-test".to_string()), + inbox: Arc::new(MeshInbox::new()), + mesh_metrics: Arc::new(MeshMetrics::new()), + model_override: Arc::new(std::sync::RwLock::new(None)), + admin_token: None, + responses_only_models: Arc::new(std::sync::RwLock::new(Default::default())), + handoff_tokens: HandoffTokenStore::new(), + handoff_session: HandoffSession::new(), + drain_state: DrainState::new(), + pending_handoff: PendingHandoffStore::new(), + policy_status, + inference_policy: kars_inference_router::inference_policy_loader::empty_handle(), + memory_binding: kars_inference_router::memory_binding_loader::empty_handle(), + egress_allowlist: kars_inference_router::egress_allowlist_loader::empty_handle(), + deployment_health: Arc::new( + kars_inference_router::deployment_health::DeploymentHealthRegistry::new(), + ), + } +} + +async fn install_anthropic_output_guardrail(state: &AppState) { + let policy = LoadedInferencePolicy { + digest: "sha256:test".into(), + source_path: "/tmp/test-policy".into(), + per_request_tokens: None, + daily_tokens: None, + monthly_tokens: None, + content_safety: Default::default(), + model_preference: None, + provider: Some("anthropic".into()), + guardrails: vec![GuardrailStageCfg { + provider: "openai-moderation".into(), + apply_to: ApplyTo::Output, + }], + raw: serde_json::json!({}), + }; + *state.inference_policy.write().await = Some(policy); +} + +/// A buffered (`stream: false`) native Anthropic reply whose text flags +/// the moderation backend must be blocked with the stable coded error +/// and the decision headers, matching the streaming/translated paths. +#[tokio::test] +async fn native_buffered_violation_is_coded_and_carries_decision_headers() { + let anthropic = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "msg_test", + "type": "message", + "role": "assistant", + "content": [{ "type": "text", "text": "here is RANSOM material" }], + "stop_reason": "end_turn", + "usage": { "input_tokens": 4, "output_tokens": 5 } + }))) + .mount(&anthropic) + .await; + + let moderation = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/moderations")) + .respond_with(move |req: &wiremock::Request| { + let flagged = String::from_utf8_lossy(&req.body).contains("RANSOM"); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "results": [{ "flagged": flagged, "categories": { "illicit": flagged } }] + })) + }) + .mount(&moderation) + .await; + + let state = test_state(anthropic.uri(), moderation.uri()); + install_anthropic_output_guardrail(&state).await; + let app = Router::new().merge(inference_routes()).with_state(state); + + let req = Request::builder() + .method("POST") + .uri("/v1/messages") + .header("content-type", "application/json") + .header("anthropic-version", "2023-06-01") + .body(Body::from( + r#"{"model":"claude-x","max_tokens":32,"stream":false,"messages":[{"role":"user","content":"hi"}]}"#, + )) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + + let status = resp.status(); + let decision = resp + .headers() + .get("x-kars-decision") + .and_then(|v| v.to_str().ok()) + .map(str::to_string); + let bytes = axum::body::to_bytes(resp.into_body(), 1_048_576) + .await + .unwrap(); + let body: Value = serde_json::from_slice(&bytes).unwrap_or(Value::Null); + + assert_eq!(status, StatusCode::FORBIDDEN, "buffered violation: {body}"); + assert_eq!( + body["error"]["code"].as_str(), + Some("guardrail_blocked"), + "stable machine code required: {body}" + ); + assert_eq!( + body["error"]["type"].as_str(), + Some("content_policy_violation"), + "Anthropic-native type preserved: {body}" + ); + assert_eq!( + decision.as_deref(), + Some("blocked"), + "x-kars-decision header must be attached" + ); +} + +/// A moderation-backend outage on the buffered native path must fail +/// closed with the specific `guardrail_unavailable` code (not a generic +/// `api_error`) and the decision headers. +#[tokio::test] +async fn native_buffered_moderation_outage_is_coded_unavailable() { + let anthropic = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "msg_test", + "type": "message", + "role": "assistant", + "content": [{ "type": "text", "text": "anything at all" }], + "stop_reason": "end_turn", + "usage": { "input_tokens": 4, "output_tokens": 3 } + }))) + .mount(&anthropic) + .await; + + // Moderation is down → the pipeline must fail closed. + let moderation = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/moderations")) + .respond_with(ResponseTemplate::new(500)) + .mount(&moderation) + .await; + + let state = test_state(anthropic.uri(), moderation.uri()); + install_anthropic_output_guardrail(&state).await; + let app = Router::new().merge(inference_routes()).with_state(state); + + let req = Request::builder() + .method("POST") + .uri("/v1/messages") + .header("content-type", "application/json") + .header("anthropic-version", "2023-06-01") + .body(Body::from( + r#"{"model":"claude-x","max_tokens":32,"stream":false,"messages":[{"role":"user","content":"hi"}]}"#, + )) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + + let status = resp.status(); + let decision = resp + .headers() + .get("x-kars-decision") + .and_then(|v| v.to_str().ok()) + .map(str::to_string); + let bytes = axum::body::to_bytes(resp.into_body(), 1_048_576) + .await + .unwrap(); + let body: Value = serde_json::from_slice(&bytes).unwrap_or(Value::Null); + + assert_eq!( + status, + StatusCode::BAD_GATEWAY, + "outage fails closed: {body}" + ); + assert_eq!( + body["error"]["code"].as_str(), + Some("guardrail_unavailable"), + "misconfig vs outage must not collapse: {body}" + ); + assert_eq!( + decision.as_deref(), + Some("blocked"), + "x-kars-decision header must be attached" + ); +} diff --git a/inference-router/tests/chat_output_guardrail_nonjson.rs b/inference-router/tests/chat_output_guardrail_nonjson.rs new file mode 100644 index 000000000..2e49d415d --- /dev/null +++ b/inference-router/tests/chat_output_guardrail_nonjson.rs @@ -0,0 +1,232 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Regression for the HIGH review finding: a buffered chat-completions +//! response whose upstream body is NOT valid JSON must still be scanned +//! by a declared output guardrail (via the raw-text fallback), not +//! returned to the client verbatim. +//! +//! Driven through the real router: an Ollama-provider policy (no +//! upstream auth needed) with an `openai-moderation` output guardrail. +//! The Ollama mock returns a non-JSON 200 body containing a flagged +//! marker; the moderation mock flags any input containing it. The +//! handler must block with 403 rather than pass the body through. + +use std::sync::Arc; + +use axum::{ + Router, + body::Body, + http::{Request, StatusCode}, +}; +use serde_json::Value; +use tower::ServiceExt; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +use kars_inference_router::auth::WorkloadIdentityAuth; +use kars_inference_router::blocklist::Blocklist; +use kars_inference_router::budget::TokenBudgetTracker; +use kars_inference_router::config::{Config, RegistryMode}; +use kars_inference_router::egress_blocked::BlockedBuffer; +use kars_inference_router::governance::Governance; +use kars_inference_router::guardrails::{ApplyTo, GuardrailStageCfg}; +use kars_inference_router::handoff::{ + DrainState, HandoffSession, HandoffTokenStore, PendingHandoffStore, +}; +use kars_inference_router::inference_policy_loader::LoadedInferencePolicy; +use kars_inference_router::mesh::{MeshInbox, MeshMetrics}; +use kars_inference_router::policy_status::PolicyStatusRegistry; +use kars_inference_router::providers::{AuditSink, PolicyDecisionProvider, SigningProvider}; +use kars_inference_router::routes::{AppState, inference_routes}; + +fn test_state(ollama_endpoint: String, moderation_endpoint: String) -> AppState { + let policy_status = Arc::new(PolicyStatusRegistry::new()); + let governance = Arc::new(Governance::new_with_status( + "sb-test", + policy_status.clone(), + )); + AppState { + auth: Arc::new(WorkloadIdentityAuth::new()), + copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), + client: reqwest::Client::new(), + config: Arc::new(Config { + port: 0, + foundry_endpoint: None, + foundry_project_endpoint: None, + azure_openai_endpoint: None, + default_model: "llama3.1".into(), + content_safety_enabled: false, + prompt_shields_enabled: false, + content_safety_endpoint: None, + token_budget_daily: 1_000_000_000, + token_budget_per_request: 1_000_000_000, + registry_mode: RegistryMode::Local, + registry_url: None, + provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: None, + ollama_endpoint: Some(ollama_endpoint), + openai_moderation_endpoint: moderation_endpoint, + openai_moderation_api_key: Some("sk-mod-test".into()), + openai_moderation_model: "omni-moderation-latest".into(), + }), + budget: TokenBudgetTracker::new(1_000_000_000, 1_000_000_000), + policy_provider: Arc::clone(&governance) as Arc, + audit_sink: Arc::clone(&governance) as Arc, + signing_provider: Arc::clone(&governance) as Arc, + governance, + blocklist: Blocklist::disabled(), + blocked_egress: Arc::new(BlockedBuffer::with_defaults()), + sandbox_name: Arc::new("sb-test".to_string()), + inbox: Arc::new(MeshInbox::new()), + mesh_metrics: Arc::new(MeshMetrics::new()), + model_override: Arc::new(std::sync::RwLock::new(None)), + admin_token: None, + responses_only_models: Arc::new(std::sync::RwLock::new(Default::default())), + handoff_tokens: HandoffTokenStore::new(), + handoff_session: HandoffSession::new(), + drain_state: DrainState::new(), + pending_handoff: PendingHandoffStore::new(), + policy_status, + inference_policy: kars_inference_router::inference_policy_loader::empty_handle(), + memory_binding: kars_inference_router::memory_binding_loader::empty_handle(), + egress_allowlist: kars_inference_router::egress_allowlist_loader::empty_handle(), + deployment_health: Arc::new( + kars_inference_router::deployment_health::DeploymentHealthRegistry::new(), + ), + } +} + +async fn install_output_guardrail_ollama(state: &AppState) { + let policy = LoadedInferencePolicy { + digest: "sha256:test".into(), + source_path: "/tmp/test-policy".into(), + per_request_tokens: None, + daily_tokens: None, + monthly_tokens: None, + content_safety: Default::default(), + model_preference: None, + provider: Some("ollama".into()), + guardrails: vec![GuardrailStageCfg { + provider: "openai-moderation".into(), + apply_to: ApplyTo::Output, + }], + raw: serde_json::json!({}), + }; + *state.inference_policy.write().await = Some(policy); +} + +async fn post_chat(app: &Router, body: &str) -> (StatusCode, Value) { + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from(body.to_owned())) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + let status = resp.status(); + let bytes = axum::body::to_bytes(resp.into_body(), 4_194_304) + .await + .unwrap(); + let json = serde_json::from_slice(&bytes).unwrap_or(Value::Null); + (status, json) +} + +#[tokio::test] +async fn non_json_upstream_body_is_still_scanned_and_blocked() { + // Ollama upstream returns a NON-JSON 200 body carrying the flagged + // marker, the exact shape that previously slipped past the + // JSON-nested output scan. + let ollama = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/plain") + .set_body_string("upstream meltdown: RANSOM instructions here "), + ) + .mount(&ollama) + .await; + + // Moderation flags any input containing "RANSOM". + let moderation = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/moderations")) + .respond_with(move |req: &wiremock::Request| { + let flagged = String::from_utf8_lossy(&req.body).contains("RANSOM"); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "results": [{ "flagged": flagged, "categories": { "illicit": flagged } }] + })) + }) + .mount(&moderation) + .await; + + let state = test_state(ollama.uri(), moderation.uri()); + install_output_guardrail_ollama(&state).await; + let app = Router::new().merge(inference_routes()).with_state(state); + + let (status, body) = post_chat( + &app, + r#"{"model":"llama3.1","messages":[{"role":"user","content":"hi"}]}"#, + ) + .await; + + assert_eq!( + status, + StatusCode::FORBIDDEN, + "non-JSON upstream body must be scanned and blocked, got {status}: {body}" + ); + assert_eq!( + body["error"]["code"].as_str(), + Some("guardrail_blocked"), + "expected a guardrail violation block: {body}" + ); + // The flagged upstream text must not have leaked to the client. + assert!( + !serde_json::to_string(&body) + .unwrap() + .contains("RANSOM instructions"), + "flagged upstream text must not reach the client: {body}" + ); +} + +#[tokio::test] +async fn non_json_clean_upstream_body_passes_through() { + // Control: a clean non-JSON body (no flagged marker) is scanned and + // released unchanged, the raw-text fallback does not over-block. + let ollama = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("content-type", "text/plain") + .set_body_string("a perfectly benign non-json reply"), + ) + .mount(&ollama) + .await; + + let moderation = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/moderations")) + .respond_with(move |req: &wiremock::Request| { + let flagged = String::from_utf8_lossy(&req.body).contains("RANSOM"); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "results": [{ "flagged": flagged, "categories": { "illicit": flagged } }] + })) + }) + .mount(&moderation) + .await; + + let state = test_state(ollama.uri(), moderation.uri()); + install_output_guardrail_ollama(&state).await; + let app = Router::new().merge(inference_routes()).with_state(state); + + let (status, _body) = post_chat( + &app, + r#"{"model":"llama3.1","messages":[{"role":"user","content":"hi"}]}"#, + ) + .await; + + assert_eq!(status, StatusCode::OK, "clean non-JSON body must pass"); +} diff --git a/inference-router/tests/egress_blocked_endpoint.rs b/inference-router/tests/egress_blocked_endpoint.rs index 48773bdaf..125e79e64 100644 --- a/inference-router/tests/egress_blocked_endpoint.rs +++ b/inference-router/tests/egress_blocked_endpoint.rs @@ -53,6 +53,12 @@ fn test_state() -> AppState { registry_mode: RegistryMode::Local, registry_url: None, provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: None, + ollama_endpoint: None, + openai_moderation_endpoint: "https://api.openai.com".into(), + openai_moderation_api_key: None, + openai_moderation_model: "omni-moderation-latest".into(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), policy_provider: Arc::clone(&governance) as Arc, diff --git a/inference-router/tests/failover_walk.rs b/inference-router/tests/failover_walk.rs index 3f7cb13ce..cb7e25790 100644 --- a/inference-router/tests/failover_walk.rs +++ b/inference-router/tests/failover_walk.rs @@ -27,6 +27,7 @@ use kars_inference_router::failover::forward_with_failover; use kars_inference_router::inference_policy_loader::{ InferencePolicySnapshot, ModelPreference, ModelRef, }; +use kars_inference_router::provider::ProviderKind; use kars_inference_router::proxy::UpstreamConfig; use serde_json::Value; use std::net::SocketAddr; @@ -133,6 +134,8 @@ async fn primary_503_falls_through_to_fallback_200() { endpoint: base, deployment: "fallback-up".into(), sandbox_name: "sbx".into(), + provider: ProviderKind::AzureOpenAI, + api_key: None, }; let snap = snapshot("primary-down", &["fallback-up"]); @@ -195,6 +198,8 @@ async fn unhealthy_primary_is_skipped_in_second_pass() { endpoint: base, deployment: "fallback-up".into(), sandbox_name: "sbx".into(), + provider: ProviderKind::AzureOpenAI, + api_key: None, }; let snap = snapshot("primary-down", &["fallback-up"]); @@ -248,6 +253,8 @@ async fn all_unhealthy_still_punches_primary_for_last_resort() { endpoint: base, deployment: "primary-down".into(), sandbox_name: "sbx".into(), + provider: ProviderKind::AzureOpenAI, + api_key: None, }; let snap = snapshot("primary-down", &["fallback-up"]); diff --git a/inference-router/tests/foundry_route_guard.rs b/inference-router/tests/foundry_route_guard.rs new file mode 100644 index 000000000..e8e3eab84 --- /dev/null +++ b/inference-router/tests/foundry_route_guard.rs @@ -0,0 +1,404 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Fail-closed guard on inference-bearing Foundry proxy routes. +//! +//! An `InferencePolicy` that declares guardrails (or selects a +//! non-Azure provider) is enforced on `/v1/chat/completions` and the +//! Anthropic Messages routes — but `foundry_proxy` implements neither. +//! These tests prove the guard makes the bypass impossible: an agent +//! blocked by a guardrail on `/v1/chat/completions` must NOT be able +//! to rerun the same inference through `/openai/responses*`, +//! `/openai/conversations*`, or `/agents*` and receive an unscanned +//! response. +//! +//! The router is assembled with the same `Router::merge` wiring +//! `main.rs` uses, so the guard is exercised through real routing. +//! `foundry_endpoint` points at the GitHub Models marketplace so +//! un-guarded requests terminate in the deterministic, network-free +//! `is_github_models()` 501 — reaching that branch proves the request +//! got PAST the guard. + +use std::sync::Arc; + +use axum::{ + Router, + body::Body, + http::{Request, StatusCode}, +}; +use serde_json::Value; +use tower::ServiceExt; + +use kars_inference_router::auth::WorkloadIdentityAuth; +use kars_inference_router::blocklist::Blocklist; +use kars_inference_router::budget::TokenBudgetTracker; +use kars_inference_router::config::{Config, RegistryMode}; +use kars_inference_router::egress_blocked::BlockedBuffer; +use kars_inference_router::governance::Governance; +use kars_inference_router::guardrails::{ApplyTo, GuardrailStageCfg}; +use kars_inference_router::handoff::{ + DrainState, HandoffSession, HandoffTokenStore, PendingHandoffStore, +}; +use kars_inference_router::inference_policy_loader::LoadedInferencePolicy; +use kars_inference_router::mesh::{MeshInbox, MeshMetrics}; +use kars_inference_router::policy_status::PolicyStatusRegistry; +use kars_inference_router::providers::{AuditSink, PolicyDecisionProvider, SigningProvider}; +use kars_inference_router::routes::{ + AppState, foundry_agent_routes, foundry_standalone_routes, inference_routes, +}; + +fn test_state() -> AppState { + let policy_status = Arc::new(PolicyStatusRegistry::new()); + let governance = Arc::new(Governance::new_with_status( + "sb-test", + policy_status.clone(), + )); + AppState { + auth: Arc::new(WorkloadIdentityAuth::new()), + copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), + client: reqwest::Client::new(), + config: Arc::new(Config { + port: 0, + // GitHub Models marketplace endpoint ⇒ foundry_proxy's own + // is_github_models() 501 fires for any request the guard + // lets through — deterministic, no network. + foundry_endpoint: Some("https://models.github.ai/inference".into()), + foundry_project_endpoint: None, + azure_openai_endpoint: None, + default_model: "gpt-4".into(), + content_safety_enabled: false, + prompt_shields_enabled: false, + content_safety_endpoint: None, + token_budget_daily: 1_000_000, + token_budget_per_request: 100_000, + registry_mode: RegistryMode::Local, + registry_url: None, + provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: None, + ollama_endpoint: None, + openai_moderation_endpoint: "https://api.openai.com".into(), + openai_moderation_api_key: Some("sk-mod-test".into()), + openai_moderation_model: "omni-moderation-latest".into(), + }), + budget: TokenBudgetTracker::new(1_000_000, 100_000), + policy_provider: Arc::clone(&governance) as Arc, + audit_sink: Arc::clone(&governance) as Arc, + signing_provider: Arc::clone(&governance) as Arc, + governance, + blocklist: Blocklist::disabled(), + blocked_egress: Arc::new(BlockedBuffer::with_defaults()), + sandbox_name: Arc::new("sb-test".to_string()), + inbox: Arc::new(MeshInbox::new()), + mesh_metrics: Arc::new(MeshMetrics::new()), + model_override: Arc::new(std::sync::RwLock::new(None)), + admin_token: None, + responses_only_models: Arc::new(std::sync::RwLock::new(Default::default())), + handoff_tokens: HandoffTokenStore::new(), + handoff_session: HandoffSession::new(), + drain_state: DrainState::new(), + pending_handoff: PendingHandoffStore::new(), + policy_status, + inference_policy: kars_inference_router::inference_policy_loader::empty_handle(), + memory_binding: kars_inference_router::memory_binding_loader::empty_handle(), + egress_allowlist: kars_inference_router::egress_allowlist_loader::empty_handle(), + deployment_health: Arc::new( + kars_inference_router::deployment_health::DeploymentHealthRegistry::new(), + ), + } +} + +/// Same public wiring as `main.rs` — the guard must hold through real +/// routing, not a handler called directly. +fn app(state: AppState) -> Router { + Router::new() + .merge(inference_routes()) + .merge(foundry_agent_routes()) + .merge(foundry_standalone_routes()) + .with_state(state) +} + +async fn install_policy(state: &AppState, provider: Option<&str>, guardrails: bool) { + let policy = LoadedInferencePolicy { + digest: "sha256:test".into(), + source_path: "/tmp/test-policy".into(), + per_request_tokens: None, + daily_tokens: None, + monthly_tokens: None, + content_safety: Default::default(), + model_preference: None, + provider: provider.map(String::from), + guardrails: if guardrails { + vec![GuardrailStageCfg { + provider: "openai-moderation".into(), + apply_to: ApplyTo::Both, + }] + } else { + vec![] + }, + raw: serde_json::json!({}), + }; + *state.inference_policy.write().await = Some(policy); +} + +async fn send(app: &Router, method: &str, uri: &str) -> (StatusCode, Value) { + let req = Request::builder() + .method(method) + .uri(uri) + .header("content-type", "application/json") + .body(Body::from("{}")) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + let status = resp.status(); + let bytes = axum::body::to_bytes(resp.into_body(), 1_048_576) + .await + .unwrap(); + let json = serde_json::from_slice(&bytes).unwrap_or(Value::Null); + (status, json) +} + +/// Every inference-bearing Foundry route must refuse a guardrail +/// policy — 403 with the same code the sibling `/v1/*` routes use. +#[tokio::test] +async fn guardrail_policy_blocks_inference_bearing_foundry_routes() { + let state = test_state(); + install_policy(&state, None, true).await; + let app = app(state); + + for (method, uri) in [ + ("POST", "/agents"), + ("GET", "/agents/a1/runs/r1"), + ("POST", "/openai/responses"), + ("GET", "/openai/responses/resp_123"), + ("POST", "/openai/conversations"), + ("GET", "/openai/conversations/c1/items"), + ] { + let (status, body) = send(&app, method, uri).await; + assert_eq!( + status, + StatusCode::FORBIDDEN, + "{method} {uri} must fail closed under a guardrail policy, got {status}: {body}" + ); + assert_eq!( + body["error"]["type"].as_str(), + Some("guardrail_route_unsupported"), + "{method} {uri}: {body}" + ); + } +} + +/// A policy selecting a non-Azure provider must also fail closed — +/// these routes only speak to the Azure/Foundry upstream. +#[tokio::test] +async fn non_azure_provider_blocks_inference_bearing_foundry_routes() { + let state = test_state(); + install_policy(&state, Some("anthropic"), false).await; + let app = app(state); + + for uri in ["/agents", "/openai/responses", "/openai/conversations"] { + let (status, body) = send(&app, "POST", uri).await; + assert_eq!( + status, + StatusCode::NOT_IMPLEMENTED, + "POST {uri} must refuse a non-Azure provider, got {status}: {body}" + ); + assert_eq!( + body["error"]["type"].as_str(), + Some("provider_unimplemented"), + "POST {uri}: {body}" + ); + } +} + +/// Back-compat: with no policy loaded (or a plain Azure policy), the +/// guard must not fire — the request reaches the proxy body, which in +/// this fixture terminates in the network-free GitHub Models 501. +#[tokio::test] +async fn plain_policy_leaves_foundry_routes_unguarded() { + let state = test_state(); + install_policy(&state, None, false).await; + let app = app(state); + + let (status, body) = send(&app, "POST", "/openai/responses").await; + assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}"); + assert_eq!( + body["error"]["type"].as_str(), + Some("unsupported_for_provider"), + "expected the proxy's own github-models 501 (proof the guard let it through): {body}" + ); +} + +/// Non-inference Foundry surfaces (storage/management) stay unguarded +/// even under a guardrail policy — enforcement scope is documented in +/// docs/api/crd-reference.md. +#[tokio::test] +async fn guardrail_policy_leaves_non_inference_foundry_routes_unguarded() { + let state = test_state(); + install_policy(&state, None, true).await; + let app = app(state); + + for uri in ["/openai/files", "/memory_stores", "/evaluations"] { + let (status, body) = send(&app, "POST", uri).await; + assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "POST {uri}: {body}"); + assert_eq!( + body["error"]["type"].as_str(), + Some("unsupported_for_provider"), + "expected the proxy's own github-models 501 (proof the guard let it through): {body}" + ); + } +} + +/// Pins the premise of the dot-segment guard: URL parsing (as done by +/// reqwest when the raw path is concatenated into the upstream URL) +/// resolves `.`/`..` segments including percent-encoded `%2e` forms. +/// If this ever stops holding, the traversal rejection below is +/// defense-in-depth rather than load-bearing, but it must hold today. +#[test] +fn url_parsing_normalizes_dot_segments() { + for (raw, normalized) in [ + ("/openai/files/../responses", "/openai/responses"), + ("/openai/files/%2e%2e/responses", "/openai/responses"), + ("/memory_stores/../openai/responses", "/openai/responses"), + ("/evaluations/../openai/responses", "/openai/responses"), + ] { + let url = reqwest::Url::parse(&format!("https://upstream.example{raw}")).unwrap(); + assert_eq!(url.path(), normalized, "raw: {raw}"); + } +} + +/// A dot-segment / encoded-dot / encoded-slash / malformed-escape path +/// routed through an UNGUARDED wildcard must not pivot into a guarded +/// inference path after normalization. These are rejected outright +/// (400) before classification or forwarding, through real router +/// wiring. (Benign empty segments are collapsed, not rejected, so the +/// `//../` case here is rejected on its `..`, not the `//`.) +#[tokio::test] +async fn traversal_paths_are_rejected_not_forwarded() { + let state = test_state(); + install_policy(&state, None, true).await; + let app = app(state); + + for uri in [ + // Literal dot segments through unguarded wildcards. + "/openai/files/../responses", + "/openai/vector_stores/../responses", + "/openai/evals/../responses", + "/memory_stores/../openai/responses", + "/evaluations/../openai/responses", + "/connections/../openai/responses", + "/openai/files/../conversations", + // Percent-encoded dot segments (lower + upper case). + "/openai/files/%2e%2e/responses", + "/openai/files/%2E%2E/responses", + "/memory_stores/%2e%2e/openai/responses", + // Single-dot segment. + "/openai/files/./responses", + // Encoded slash smuggled inside one segment. + "/openai/files/..%2fresponses", + // Double slash + dot segment: rejected on the `..` (the `//` + // itself is collapsed, not rejected). + "/openai/files//../responses", + // Malformed percent escape. + "/openai/files/%zz/responses", + ] { + let (status, body) = send(&app, "POST", uri).await; + assert_eq!( + status, + StatusCode::BAD_REQUEST, + "POST {uri} must be rejected before forwarding, got {status}: {body}" + ); + assert_eq!( + body["error"]["type"].as_str(), + Some("invalid_path"), + "POST {uri}: {body}" + ); + } +} + +/// Double-slash and trailing-slash variants of guarded families must +/// never reach the proxy body, they either 404 at routing or are +/// rejected/blocked at the top of `foundry_proxy`. The one outcome +/// that would be a bug is the github-models 501 marker (proof of +/// reaching the proxy body) or any 2xx. +#[tokio::test] +async fn slash_variants_of_guarded_families_never_reach_proxy_body() { + let state = test_state(); + install_policy(&state, None, true).await; + let app = app(state); + + for uri in [ + "/openai/responses/", + "/openai//responses", + "/agents/", + "/agents//runs", + "/openai/conversations/", + ] { + let (status, body) = send(&app, "POST", uri).await; + assert!( + matches!( + status, + StatusCode::BAD_REQUEST | StatusCode::FORBIDDEN | StatusCode::NOT_FOUND + ), + "POST {uri} must not reach the proxy body, got {status}: {body}" + ); + assert_ne!( + body["error"]["type"].as_str(), + Some("unsupported_for_provider"), + "POST {uri} reached the proxy body (github-models 501 marker): {body}" + ); + } +} + +/// Compatibility: trailing/double-slash management paths that live +/// customers have always used must NOT be rejected by the path guard, +/// even under an active guardrail policy (they are exempt) and even +/// with no policy at all. They collapse to canonical form and reach the +/// proxy body (here the network-free github-models 501). Paths carry a +/// real segment before the slash so they match the `{*path}` wildcard +/// route (a bare `/openai/files/` 404s at axum routing, which is +/// pre-existing and orthogonal to the guard). +#[tokio::test] +async fn benign_trailing_slash_paths_are_not_rejected() { + for guardrails in [false, true] { + let state = test_state(); + install_policy(&state, None, guardrails).await; + let app = app(state); + for uri in [ + "/openai/files/f1/", // trailing slash + "/memory_stores/s1/", // trailing slash + "/openai/files//f1", // double slash + "/evaluations/e1/", // trailing slash + ] { + let (status, body) = send(&app, "POST", uri).await; + assert_ne!( + body["error"]["type"].as_str(), + Some("invalid_path"), + "POST {uri} (guardrails={guardrails}) must not be rejected by the path guard: {status} {body}" + ); + assert_eq!( + body["error"]["type"].as_str(), + Some("unsupported_for_provider"), + "POST {uri} (guardrails={guardrails}) should reach the proxy body: {status} {body}" + ); + } + } +} + +/// Default-deny: with the guard inverted to an exempt list, an +/// unknown-but-routed path family fails closed rather than slipping +/// through. `/agents*` wildcards cover arbitrary suffixes, so use a +/// suffix no explicit rule names. +#[tokio::test] +async fn unknown_wildcard_suffixes_fail_closed() { + let state = test_state(); + install_policy(&state, None, true).await; + let app = app(state); + + let (status, body) = send(&app, "POST", "/agents/a1/some/new/api").await; + assert_eq!(status, StatusCode::FORBIDDEN, "{body}"); + assert_eq!( + body["error"]["type"].as_str(), + Some("guardrail_route_unsupported"), + "{body}" + ); +} diff --git a/inference-router/tests/multi_provider_guardrails.rs b/inference-router/tests/multi_provider_guardrails.rs new file mode 100644 index 000000000..e839db6c8 --- /dev/null +++ b/inference-router/tests/multi_provider_guardrails.rs @@ -0,0 +1,259 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! End-to-end tests for the multi-provider slice: `proxy::forward` +//! against fake Anthropic / Ollama upstreams, and the OpenAI +//! Moderation guardrail against a fake moderation endpoint. +//! +//! No env mutation — provider endpoints and credentials are injected +//! via `UpstreamConfig` / `Config` struct literals, which is exactly +//! how the production path receives them after +//! `routes::apply_provider_resolution`. + +use axum::http::{HeaderMap, Method}; +use bytes::Bytes; +use kars_inference_router::auth::WorkloadIdentityAuth; +use kars_inference_router::config::{Config, RegistryMode}; +use kars_inference_router::guardrails::{ApplyTo, Direction, GuardrailPipeline, GuardrailStageCfg}; +use kars_inference_router::provider::ProviderKind; +use kars_inference_router::proxy::{UpstreamConfig, forward}; +use wiremock::matchers::{method, path}; +use wiremock::{Mock, MockServer, ResponseTemplate}; + +fn config_with_moderation(endpoint: &str, api_key: Option<&str>) -> Config { + Config { + port: 8443, + foundry_endpoint: None, + foundry_project_endpoint: None, + azure_openai_endpoint: None, + default_model: "gpt-4o-mini".into(), + content_safety_enabled: false, + prompt_shields_enabled: false, + content_safety_endpoint: None, + token_budget_daily: 0, + token_budget_per_request: 0, + registry_mode: RegistryMode::Local, + registry_url: None, + provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: None, + ollama_endpoint: None, + openai_moderation_endpoint: endpoint.to_string(), + openai_moderation_api_key: api_key.map(String::from), + openai_moderation_model: "omni-moderation-latest".into(), + } +} + +#[tokio::test] +async fn ollama_provider_forwards_openai_compat_without_auth() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/chat/completions")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "chatcmpl-ollama", + "choices": [{ "message": { "role": "assistant", "content": "hi" }, + "finish_reason": "stop" }], + "usage": { "prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5 } + }))) + .expect(1) + .mount(&server) + .await; + + let upstream = UpstreamConfig { + endpoint: server.uri(), + deployment: "llama3.1".into(), + sandbox_name: "test-sandbox".into(), + provider: ProviderKind::Ollama, + api_key: None, + }; + + let (status, _headers, resp) = forward( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::new(), + &upstream, + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from(r#"{"messages":[{"role":"user","content":"hello"}]}"#), + ) + .await + .expect("forward to fake ollama"); + + assert_eq!(status.as_u16(), 200); + let v: serde_json::Value = serde_json::from_slice(&resp).unwrap(); + assert_eq!(v["choices"][0]["message"]["content"], "hi"); + + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + let req = &requests[0]; + assert_eq!(req.url.path(), "/v1/chat/completions"); + assert!( + !req.headers.contains_key("authorization") && !req.headers.contains_key("x-api-key"), + "ollama upstream must receive no credentials" + ); + // Deployment injected as the model. + let body: serde_json::Value = serde_json::from_slice(&req.body).unwrap(); + assert_eq!(body["model"], "llama3.1"); +} + +#[tokio::test] +async fn anthropic_provider_forwards_messages_with_router_held_key() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "id": "msg_test", + "type": "message", + "role": "assistant", + "content": [{ "type": "text", "text": "hello back" }], + "stop_reason": "end_turn", + "usage": { "input_tokens": 4, "output_tokens": 3 } + }))) + .expect(1) + .mount(&server) + .await; + + let upstream = UpstreamConfig { + endpoint: server.uri(), + deployment: "claude-sonnet-4-5".into(), + sandbox_name: "test-sandbox".into(), + provider: ProviderKind::Anthropic, + api_key: Some("sk-ant-router-held".into()), + }; + + // The inbound request carries an agent-supplied x-api-key that + // must be stripped — only the router-held key may reach upstream. + let mut inbound = HeaderMap::new(); + inbound.insert("x-api-key", "agent-smuggled-key".parse().unwrap()); + + let (status, _headers, _resp) = forward( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::new(), + &upstream, + Method::POST, + "v1/messages", + &inbound, + Bytes::from(r#"{"model":"claude-sonnet-4-5","max_tokens":64,"messages":[{"role":"user","content":"hi"}]}"#), + ) + .await + .expect("forward to fake anthropic"); + + assert_eq!(status.as_u16(), 200); + let requests = server.received_requests().await.unwrap(); + assert_eq!(requests.len(), 1); + let req = &requests[0]; + assert_eq!(req.url.path(), "/v1/messages"); + assert_eq!( + req.headers + .get("x-api-key") + .and_then(|v| v.to_str().ok()) + .unwrap_or_default(), + "sk-ant-router-held", + "router-held key must replace any agent-supplied key" + ); + assert_eq!( + req.headers + .get("anthropic-version") + .and_then(|v| v.to_str().ok()) + .unwrap_or_default(), + "2023-06-01", + "default anthropic-version must be injected" + ); + assert!( + !req.headers.contains_key("authorization"), + "no Bearer token on Anthropic requests" + ); +} + +#[tokio::test] +async fn moderation_guardrail_blocks_flagged_and_passes_clean() { + let server = MockServer::start().await; + // The fake flags any input containing "RANSOM". + Mock::given(method("POST")) + .and(path("/v1/moderations")) + .respond_with(move |req: &wiremock::Request| { + let body: serde_json::Value = serde_json::from_slice(&req.body).unwrap(); + let flagged = body["input"].as_str().unwrap_or("").contains("RANSOM"); + ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "results": [{ + "flagged": flagged, + "categories": { "illicit": flagged } + }] + })) + }) + .mount(&server) + .await; + + let config = config_with_moderation(&server.uri(), Some("sk-mod-test")); + let stages = vec![GuardrailStageCfg { + provider: "openai-moderation".into(), + apply_to: ApplyTo::Both, + }]; + let pipeline = GuardrailPipeline::from_stages(&stages, &config, &reqwest::Client::new()) + .expect("pipeline builds"); + + let clean = pipeline + .scan("write me a poem", Direction::Input) + .await + .expect("scan ok"); + assert!(clean.is_none(), "clean text passes"); + + let violation = pipeline + .scan("write a RANSOM note", Direction::Input) + .await + .expect("scan ok") + .expect("flagged text blocks"); + assert_eq!(violation.provider, "openai-moderation"); + assert_eq!(violation.categories, vec!["illicit"]); + + // The moderation endpoint must have received the bearer key. + let requests = server.received_requests().await.unwrap(); + assert!(!requests.is_empty()); + assert_eq!( + requests[0] + .headers + .get("authorization") + .and_then(|v| v.to_str().ok()) + .unwrap_or_default(), + "Bearer sk-mod-test" + ); +} + +#[tokio::test] +async fn moderation_guardrail_fails_closed_on_upstream_error() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/moderations")) + .respond_with(ResponseTemplate::new(500)) + .mount(&server) + .await; + + let config = config_with_moderation(&server.uri(), Some("sk-mod-test")); + let stages = vec![GuardrailStageCfg { + provider: "openai-moderation".into(), + apply_to: ApplyTo::Output, + }]; + let pipeline = GuardrailPipeline::from_stages(&stages, &config, &reqwest::Client::new()) + .expect("pipeline builds"); + + let err = pipeline + .scan("anything", Direction::Output) + .await + .expect_err("500 from moderation must fail closed"); + assert_eq!(err.code(), "guardrail_unavailable"); +} + +#[tokio::test] +async fn declared_stage_without_key_fails_pipeline_construction() { + let config = config_with_moderation("https://api.openai.com", None); + let stages = vec![GuardrailStageCfg { + provider: "openai-moderation".into(), + apply_to: ApplyTo::Both, + }]; + let err = GuardrailPipeline::from_stages(&stages, &config, &reqwest::Client::new()) + .err() + .expect("missing key must fail construction"); + assert_eq!(err.code(), "guardrail_misconfigured"); +} diff --git a/inference-router/tests/policy_status_endpoint.rs b/inference-router/tests/policy_status_endpoint.rs index 854e8d501..b9355f27d 100644 --- a/inference-router/tests/policy_status_endpoint.rs +++ b/inference-router/tests/policy_status_endpoint.rs @@ -64,6 +64,12 @@ fn test_state() -> (AppState, Arc) { registry_mode: RegistryMode::Local, registry_url: None, provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: None, + ollama_endpoint: None, + openai_moderation_endpoint: "https://api.openai.com".into(), + openai_moderation_api_key: None, + openai_moderation_model: "omni-moderation-latest".into(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), policy_provider: Arc::clone(&governance) as Arc, diff --git a/inference-router/tests/proxy_fake_upstream.rs b/inference-router/tests/proxy_fake_upstream.rs index 910966627..28c92f13c 100644 --- a/inference-router/tests/proxy_fake_upstream.rs +++ b/inference-router/tests/proxy_fake_upstream.rs @@ -22,6 +22,7 @@ use axum::http::{HeaderMap, Method}; use bytes::Bytes; use common::{FakeAd, FakeAzure, FakeImds, FixtureRoute}; use kars_inference_router::auth::WorkloadIdentityAuth; +use kars_inference_router::provider::ProviderKind; use kars_inference_router::proxy::{UpstreamConfig, forward}; use std::sync::Mutex; @@ -71,6 +72,8 @@ async fn api_key_mode_proxies_chat_completion_with_filter_results() { endpoint: azure.base_url(), deployment: "gpt-4o".to_string(), sandbox_name: "test-sandbox".to_string(), + provider: ProviderKind::AzureOpenAI, + api_key: None, }; let client = reqwest::Client::new(); @@ -149,6 +152,8 @@ async fn wi_mode_falls_back_to_imds_and_proxies_embeddings() { endpoint: azure.base_url(), deployment: "text-embedding-3-small".to_string(), sandbox_name: "test-sandbox-wi".to_string(), + provider: ProviderKind::AzureOpenAI, + api_key: None, }; let client = reqwest::Client::new(); let body = Bytes::from(r#"{"input":"hello"}"#.as_bytes().to_vec()); @@ -218,6 +223,8 @@ async fn upstream_error_status_is_propagated() { endpoint: azure.base_url(), deployment: "gpt-4o".to_string(), sandbox_name: "test-sandbox-429".to_string(), + provider: ProviderKind::AzureOpenAI, + api_key: None, }; let client = reqwest::Client::new();