From b277cc750b1cf6b5080da18d54882be211960a01 Mon Sep 17 00:00:00 2001 From: johnseong Date: Thu, 30 Jul 2026 20:18:56 -0400 Subject: [PATCH 01/18] feat(inference): multi-provider LLM upstreams + pluggable guardrail pipeline MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Roadmap slice 1 of 'multi-cloud LLM providers + native guardrails': Anthropic + Ollama providers and an OpenAI Moderation guardrail stage, policy-driven via InferencePolicy. Provider credentials stay on the router sidecar (secret mount / env) — the agent process never sees them. Controller: - InferencePolicy spec.provider (typed InferenceProvider enum with kebab-case wire tags matching the existing ModelRef.provider strings) and spec.guardrails[] (openai-moderation, applyTo input|output|both). - Compile step emits 'provider' + 'guardrails' in the compiled profile; CEL + reconciler keep both mutually exclusive with bundleRef. - Reconciler forwards ANTHROPIC_API_KEY/ANTHROPIC_ENDPOINT/ OLLAMA_ENDPOINT/OPENAI_MODERATION_* to router sidecars when set; helm CRD template regenerated. Router: - New provider module: fail-closed resolution (unimplemented bedrock -> 501, missing endpoint/credential -> 503; never a silent Azure fallback). Provider-aware UpstreamConfig, URL shapes and auth schemes (Anthropic x-api-key + anthropic-version; Ollama unauthenticated OpenAI-compat under /v1/). - Anthropic Messages native pass-through (streaming + tool use) on /anthropic/v1/messages; Ollama chat completions buffered + SSE. - New guardrails module: Guardrail trait + OpenAI Moderation backend; input pre-flight, buffered output, and hold-and-release SSE scanning (no model text delivered before a scan covers it). Fail-closed on misconfigured stages and backend outages; kars_guardrail_scans_total metric. Tests: 857 controller + 997 router unit tests green; new wiremock integration suite (fake Anthropic/Ollama/moderation upstreams). --- CHANGELOG.md | 72 + controller/src/config_hash.rs | 6 + controller/src/crd_validations.rs | 13 +- controller/src/inference_policy.rs | 119 ++ controller/src/inference_policy_compile.rs | 90 +- controller/src/inference_policy_reconciler.rs | 13 +- controller/src/reconciler/mod.rs | 43 + .../kars/templates/crd-inferencepolicy.yaml | 72 +- docs/api/crd-reference.md | 12 +- docs/architecture.md | 2 +- inference-router/src/config.rs | 87 ++ inference-router/src/failover.rs | 2 + inference-router/src/guardrails.rs | 1306 +++++++++++++++++ .../src/inference_policy_loader.rs | 114 ++ inference-router/src/lib.rs | 2 + inference-router/src/metrics.rs | 13 + inference-router/src/provider.rs | 344 +++++ inference-router/src/proxy.rs | 166 ++- .../src/routes/anthropic_messages.rs | 204 ++- .../src/routes/chat_completions.rs | 291 +++- inference-router/src/routes/mod.rs | 53 +- .../tests/agt_governance_integration.rs | 6 + .../tests/egress_blocked_endpoint.rs | 6 + inference-router/tests/failover_walk.rs | 7 + .../tests/multi_provider_guardrails.rs | 259 ++++ .../tests/policy_status_endpoint.rs | 6 + inference-router/tests/proxy_fake_upstream.rs | 7 + 27 files changed, 3265 insertions(+), 50 deletions(-) create mode 100644 inference-router/src/guardrails.rs create mode 100644 inference-router/src/provider.rs create mode 100644 inference-router/tests/multi_provider_guardrails.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index a22c06412..a8c87f12f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,78 @@ 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 + (`modelPreference.primary.provider` wins over `spec.provider`; + unknown tags warn and fall through; missing endpoint/credential → + 503, unimplemented `bedrock` → 501 — never a silent Azure fallback). +- `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; usage metered from Anthropic + `usage` fields. 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. + +**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 at every governed exchange: request pre-flight (input + stages), buffered responses (including the Responses-API recovery + paths), and SSE streams via **hold-and-release** windows + (`GUARDRAIL_STREAM_SCAN_CHARS`, default 1000 chars) — no model text + is delivered before a scan has covered it; flagged streams are cut + with a structured SSE error frame + `data: [DONE]`. +- Fail-closed contract: declared-but-unbuildable stages reject the + request (503 `guardrail_misconfigured`); backend outages block (502 + `guardrail_unavailable`); scan-text truncation (16k-char cap) is + logged, never silent. 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..f49dc7d82 100644 --- a/controller/src/config_hash.rs +++ b/controller/src/config_hash.rs @@ -38,6 +38,12 @@ 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 (credentials deliberately excluded — + // this list never hashes secret material, matching 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..f027c6ca7 100644 --- a/controller/src/crd_validations.rs +++ b/controller/src/crd_validations.rs @@ -277,8 +277,17 @@ 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() + }, + // Guardrail pipeline shape: an explicitly-empty list is an + // authoring mistake (delete the field to mean "no pipeline"), + // and a small cap keeps per-request scan fan-out bounded. + 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..a28056093 100644 --- a/controller/src/inference_policy.rs +++ b/controller/src/inference_policy.rs @@ -54,6 +54,59 @@ use serde::{Deserialize, Serialize}; use crate::mcp_server::LocalObjectRef; +/// Inference provider selector — the typed form of the kebab-case +/// provider tags this CRD has always documented on [`ModelRef`] +/// (`azure-openai` / `anthropic` / `bedrock` / `ollama`). +/// +/// Serialized with explicit kebab-case renames (NOT `rename_all`) so +/// the wire tags match the existing free-form `ModelRef.provider` +/// strings byte-for-byte — a CR that says `provider: anthropic` under +/// `modelPreference` today can move to the typed field without a +/// migration. +/// +/// Router-side consumption (Slice: multi-provider): `azure-openai` +/// keeps the env-configured Foundry/AOAI upstream (back-compat +/// default); `anthropic` targets the Anthropic Messages API +/// (`ANTHROPIC_ENDPOINT`, key from the router-side secret mount — +/// never visible to the agent process); `ollama` targets an +/// OpenAI-compatible Ollama server (`OLLAMA_ENDPOINT`, no auth). +/// `bedrock` is accepted by the schema for forward-compat but the +/// router rejects it with a clear 501 until the Bedrock client lands +/// — declaring it here keeps the CRD stable across that slice. +#[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 — same string serde emits. Exposed for + /// log lines and compile-step JSON so call sites never hand-roll + /// the mapping. + #[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 +155,27 @@ pub struct InferencePolicySpec { /// [`Self::bundle_ref`]. pub model_preference: Option, + /// Default inference provider for call sites this policy governs. + /// Optional — absent ⇒ the router keeps its env-configured Azure + /// OpenAI / Foundry upstream (back-compat). When set to a + /// non-Azure provider the router swaps the upstream base URL and + /// auth scheme accordingly; provider credentials stay inside the + /// router sidecar (secret mount / env), never in the agent + /// process. `modelPreference.primary.provider`, when it names a + /// recognised tag, takes precedence over this field so a fallback + /// chain can pin its own route. Mutually exclusive with + /// [`Self::bundle_ref`]. + pub provider: Option, + + /// Ordered guardrail pipeline stages the router runs around each + /// inference call (request pre-flight and response — buffered and + /// streaming). Optional — absent ⇒ only the Phase 1 substrate + /// (Foundry guardrail annotations + `contentSafety` floors) + /// applies. Stages run in declaration order; the first stage that + /// flags content blocks the call. 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 +302,51 @@ pub struct ModelRef { pub deployment: String, } +/// A single stage of the router-side guardrail pipeline. The router +/// materialises each stage into a scanner (network client + policy) +/// at policy load time; a stage whose backend is not configured on +/// the router (e.g. missing moderation API key) fails the *request* +/// closed with an explicit error rather than silently skipping — a +/// declared guardrail that cannot run must never be an open gate. +#[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. One variant today; Bedrock Guardrails and +/// Model Armor are declared roadmap follow-ups and will extend this +/// enum (adding a variant is a non-breaking CRD change). +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, JsonSchema)] +pub enum GuardrailProvider { + /// OpenAI Moderation API (`omni-moderation-latest`). Router-side + /// key via `OPENAI_MODERATION_API_KEY` (falls back to + /// `OPENAI_API_KEY`); endpoint override via + /// `OPENAI_MODERATION_ENDPOINT` for Azure-hosted equivalents. + #[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..28310c007 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,32 @@ pub fn compile_to_profile(spec: &InferencePolicySpec) -> Value { }) }); + // Provider + guardrails travel as the same kebab-case wire tags + // the CRD serde emits (`InferenceProvider::as_tag` / + // `GuardrailProvider` renames) so the router-side loader parses + // one vocabulary for both the typed field and the free-form + // `modelPreference.*.provider` strings. + 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 +218,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 +254,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 +278,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 +308,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..fda976b30 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,14 @@ fn merge_bundle_with_selector( token_budget, content_safety, model_preference, + // The signed-bundle canonical format + // (`policy_canonical::inference`) does not carry + // `provider`/`guardrails` yet — extending that wire contract + // is a coordinated change with the bundle tooling. Until + // then, 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..a288519ed 100644 --- a/controller/src/reconciler/mod.rs +++ b/controller/src/reconciler/mod.rs @@ -311,6 +311,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 +1948,24 @@ 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/deploy/helm/kars/templates/crd-inferencepolicy.yaml b/deploy/helm/kars/templates/crd-inferencepolicy.yaml index fceb79c29..1dd9bb4ec 100644 --- a/deploy/helm/kars/templates/crd-inferencepolicy.yaml +++ b/deploy/helm/kars/templates/crd-inferencepolicy.yaml @@ -181,6 +181,42 @@ spec: comes from the signed bundle. nullable: true type: string + guardrails: + description: |- + Ordered guardrail pipeline stages the router runs around each + inference call (request pre-flight and response — buffered and + streaming). Optional — absent ⇒ only the Phase 1 substrate + (Foundry guardrail annotations + `contentSafety` floors) + applies. Stages run in declaration order; the first stage that + flags content blocks the call. Mutually exclusive with + [`Self::bundle_ref`]. + items: + description: |- + A single stage of the router-side guardrail pipeline. The router + materialises each stage into a scanner (network client + policy) + at policy load time; a stage whose backend is not configured on + the router (e.g. missing moderation API key) fails the *request* + closed with an explicit error rather than silently skipping — a + declared guardrail that cannot run must never be an open gate. + 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 +262,34 @@ spec: required: - primary type: object + provider: + description: |- + Inference provider selector — the typed form of the kebab-case + provider tags this CRD has always documented on [`ModelRef`] + (`azure-openai` / `anthropic` / `bedrock` / `ollama`). + + Serialized with explicit kebab-case renames (NOT `rename_all`) so + the wire tags match the existing free-form `ModelRef.provider` + strings byte-for-byte — a CR that says `provider: anthropic` under + `modelPreference` today can move to the typed field without a + migration. + + Router-side consumption (Slice: multi-provider): `azure-openai` + keeps the env-configured Foundry/AOAI upstream (back-compat + default); `anthropic` targets the Anthropic Messages API + (`ANTHROPIC_ENDPOINT`, key from the router-side secret mount — + never visible to the agent process); `ollama` targets an + OpenAI-compatible Ollama server (`OLLAMA_ENDPOINT`, no auth). + `bedrock` is accepted by the schema for forward-compat but the + router rejects it with a clear 501 until the Bedrock client lands + — declaring it here keeps the CRD stable across that slice. + enum: + - azure-openai + - anthropic + - ollama + - bedrock + nullable: true + type: string tokenBudget: description: |- Token-budget caps. Optional — absent ⇒ no budget enforcement. @@ -262,9 +326,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.displayName))' + 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.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 +470,3 @@ spec: storage: true subresources: status: {} - diff --git a/docs/api/crd-reference.md b/docs/api/crd-reference.md index 3f7c49029..51f22315a 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,15 @@ 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`). 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. Credentials/endpoints are router-sidecar config only — never visible to the agent. | +| `spec.guardrails[]` | Optional ordered guardrail pipeline (1–8 stages) the router runs around each governed call. 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 model text reaches the client before a scan has covered it. | +| `spec.modelPreference.primary` | `{provider, deployment}`. `provider` is one of `azure-openai`, `anthropic`, `gemini`, `bedrock`, `ollama`. A recognised tag here overrides `spec.provider` for routing. | +| `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. | > **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/inference-router/src/config.rs b/inference-router/src/config.rs index 2707674a6..d81300b01 100644 --- a/inference-router/src/config.rs +++ b/inference-router/src/config.rs @@ -81,6 +81,55 @@ 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 var or the secret + /// mounts `/etc/kars/secrets/anthropic-api-key` / + /// `/run/secrets/anthropic-api-key`. Lives ONLY in the router + /// sidecar; the agent process talks to the localhost proxy and + /// never sees it. `None` ⇒ policies selecting Anthropic fail + /// closed with an operator-actionable error. + 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, + + /// API key for the OpenAI Moderation guardrail — + /// `OPENAI_MODERATION_API_KEY` env var (falls back to + /// `OPENAI_API_KEY`, then the secret mounts + /// `/etc/kars/secrets/openai-moderation-api-key` / + /// `/run/secrets/openai-moderation-api-key`). `None` ⇒ policies + /// declaring an `openai-moderation` guardrail stage 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 +194,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 +293,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/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.rs b/inference-router/src/guardrails.rs new file mode 100644 index 000000000..f415c4451 --- /dev/null +++ b/inference-router/src/guardrails.rs @@ -0,0 +1,1306 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Pluggable guardrail pipeline (multi-cloud guardrails slice). +//! +//! An `InferencePolicy` can declare an ordered list of guardrail +//! stages (`spec.guardrails[]`) that the router runs around every +//! inference call it governs — on the request text before the +//! upstream forward, and on the response text both buffered and +//! streaming. The first backend today is the OpenAI Moderation API; +//! Bedrock Guardrails and Model Armor extend the same [`Guardrail`] +//! trait in follow-up slices. +//! +//! ## Fail-closed contract +//! +//! A *declared* guardrail that cannot run must never become an open +//! gate: +//! +//! - A stage whose backend the router does not recognise, or whose +//! credential/endpoint is missing, fails pipeline construction — +//! the handler rejects the request with an operator-actionable +//! error before any prompt bytes leave the pod. +//! - A scan that errors at runtime (transport failure, non-2xx, +//! unparseable verdict) blocks the request with +//! `guardrail_unavailable` rather than passing unscanned content. +//! +//! ## Streaming semantics (hold-and-release) +//! +//! SSE responses are guarded with a hold-and-release window: chunks +//! are buffered until the accumulated new text reaches +//! [`STREAM_SCAN_THRESHOLD_CHARS`] (or the stream ends), the +//! accumulated text is scanned, and only then is the held window +//! released to the client. No model output is ever delivered before +//! some scan has covered it. On a flagged scan, the client receives a +//! structured SSE error frame + `data: [DONE]` and the upstream +//! stream is dropped. The cost is scan-sized delivery granularity +//! (one moderation round-trip per window), which is the standard +//! trade-off for streaming guardrails. +//! +//! Scanned text is capped at [`MAX_SCAN_CHARS`] (most recent chars +//! for output, leading chars for input) to stay under moderation +//! input limits; truncation is logged at WARN — never silent. + +use std::sync::Arc; + +use async_trait::async_trait; +use bytes::Bytes; +use futures::stream::BoxStream; +use futures::stream::StreamExt; +use reqwest::Client; + +use crate::config::Config; +use crate::metrics; + +/// 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 stage as it travels through the compiled policy JSON +/// (`{"provider": "...", "applyTo": "..." | null}`). Parsed liberally +/// by the loader; strictness (unknown backend ⇒ fail closed) applies +/// at pipeline *construction*, where a request is available to +/// reject. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct GuardrailStageCfg { + pub provider: String, + pub apply_to: ApplyTo, +} + +impl GuardrailStageCfg { + /// Parse the compiled `guardrails` block (array | null | absent). + /// Entries without a string `provider` are dropped with a WARN — + /// they cannot be built into anything enforceable and the + /// controller schema rejects them at admission anyway. + #[must_use] + pub fn from_compiled_json(v: &serde_json::Value) -> Vec { + let Some(arr) = v.as_array() else { + return Vec::new(); + }; + arr.iter() + .filter_map(|stage| { + let Some(provider) = stage.get("provider").and_then(|p| p.as_str()) else { + tracing::warn!( + stage = %stage, + "guardrail stage missing string 'provider' — dropped" + ); + return None; + }; + Some(Self { + provider: provider.to_string(), + apply_to: ApplyTo::parse(stage.get("applyTo").and_then(|a| a.as_str())), + }) + }) + .collect() + } +} + +/// 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; +} + +// ─── 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 ──────────────────────────────────────────────────────────────── + +struct BuiltStage { + apply_to: ApplyTo, + guard: Arc, +} + +/// Ordered guardrail stages materialised from a policy snapshot. +/// Cheap to build per request (clones a shared `reqwest::Client`). +pub struct GuardrailPipeline { + 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`, in order. + /// First flagged verdict wins. Empty text short-circuits to pass. + pub async fn scan( + &self, + text: &str, + direction: Direction, + ) -> Result, GuardrailError> { + if text.is_empty() { + return Ok(None); + } + let capped = cap_for_scan(text, direction); + for stage in self.stages.iter().filter(|s| s.apply_to.covers(direction)) { + let outcome = stage.guard.scan(capped).await; + match outcome { + 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) + } +} + +/// Cap text to [`MAX_SCAN_CHARS`]: leading chars for input (the +/// system prompt + earliest instructions), trailing chars for output +/// (the newest generated text — earlier output was already scanned by +/// previous windows in the streaming path). Logs at WARN on +/// truncation. +fn cap_for_scan(text: &str, direction: Direction) -> &str { + if text.chars().count() <= MAX_SCAN_CHARS { + return text; + } + tracing::warn!( + direction = direction.as_str(), + total_chars = text.chars().count(), + scanned_chars = MAX_SCAN_CHARS, + "guardrail scan text exceeds cap — scanning a truncated window" + ); + match direction { + Direction::Input => { + let end = text + .char_indices() + .nth(MAX_SCAN_CHARS) + .map_or(text.len(), |(i, _)| i); + &text[..end] + } + Direction::Output => { + let start = text + .char_indices() + .rev() + .nth(MAX_SCAN_CHARS - 1) + .map_or(0, |(i, _)| i); + &text[start..] + } + } +} + +// ─── 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") +} + +// ─── 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] +fn delta_text_from_event(dialect: StreamDialect, event: &serde_json::Value) -> Option { + match dialect { + StreamDialect::OpenAiChat => event + .get("choices") + .and_then(|c| c.as_array()) + .and_then(|c| c.first()) + .and_then(|c| c.get("delta")) + .and_then(|d| d.get("content")) + .and_then(|t| t.as_str()) + .filter(|t| !t.is_empty()) + .map(str::to_string), + 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 + } + } + } +} + +/// The client-facing SSE error frame emitted when a stream is cut by +/// a guardrail. OpenAI-style error object works for both dialects' +/// SDK error paths and is what the existing content-safety stream cut +/// emits too. +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() + } + }) + )) +} + +fn unavailable_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. +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) +} + +/// 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`]. +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 `data:` lines split across chunk boundaries. + line_carry: String, + /// All delta text accumulated so far (scan context). + accumulated: String, + /// Chars of `accumulated` not yet covered by a scan. + unscanned: usize, +} + +/// What the state machine wants the adaptor to emit next. +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 { + fn new(pipeline: Arc, dialect: StreamDialect, threshold: usize) -> Self { + Self { + pipeline, + dialect, + threshold, + held: Vec::new(), + line_carry: String::new(), + accumulated: String::new(), + unscanned: 0, + } + } + + /// Pull complete lines out of `chunk` (+ carry), extract delta + /// text, and account it as unscanned. + fn ingest_text(&mut self, chunk: &[u8]) { + self.line_carry.push_str(&String::from_utf8_lossy(chunk)); + // Keep the trailing partial line (no '\n' yet) in the carry. + 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() { + let Some(payload) = line.trim().strip_prefix("data: ") else { + continue; + }; + if payload == "[DONE]" { + continue; + } + if let Ok(event) = serde_json::from_str::(payload) + && let Some(text) = delta_text_from_event(self.dialect, &event) + { + self.unscanned += text.chars().count(); + self.accumulated.push_str(&text); + } + } + } + + async fn on_chunk(&mut self, chunk: Bytes) -> SseGuardStep { + self.ingest_text(&chunk); + self.held.push(chunk); + if self.unscanned < self.threshold { + // Fast path: window not full. Chunks carrying no delta + // text at all (keepalives, role/annotation frames) are + // safe to release immediately when nothing text-bearing + // is being held alongside them. + 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 { + 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; + SseGuardStep::Release(std::mem::take(&mut self.held)) + } + Ok(Some(violation)) => SseGuardStep::Cut(violation_sse_frame(&violation)), + Err(e) => SseGuardStep::Cut(unavailable_sse_frame(&e)), + } + } +} + +/// 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) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + // ---- 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 }, + { "applyTo": "input" } // dropped: no provider + ]); + let stages = GuardrailStageCfg::from_compiled_json(&v); + 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_handles_null_and_absent() { + assert!(GuardrailStageCfg::from_compiled_json(&serde_json::Value::Null).is_empty()); + assert!(GuardrailStageCfg::from_compiled_json(&serde_json::json!({})).is_empty()); + } + + // ---- 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 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 + ); + } + + // ---- cap ---- + + #[test] + fn cap_keeps_short_text_intact() { + assert_eq!(cap_for_scan("short", Direction::Input), "short"); + } + + #[test] + fn cap_truncates_head_for_input_and_tail_for_output() { + let long: String = "a".repeat(MAX_SCAN_CHARS) + "TAIL"; + let capped_in = cap_for_scan(&long, Direction::Input); + assert_eq!(capped_in.len(), MAX_SCAN_CHARS); + assert!(capped_in.starts_with('a') && !capped_in.contains("TAIL")); + let long2 = "HEAD".to_string() + &"b".repeat(MAX_SCAN_CHARS); + let capped_out = cap_for_scan(&long2, Direction::Output); + assert_eq!(capped_out.len(), MAX_SCAN_CHARS); + assert!(!capped_out.contains("HEAD")); + } + + // ---- pipeline + streaming with a fake backend ---- + + /// Test backend: flags any text containing the marker. Counts + /// scans so tests can assert hold-and-release windowing. + 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_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")); + } +} diff --git a/inference-router/src/inference_policy_loader.rs b/inference-router/src/inference_policy_loader.rs index 30defe06c..1e8b0ceed 100644 --- a/inference-router/src/inference_policy_loader.rs +++ b/inference-router/src/inference_policy_loader.rs @@ -165,6 +165,22 @@ pub struct LoadedInferencePolicy { /// back to the env-driven default deployment (back-compat). pub model_preference: Option, + /// `spec.provider` — multi-provider slice. Raw kebab-case tag + /// (`azure-openai` / `anthropic` / `ollama` / `bedrock`); + /// interpretation (including the unimplemented-provider + /// fail-closed path) happens per request in + /// [`crate::provider::resolve`] so a policy naming a provider + /// this build can't serve degrades that *request*, not the whole + /// policy load. `None` ⇒ env-driven Azure upstream (back-compat). + pub provider: Option, + + /// `spec.guardrails[]` — pluggable guardrail pipeline stages. + /// Empty when the CR omits the block. Stage validity (known + /// backend, credential present) is checked at pipeline build + /// time in [`crate::guardrails::GuardrailPipeline::from_stages`], + /// where a request exists to fail closed. + 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 +343,19 @@ 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); + let guardrails = crate::guardrails::GuardrailStageCfg::from_compiled_json( + parsed.get("guardrails").unwrap_or(&serde_json::Value::Null), + ); + // 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 +371,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 +414,8 @@ pub fn load_inference_policy_from_dir( monthly_tokens, content_safety, model_preference, + provider, + guardrails, raw: parsed, }) } @@ -507,6 +540,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 +569,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 +785,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 +808,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..6bf7871db --- /dev/null +++ b/inference-router/src/provider.rs @@ -0,0 +1,344 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Multi-provider upstream resolution. +//! +//! Maps the `InferencePolicy` provider tags (`azure-openai` / +//! `anthropic` / `ollama` / `bedrock`) onto a concrete upstream +//! target: base URL shape + auth scheme. The provider tag travels in +//! the compiled policy JSON (`spec.provider`, and +//! `spec.modelPreference.*.provider`); the *credentials and endpoints* +//! come exclusively from the router's own environment / secret mounts +//! — the agent process never sees a provider API key, exactly as with +//! the Azure Workload Identity path. +//! +//! ## Resolution precedence +//! +//! 1. `modelPreference.primary.provider`, when it parses to a known +//! tag — an explicit route preference wins over the policy-level +//! default. +//! 2. `spec.provider` (policy-level default). +//! 3. `azure-openai` (absent / unknown tags — matches the pre-slice +//! behaviour where provider tags were informational-only). +//! +//! `bedrock` is recognised but not yet implemented: a policy that +//! declares it gets an explicit 501-style error instead of a silent +//! reroute to Azure — declared intent must never be silently ignored. +//! +//! ## Failure semantics +//! +//! A resolvable provider with missing router-side configuration +//! (no `ANTHROPIC_API_KEY`, no `OLLAMA_ENDPOINT`) fails the request +//! closed with a specific, operator-actionable error. Falling back to +//! the Azure upstream would send prompts to a provider the operator +//! didn't select. + +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 effective provider for a request. +/// +/// `model_pref_tag` is `modelPreference.primary.provider` (when a +/// policy carries a model preference), `policy_tag` is the top-level +/// `spec.provider`. See module docs for precedence. Unknown tags log +/// at WARN and fall through to the next precedence level. +pub fn resolve( + policy_tag: Option<&str>, + model_pref_tag: Option<&str>, + config: &Config, +) -> Result { + let kind = effective_kind(policy_tag, model_pref_tag)?; + target_for(kind, config) +} + +fn effective_kind( + policy_tag: Option<&str>, + model_pref_tag: Option<&str>, +) -> Result { + if let Some(tag) = model_pref_tag.filter(|t| !t.trim().is_empty()) { + match parse_tag(tag)? { + Some(kind) => return Ok(kind), + None => { + tracing::warn!( + tag, + "InferencePolicy modelPreference.primary.provider tag not recognised — \ + falling back to spec.provider / default" + ); + } + } + } + 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 no_tags_resolves_to_azure() { + let t = resolve(None, None, &cfg(None, None)).unwrap(); + assert_eq!(t, ProviderTarget::AzureOpenAI); + } + + #[test] + fn model_pref_tag_wins_over_policy_tag() { + let t = resolve( + Some("anthropic"), + Some("azure-openai"), + &cfg(Some("sk-x"), None), + ) + .unwrap(); + assert_eq!(t, ProviderTarget::AzureOpenAI); + } + + #[test] + fn unknown_model_pref_tag_falls_back_to_policy_tag() { + let t = resolve(Some("anthropic"), Some("gemini"), &cfg(Some("sk-x"), None)).unwrap(); + assert_eq!( + t, + 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"), None, &cfg(None, None)), + Err(ProviderError::MissingCredential { + provider: "anthropic", + .. + }) + )); + } + + #[test] + fn ollama_without_endpoint_fails_closed() { + assert!(matches!( + resolve(Some("ollama"), None, &cfg(None, None)), + Err(ProviderError::MissingEndpoint { + provider: "ollama", + .. + }) + )); + } + + #[test] + fn ollama_with_endpoint_resolves() { + let t = resolve( + Some("ollama"), + None, + &cfg(None, Some("http://ollama.ollama.svc:11434")), + ) + .unwrap(); + assert_eq!( + t, + ProviderTarget::Ollama { + endpoint: "http://ollama.ollama.svc:11434".into() + } + ); + } + + #[test] + fn bedrock_anywhere_is_unimplemented_not_silent() { + assert!(matches!( + resolve(Some("bedrock"), None, &cfg(None, None)), + Err(ProviderError::Unimplemented { .. }) + )); + assert!(matches!( + resolve(None, 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..42b4046e6 100644 --- a/inference-router/src/routes/anthropic_messages.rs +++ b/inference-router/src/routes/anthropic_messages.rs @@ -25,7 +25,10 @@ use futures::stream::StreamExt; use serde_json::{Value, json}; use super::AppState; +use crate::guardrails::{self, Direction, GuardrailPipeline}; +use crate::provider::{ProviderError, ProviderKind}; use crate::proxy; +use std::sync::Arc; fn deny_response(status: StatusCode, message: &str, code: &str) -> axum::response::Response { ( @@ -264,11 +267,98 @@ 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; + // Multi-provider slice: retarget at the policy-selected provider. + // Fails closed — see `routes::apply_provider_resolution`. + 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_response(status, &e.to_string(), "api_error"); + } + + // Guardrail pipeline — build fails closed on declared-but- + // unbuildable stages, input scan runs before any upstream forward. + 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_response(StatusCode::SERVICE_UNAVAILABLE, &e.to_string(), "api_error"); + } + }; + 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_response( + StatusCode::FORBIDDEN, + &v.message(), + "content_policy_violation", + ); + } + 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_response(StatusCode::BAD_GATEWAY, &e.to_string(), "api_error"); + } + } + } + + // Native Anthropic Messages pass-through — either the policy + // selected `provider: anthropic` (upstream is api.anthropic.com + // with the router-held API key) or the endpoint is GitHub Copilot + // (native /v1/messages). No translation: streaming, tool_use and + // multi-modal content flow through unchanged. + 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 +427,46 @@ 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_response( + StatusCode::FORBIDDEN, + &v.message(), + "content_policy_violation", + ); + } + 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_response(StatusCode::BAD_GATEWAY, &e.to_string(), "api_error"); + } + } + } + (StatusCode::OK, Json(anthropic_resp)).into_response() } Err(e) => { @@ -362,6 +492,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,7 +525,22 @@ 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))); + // Guardrail output scan (streaming, Anthropic event + // dialect): hold-and-release — see guardrails.rs. + let guarded = match guardrail_pipeline + .as_ref() + .filter(|p| p.covers(Direction::Output)) + { + Some(p) => guardrails::guard_sse_stream( + stream, + p.clone(), + guardrails::StreamDialect::AnthropicMessages, + sandbox_name.to_string(), + policy_digest.clone(), + ), + None => stream, + }; + let body = Body::from_stream(guarded.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() { @@ -454,6 +601,51 @@ 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 Ok(body_json) = serde_json::from_slice::(&resp_body) + { + let text = guardrails::extract_anthropic_output_text(&body_json); + 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_response( + StatusCode::FORBIDDEN, + &v.message(), + "content_policy_violation", + ); + } + 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_response( + StatusCode::BAD_GATEWAY, + &e.to_string(), + "api_error", + ); + } + } + } + let mut resp = axum::response::Response::builder().status(status); if let Some(h) = resp.headers_mut() { for (n, v) in resp_headers.iter() { diff --git a/inference-router/src/routes/chat_completions.rs b/inference-router/src/routes/chat_completions.rs index 2b53a7e20..5b0c12da2 100644 --- a/inference-router/src/routes/chat_completions.rs +++ b/inference-router/src/routes/chat_completions.rs @@ -20,8 +20,11 @@ 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, @@ -46,6 +49,125 @@ 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))) +} + +/// Run the output-direction guardrail stages over a buffered +/// OpenAI-shaped response body. `Ok(())` when there is nothing to do +/// or the scan passes; `Err(response)` carries the ready-made block +/// response. +pub(super) async fn enforce_openai_output_guardrails( + pipeline: Option<&Arc>, + resp_body: &[u8], + sandbox_name: &str, + policy_digest: &str, +) -> Result<(), axum::response::Response> { + let Some(p) = pipeline else { + return Ok(()); + }; + if !p.covers(Direction::Output) { + return Ok(()); + } + let Ok(body_json) = serde_json::from_slice::(resp_body) else { + return Ok(()); + }; + let text = guardrails::extract_openai_output_text(&body_json); + match p.scan(&text, Direction::Output).await { + Ok(None) => Ok(()), + 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" + ); + Err(guardrail_violation_response(&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" + ); + Err(guardrail_error_response(&e)) + } + } +} + /// POST /v1/chat/completions — the primary inference endpoint. pub(super) async fn chat_completions( State(state): State, @@ -197,6 +319,95 @@ pub(super) async fn chat_completions( // Slice 2d.1: honour `InferencePolicy.modelPreference.primary.deployment`. crate::routes::apply_model_preference_override(&mut upstream, &policy); + // Multi-provider slice: retarget the upstream when the policy + // selects a non-Azure provider. Fails closed on unimplemented / + // unconfigured providers — never a silent Azure fallback. + 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); + } + + // The Anthropic upstream speaks the Messages API, not + // chat-completions. Anthropic-native runtimes (Claude Agent SDK) + // already target the router's `/anthropic/v1/messages` surface, + // which forwards natively; an OpenAI-shaped client under an + // Anthropic-provider policy gets an explicit 501 (mirrors the + // GitHub-Models 501s for Foundry-only routes) instead of a + // confusing upstream 404. + if upstream.provider == ProviderKind::Anthropic { + return errors::openai( + StatusCode::NOT_IMPLEMENTED, + "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)", + "provider_unimplemented", + ) + .into_response(); + } + + // Guardrail pipeline (multi-cloud guardrails slice). Built per + // request from the policy snapshot; a declared stage that cannot + // be materialised blocks the request. + 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 = serde_json::from_slice::(&body) + .map(|v| guardrails::extract_openai_input_text(&v)) + .unwrap_or_default(); + 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 +465,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 +506,27 @@ pub(super) async fn chat_completions( { budget.record_usage(&sandbox_owned, total).await; } + // Guardrail output scan — the Responses-API + // recovery path must not bypass a declared + // pipeline. Violations surface as an SSE + // error frame (this branch already committed + // to text/event-stream). + if let Err(resp) = enforce_openai_output_guardrails( + guardrail_for_task.as_ref(), + &chat_body, + &sandbox_owned, + &digest_for_task, + ) + .await + { + let _ = resp; // structured 4xx/5xx body not sendable mid-SSE + let err_sse = format!( + "data: {}\n\ndata: [DONE]\n\n", + serde_json::json!({"error":{"message":"Blocked by guardrail pipeline","type":"content_policy_violation","code":"guardrail_blocked"}}) + ); + let _ = tx.send(Ok(bytes::Bytes::from(err_sse))).await; + return; + } let sse_data = format!( "data: {}\n\ndata: [DONE]\n\n", String::from_utf8_lossy(&chat_body) @@ -347,6 +581,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 +692,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 +747,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 +836,25 @@ pub(super) async fn chat_completions( } chunk }); - let body = Body::from_stream(wrapped); + // Guardrail output scan (streaming): hold-and-release + // windows — no model text reaches the client before a + // scan has covered it. Skipped entirely (zero cost) + // when the policy declares 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()); @@ -794,6 +1067,22 @@ pub(super) async fn chat_completions( return resp; } + // Guardrail output scan (buffered) — runs + // beside the contentSafety floor: the floor + // polices upstream-native annotations, the + // pipeline runs the policy's external + // backends (e.g. OpenAI Moderation). + if let Err(block) = enforce_openai_output_guardrails( + guardrail_pipeline.as_ref(), + &resp_body, + sandbox_name, + &policy.digest, + ) + .await + { + return block; + } + // AGT output pipeline: redact → scan → policy check (blocking) let response_text = body_json .get("choices") diff --git a/inference-router/src/routes/mod.rs b/inference-router/src/routes/mod.rs index 38428bd5f..c8148c1ba 100644 --- a/inference-router/src/routes/mod.rs +++ b/inference-router/src/routes/mod.rs @@ -366,12 +366,57 @@ 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()) + } +} + +/// Multi-provider slice — retarget `upstream` at the provider the +/// loaded `InferencePolicy` selects. No-op (Azure) when the policy +/// carries no provider opinion, keeping the historic behaviour for +/// every sandbox without a policy. +/// +/// Fails closed: a policy that names a provider the router cannot +/// serve (unimplemented `bedrock`, or missing endpoint/credential +/// config) yields an error the handler must surface — silently +/// falling back to Azure would ship prompts to a provider the +/// operator didn't select. +pub(crate) fn apply_provider_resolution( + state: &AppState, + upstream: &mut UpstreamConfig, + policy: &crate::inference_policy_loader::InferencePolicySnapshot, +) -> Result<(), crate::provider::ProviderError> { + let model_pref_tag = policy + .model_preference + .as_ref() + .map(|m| m.primary.provider.as_str()); + let target = + crate::provider::resolve(policy.provider.as_deref(), model_pref_tag, &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/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/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(); From e7fb2574c7436f325155e8ee91871bf93dd7c1b9 Mon Sep 17 00:00:00 2001 From: johnseong Date: Fri, 31 Jul 2026 17:35:16 -0400 Subject: [PATCH 02/18] fix(guardrails): accurate SSE block frames + bounded stream scan context MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address Copilot review on #488: - Responses-API recovery SSE branch now emits the outcome-accurate frame (content_policy_violation vs guardrail_unavailable/ guardrail_misconfigured) instead of a hard-coded violation frame — scan_openai_output_guardrails returns the block as data and each transport picks its wire shape. - SseGuardState trims retained scan context to MAX_SCAN_CHARS after every clean scan; output scans only ever submit the trailing MAX_SCAN_CHARS anyway, so per-connection memory stays bounded on long streams without weakening scanned-before-delivery. Regression test added. --- inference-router/src/guardrails.rs | 58 ++++++++++++- .../src/routes/chat_completions.rs | 81 +++++++++++++------ 2 files changed, 110 insertions(+), 29 deletions(-) diff --git a/inference-router/src/guardrails.rs b/inference-router/src/guardrails.rs index f415c4451..986377645 100644 --- a/inference-router/src/guardrails.rs +++ b/inference-router/src/guardrails.rs @@ -622,8 +622,10 @@ fn delta_text_from_event(dialect: StreamDialect, event: &serde_json::Value) -> O /// The client-facing SSE error frame emitted when a stream is cut by /// a guardrail. OpenAI-style error object works for both dialects' /// SDK error paths and is what the existing content-safety stream cut -/// emits too. -fn violation_sse_frame(violation: &GuardrailViolation) -> Bytes { +/// emits too. Public so buffered-to-SSE conversion paths (e.g. the +/// Responses-API recovery branch) can emit the same frame shape. +#[must_use] +pub fn violation_sse_frame(violation: &GuardrailViolation) -> Bytes { Bytes::from(format!( "data: {}\n\ndata: [DONE]\n\n", serde_json::json!({ @@ -636,7 +638,11 @@ fn violation_sse_frame(violation: &GuardrailViolation) -> Bytes { )) } -fn unavailable_sse_frame(err: &GuardrailError) -> Bytes { +/// 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!({ @@ -757,11 +763,31 @@ impl SseGuardState { { 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(unavailable_sse_frame(&e)), + Err(e) => SseGuardStep::Cut(error_sse_frame(&e)), + } + } + + /// Bound the retained scan context after a clean scan. Output + /// scans only ever submit the trailing [`MAX_SCAN_CHARS`] chars + /// (see [`cap_for_scan`]), so anything older cannot influence a + /// future scan's input — dropping it keeps per-connection memory + /// bounded on long-lived streams without weakening the + /// scanned-before-delivery contract. + fn trim_scan_context(&mut self) { + if self.accumulated.chars().count() <= MAX_SCAN_CHARS { + return; } + let start = self + .accumulated + .char_indices() + .rev() + .nth(MAX_SCAN_CHARS - 1) + .map_or(0, |(i, _)| i); + self.accumulated = self.accumulated.split_off(start); } } @@ -1281,6 +1307,30 @@ mod tests { ); } + #[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)); diff --git a/inference-router/src/routes/chat_completions.rs b/inference-router/src/routes/chat_completions.rs index 5b0c12da2..c11372772 100644 --- a/inference-router/src/routes/chat_completions.rs +++ b/inference-router/src/routes/chat_completions.rs @@ -119,28 +119,46 @@ pub(super) fn build_guardrail_pipeline( .map(|p| Some(Arc::new(p))) } +/// A blocked output scan: either a confirmed content violation or a +/// fail-closed guardrail failure. Kept as data (not a ready-made +/// `Response`) so each transport picks the accurate wire shape — +/// buffered paths map to 403/502/503 HTTP responses, SSE paths to +/// the matching `guardrails::{violation,error}_sse_frame`. +pub(super) enum OutputGuardrailBlock { + Violation(GuardrailViolation), + Error(GuardrailError), +} + +impl OutputGuardrailBlock { + /// SSE frame carrying this block's own type/code — a backend + /// outage must surface as `guardrail_unavailable`, never be + /// mislabelled `content_policy_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. `Ok(())` when there is nothing to do -/// or the scan passes; `Err(response)` carries the ready-made block -/// response. -pub(super) async fn enforce_openai_output_guardrails( +/// 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, -) -> Result<(), axum::response::Response> { - let Some(p) = pipeline else { - return Ok(()); - }; +) -> Option { + let p = pipeline?; if !p.covers(Direction::Output) { - return Ok(()); + return None; } - let Ok(body_json) = serde_json::from_slice::(resp_body) else { - return Ok(()); - }; + let body_json = serde_json::from_slice::(resp_body).ok()?; let text = guardrails::extract_openai_output_text(&body_json); match p.scan(&text, Direction::Output).await { - Ok(None) => Ok(()), + Ok(None) => None, Ok(Some(v)) => { tracing::warn!( target: "inference.audit", @@ -151,7 +169,7 @@ pub(super) async fn enforce_openai_output_guardrails( categories = ?v.categories, "guardrail pipeline blocked buffered response" ); - Err(guardrail_violation_response(&v)) + Some(OutputGuardrailBlock::Violation(v)) } Err(e) => { tracing::warn!( @@ -163,11 +181,27 @@ pub(super) async fn enforce_openai_output_guardrails( error = %e, "guardrail pipeline unavailable — failing closed" ); - Err(guardrail_error_response(&e)) + 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). +pub(super) async fn enforce_openai_output_guardrails( + pipeline: Option<&Arc>, + resp_body: &[u8], + sandbox_name: &str, + policy_digest: &str, +) -> Result<(), axum::response::Response> { + match scan_openai_output_guardrails(pipeline, resp_body, sandbox_name, policy_digest).await { + None => Ok(()), + Some(OutputGuardrailBlock::Violation(v)) => Err(guardrail_violation_response(&v)), + Some(OutputGuardrailBlock::Error(e)) => Err(guardrail_error_response(&e)), + } +} + /// POST /v1/chat/completions — the primary inference endpoint. pub(super) async fn chat_completions( State(state): State, @@ -508,10 +542,12 @@ pub(super) async fn chat_completions( } // Guardrail output scan — the Responses-API // recovery path must not bypass a declared - // pipeline. Violations surface as an SSE - // error frame (this branch already committed - // to text/event-stream). - if let Err(resp) = enforce_openai_output_guardrails( + // pipeline. This branch already committed to + // text/event-stream, so the block surfaces as + // the outcome-accurate SSE frame (violation + // vs guardrail_unavailable/misconfigured — + // never a mislabelled content violation). + if let Some(block) = scan_openai_output_guardrails( guardrail_for_task.as_ref(), &chat_body, &sandbox_owned, @@ -519,12 +555,7 @@ pub(super) async fn chat_completions( ) .await { - let _ = resp; // structured 4xx/5xx body not sendable mid-SSE - let err_sse = format!( - "data: {}\n\ndata: [DONE]\n\n", - serde_json::json!({"error":{"message":"Blocked by guardrail pipeline","type":"content_policy_violation","code":"guardrail_blocked"}}) - ); - let _ = tx.send(Ok(bytes::Bytes::from(err_sse))).await; + let _ = tx.send(Ok(block.sse_frame())).await; return; } let sse_data = format!( From 17a2cf181df62fffbe206d08bf85f872dc650e68 Mon Sep 17 00:00:00 2001 From: johnseong Date: Fri, 31 Jul 2026 17:45:12 -0400 Subject: [PATCH 03/18] fix(guardrails): tolerant SSE data: prefix + raw-text scan fallback MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second Copilot review round on #488: - SSE guard accepts 'data:' with or without whitespace, so spaceless events can't slip through the stream scan unrecognised. - New guardrails::scan_text_or_raw: when a declared input/output guardrail is active and the body fails to parse as JSON, scan the raw (lossy-UTF-8) bytes instead of skipping — applied to the chat-completions input scan, buffered output scan, and the Anthropic pass-through buffered output scan. --- inference-router/src/guardrails.rs | 53 ++++++++++++++++++- .../src/routes/anthropic_messages.rs | 6 ++- .../src/routes/chat_completions.rs | 7 +-- 3 files changed, 58 insertions(+), 8 deletions(-) diff --git a/inference-router/src/guardrails.rs b/inference-router/src/guardrails.rs index 986377645..41a978eda 100644 --- a/inference-router/src/guardrails.rs +++ b/inference-router/src/guardrails.rs @@ -578,6 +578,21 @@ pub fn extract_anthropic_output_text(body: &serde_json::Value) -> String { out.join("\n") } +/// Extract scan text from a body via `extract`, falling back to the +/// raw (lossy-UTF-8) bytes when the body is not JSON. A declared +/// guardrail must never be skipped because a body failed to parse — +/// the raw fallback keeps unknown/malformed shapes covered instead of +/// letting them through unscanned. (A body that parses but yields no +/// extractable text — e.g. tool-call-only responses — is a deliberate +/// pass: the extractors define the scannable surface.) +#[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(), + } +} + // ─── Streaming (SSE) guard ─────────────────────────────────────────────────── /// SSE wire dialect of the guarded stream — decides how delta text is @@ -717,9 +732,11 @@ impl SseGuardState { }; self.line_carry = rest; for line in complete.lines() { - let Some(payload) = line.trim().strip_prefix("data: ") else { + // SSE permits `data:` with or without a following space. + let Some(payload) = line.trim().strip_prefix("data:") else { continue; }; + let payload = payload.trim_start(); if payload == "[DONE]" { continue; } @@ -986,6 +1003,17 @@ mod tests { // ---- 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!({ @@ -1307,6 +1335,29 @@ mod tests { ); } + #[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 diff --git a/inference-router/src/routes/anthropic_messages.rs b/inference-router/src/routes/anthropic_messages.rs index 42b4046e6..8919967bc 100644 --- a/inference-router/src/routes/anthropic_messages.rs +++ b/inference-router/src/routes/anthropic_messages.rs @@ -606,9 +606,11 @@ async fn forward_anthropic_passthrough( && let Some(p) = guardrail_pipeline .as_ref() .filter(|p| p.covers(Direction::Output)) - && let Ok(body_json) = serde_json::from_slice::(&resp_body) { - let text = guardrails::extract_anthropic_output_text(&body_json); + 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)) => { diff --git a/inference-router/src/routes/chat_completions.rs b/inference-router/src/routes/chat_completions.rs index c11372772..d2fc8c416 100644 --- a/inference-router/src/routes/chat_completions.rs +++ b/inference-router/src/routes/chat_completions.rs @@ -155,8 +155,7 @@ pub(super) async fn scan_openai_output_guardrails( if !p.covers(Direction::Output) { return None; } - let body_json = serde_json::from_slice::(resp_body).ok()?; - let text = guardrails::extract_openai_output_text(&body_json); + 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)) => { @@ -410,9 +409,7 @@ pub(super) async fn chat_completions( if let Some(ref p) = guardrail_pipeline && p.covers(Direction::Input) { - let input_text = serde_json::from_slice::(&body) - .map(|v| guardrails::extract_openai_input_text(&v)) - .unwrap_or_default(); + 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)) => { From 5d256dbe9c50e5a620fb4ef3ea1e50fb216aa2d9 Mon Sep 17 00:00:00 2001 From: johnseong Date: Fri, 31 Jul 2026 18:43:59 -0400 Subject: [PATCH 04/18] fix(anthropic): strip hop-by-hop headers when relaying pass-through responses MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Found in live testing against api.anthropic.com: when the upstream connection negotiates HTTP/1.1, Anthropic responds with transfer-encoding: chunked. The pass-through handler copied all upstream headers onto the rebuilt axum response, and hyper refuses to serialize a response carrying a stale framing header — the client got an empty reply despite a 200 from upstream. The pre-existing Copilot pass-through never hit this (h2 end-to-end), and wiremock tests don't (simple headers). Both relay loops (buffered + streaming) now skip the RFC 9110 connection-specific headers; hyper re-frames the body itself. Verified live: non-streaming + SSE Messages against api.anthropic.com, chat completions (buffered + SSE) against local Ollama, policy hot-reload provider swap, guardrail fail-closed without a key (503), and input-block + mid-stream output cut against a local moderation stub. --- .../src/routes/anthropic_messages.rs | 25 +++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/inference-router/src/routes/anthropic_messages.rs b/inference-router/src/routes/anthropic_messages.rs index 8919967bc..b90bffb59 100644 --- a/inference-router/src/routes/anthropic_messages.rs +++ b/inference-router/src/routes/anthropic_messages.rs @@ -30,6 +30,25 @@ 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 { ( status, @@ -544,6 +563,9 @@ async fn forward_anthropic_passthrough( 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( @@ -651,6 +673,9 @@ async fn forward_anthropic_passthrough( 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) { From a447a9a446bd03caa99a9ffefeb7a8dbb93a6b3f Mon Sep 17 00:00:00 2001 From: johnseong Date: Fri, 31 Jul 2026 21:13:46 -0400 Subject: [PATCH 05/18] fix(guardrails,provider): close review blockers B1-B4 + windowed scanning MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Independent review of #488 found four blockers; all fixed with tests and verified live against real Anthropic/Ollama + a moderation stub. B1 — SSE hold-and-release leaked unscanned model text on a chunk boundary that fell inside a data: line: the partial line's bytes were already in `held` but its delta text was uncounted, and the unscanned==0 fast-release shipped them. Now never release while a partial line is buffered; on_end flushes a trailing unterminated line before the final scan. Regression test splits an event mid-content; live test with a 40-char threshold delivered only the error frame. B2 — buffered Responses-API recovery path (400 'unsupported' → /responses) returned the completion with no output scan. Added the enforce_openai_output_guardrails call, matching the other two recovery paths. B3 — routing was retroactively driven by the pre-existing modelPreference.primary.provider tag, so an unchanged CR could 503 or silently cross clouds on upgrade. spec.provider is now the sole routing selector; modelPreference.provider stays informational (drives deployment failover only). Verified: modelPreference provider=anthropic with no spec.provider stays on the Azure upstream. B4 — the sibling inference routes (/v1/completions, /v1/responses, /v1/embeddings, image generation) didn't consult provider/guardrail policy, so an agent could bypass both by not using chat/completions. They now fail closed: 501 on a non-Azure spec.provider, 403 when guardrails are declared. Pure classifier unit-tested; live-verified 501 on /v1/embeddings and /v1/responses under an ollama policy. M1 — 16k scan cap truncated instead of windowing, letting content hide past the cap; GUARDRAIL_STREAM_SCAN_CHARS was unclamped. Now scan_windows() covers all text in successive MAX_SCAN_CHARS windows and the stream threshold is clamped to the cap. Docs/CHANGELOG corrected: routing precedence, sibling-route scope, windowed scanning, and the tool-arg / thinking-delta output-scan gaps (roadmap). 1958 tests green. --- CHANGELOG.md | 43 +-- docs/api/crd-reference.md | 8 +- inference-router/src/guardrails.rs | 254 ++++++++++-------- inference-router/src/provider.rs | 130 ++++----- .../src/routes/chat_completions.rs | 141 ++++++++++ inference-router/src/routes/inference.rs | 25 ++ inference-router/src/routes/mod.rs | 7 +- 7 files changed, 398 insertions(+), 210 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a8c87f12f..0bddf172e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,10 +37,14 @@ never sees a provider key. **Inference router — multi-provider upstreams** -- New `provider` module: tag parsing + fail-closed resolution - (`modelPreference.primary.provider` wins over `spec.provider`; - unknown tags warn and fall through; missing endpoint/credential → - 503, unimplemented `bedrock` → 501 — never a silent Azure fallback). +- 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: @@ -48,28 +52,37 @@ never sees a provider key. 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; usage metered from Anthropic - `usage` fields. OpenAI-shaped `/v1/chat/completions` under an - Anthropic policy returns an explicit 501 pointing at the Messages - surface (mirrors the GitHub-Models 501 precedent). + 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. The sibling inference routes + (`/v1/completions`, `/v1/responses`, `/v1/embeddings`, image + generation) don't implement either and **fail 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. **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 at every governed exchange: request pre-flight (input - stages), buffered responses (including the Responses-API recovery +- 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 chars) — no model text - is delivered before a scan has covered it; flagged streams are cut - with a structured SSE error frame + `data: [DONE]`. + (`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`); scan-text truncation (16k-char cap) is - logged, never silent. New `kars_guardrail_scans_total` metric + `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. diff --git a/docs/api/crd-reference.md b/docs/api/crd-reference.md index 51f22315a..a2b77ca1d 100644 --- a/docs/api/crd-reference.md +++ b/docs/api/crd-reference.md @@ -484,9 +484,9 @@ spec: | Field | Notes | |---|---| | `spec.appliesTo` | Required selector — AND of `sandboxName`, `sandboxMatchLabels`, `action`. | -| `spec.provider` | Optional typed default provider (`azure-openai` \| `anthropic` \| `ollama` \| `bedrock`). 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. Credentials/endpoints are router-sidecar config only — never visible to the agent. | -| `spec.guardrails[]` | Optional ordered guardrail pipeline (1–8 stages) the router runs around each governed call. 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 model text reaches the client before a scan has covered it. | -| `spec.modelPreference.primary` | `{provider, deployment}`. `provider` is one of `azure-openai`, `anthropic`, `gemini`, `bedrock`, `ollama`. A recognised tag here overrides `spec.provider` for routing. | +| `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`. | @@ -494,6 +494,8 @@ spec: | `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` / `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`). The sibling inference routes — `/v1/completions`, `/v1/responses`, `/v1/embeddings`, and image generation — do **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. 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/inference-router/src/guardrails.rs b/inference-router/src/guardrails.rs index 41a978eda..6e654cbc5 100644 --- a/inference-router/src/guardrails.rs +++ b/inference-router/src/guardrails.rs @@ -37,9 +37,9 @@ //! (one moderation round-trip per window), which is the standard //! trade-off for streaming guardrails. //! -//! Scanned text is capped at [`MAX_SCAN_CHARS`] (most recent chars -//! for output, leading chars for input) to stay under moderation -//! input limits; truncation is logged at WARN — never silent. +//! Text longer than [`MAX_SCAN_CHARS`] is scanned in successive +//! windows (see [`scan_windows`]), never truncated, so content can't +//! be hidden past the per-call cap. use std::sync::Arc; @@ -387,8 +387,10 @@ impl GuardrailPipeline { self.stages.iter().any(|s| s.apply_to.covers(direction)) } - /// Run every stage covering `direction` over `text`, in order. - /// First flagged verdict wins. Empty text short-circuits to pass. + /// 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, @@ -397,30 +399,31 @@ impl GuardrailPipeline { if text.is_empty() { return Ok(None); } - let capped = cap_for_scan(text, direction); + let windows = scan_windows(text); for stage in self.stages.iter().filter(|s| s.apply_to.covers(direction)) { - let outcome = stage.guard.scan(capped).await; - match outcome { - 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); + 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); + } } } } @@ -428,38 +431,25 @@ impl GuardrailPipeline { } } -/// Cap text to [`MAX_SCAN_CHARS`]: leading chars for input (the -/// system prompt + earliest instructions), trailing chars for output -/// (the newest generated text — earlier output was already scanned by -/// previous windows in the streaming path). Logs at WARN on -/// truncation. -fn cap_for_scan(text: &str, direction: Direction) -> &str { - if text.chars().count() <= MAX_SCAN_CHARS { - return text; - } - tracing::warn!( - direction = direction.as_str(), - total_chars = text.chars().count(), - scanned_chars = MAX_SCAN_CHARS, - "guardrail scan text exceeds cap — scanning a truncated window" - ); - match direction { - Direction::Input => { - let end = text - .char_indices() - .nth(MAX_SCAN_CHARS) - .map_or(text.len(), |(i, _)| i); - &text[..end] - } - Direction::Output => { - let start = text - .char_indices() - .rev() - .nth(MAX_SCAN_CHARS - 1) - .map_or(0, |(i, _)| i); - &text[start..] +/// Split `text` into consecutive windows of at most [`MAX_SCAN_CHARS`] +/// chars (never mid-char) so scanning all windows covers everything. +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 ────────────────────────────────────── @@ -670,13 +660,16 @@ pub fn error_sse_frame(err: &GuardrailError) -> Bytes { )) } -/// Effective hold-and-release window size. +/// 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 @@ -718,11 +711,8 @@ impl SseGuardState { } } - /// Pull complete lines out of `chunk` (+ carry), extract delta - /// text, and account it as unscanned. fn ingest_text(&mut self, chunk: &[u8]) { self.line_carry.push_str(&String::from_utf8_lossy(chunk)); - // Keep the trailing partial line (no '\n' yet) in the carry. let (complete, rest) = match self.line_carry.rfind('\n') { Some(idx) => { let (c, r) = self.line_carry.split_at(idx + 1); @@ -732,31 +722,37 @@ impl SseGuardState { }; self.line_carry = rest; for line in complete.lines() { - // SSE permits `data:` with or without a following space. - let Some(payload) = line.trim().strip_prefix("data:") else { - continue; - }; - let payload = payload.trim_start(); - if payload == "[DONE]" { - continue; - } - if let Ok(event) = serde_json::from_str::(payload) - && let Some(text) = delta_text_from_event(self.dialect, &event) - { - self.unscanned += text.chars().count(); - self.accumulated.push_str(&text); - } + 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 == "[DONE]" { + return; + } + if let Ok(event) = serde_json::from_str::(payload) + && let Some(text) = delta_text_from_event(self.dialect, &event) + { + self.unscanned += text.chars().count(); + self.accumulated.push_str(&text); } } async fn on_chunk(&mut self, chunk: Bytes) -> SseGuardStep { self.ingest_text(&chunk); self.held.push(chunk); + // A buffered partial line's bytes are already in `held` but + // its delta text has not been counted or scanned — never + // release while one is pending, or that text ships unscanned. + if !self.line_carry.is_empty() { + return SseGuardStep::Release(Vec::new()); + } if self.unscanned < self.threshold { - // Fast path: window not full. Chunks carrying no delta - // text at all (keepalives, role/annotation frames) are - // safe to release immediately when nothing text-bearing - // is being held alongside them. if self.unscanned == 0 { return SseGuardStep::Release(std::mem::take(&mut self.held)); } @@ -766,6 +762,12 @@ impl SseGuardState { } async fn on_end(&mut self) -> SseGuardStep { + // Flush a trailing unterminated line (truncated upstream) so + // its text is scanned before the held bytes are released. + 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)); } @@ -788,21 +790,20 @@ impl SseGuardState { } } - /// Bound the retained scan context after a clean scan. Output - /// scans only ever submit the trailing [`MAX_SCAN_CHARS`] chars - /// (see [`cap_for_scan`]), so anything older cannot influence a - /// future scan's input — dropping it keeps per-connection memory - /// bounded on long-lived streams without weakening the - /// scanned-before-delivery contract. + /// Retain `MAX_SCAN_CHARS - threshold` chars of scanned context + /// after a clean scan: bounds per-connection memory and keeps the + /// next scan (context + up to `threshold` new chars) within one + /// window, with overlap for cross-boundary detection. fn trim_scan_context(&mut self) { - if self.accumulated.chars().count() <= MAX_SCAN_CHARS { + 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(MAX_SCAN_CHARS - 1) + .nth(keep - 1) .map_or(0, |(i, _)| i); self.accumulated = self.accumulated.split_off(start); } @@ -1080,25 +1081,6 @@ mod tests { ); } - // ---- cap ---- - - #[test] - fn cap_keeps_short_text_intact() { - assert_eq!(cap_for_scan("short", Direction::Input), "short"); - } - - #[test] - fn cap_truncates_head_for_input_and_tail_for_output() { - let long: String = "a".repeat(MAX_SCAN_CHARS) + "TAIL"; - let capped_in = cap_for_scan(&long, Direction::Input); - assert_eq!(capped_in.len(), MAX_SCAN_CHARS); - assert!(capped_in.starts_with('a') && !capped_in.contains("TAIL")); - let long2 = "HEAD".to_string() + &"b".repeat(MAX_SCAN_CHARS); - let capped_out = cap_for_scan(&long2, Direction::Output); - assert_eq!(capped_out.len(), MAX_SCAN_CHARS); - assert!(!capped_out.contains("HEAD")); - } - // ---- pipeline + streaming with a fake backend ---- /// Test backend: flags any text containing the marker. Counts @@ -1404,4 +1386,62 @@ mod tests { ); assert!(out.contains("guardrail_blocked")); } + + #[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"); + } } diff --git a/inference-router/src/provider.rs b/inference-router/src/provider.rs index 6bf7871db..fa5f13d39 100644 --- a/inference-router/src/provider.rs +++ b/inference-router/src/provider.rs @@ -3,35 +3,26 @@ //! Multi-provider upstream resolution. //! -//! Maps the `InferencePolicy` provider tags (`azure-openai` / -//! `anthropic` / `ollama` / `bedrock`) onto a concrete upstream -//! target: base URL shape + auth scheme. The provider tag travels in -//! the compiled policy JSON (`spec.provider`, and -//! `spec.modelPreference.*.provider`); the *credentials and endpoints* -//! come exclusively from the router's own environment / secret mounts -//! — the agent process never sees a provider API key, exactly as with -//! the Azure Workload Identity path. +//! Maps `InferencePolicy.spec.provider` (`azure-openai` / `anthropic` +//! / `ollama` / `bedrock`) onto a concrete upstream target: base URL +//! shape + auth scheme. Credentials and endpoints come exclusively +//! from the router's own environment / secret mounts — the agent +//! process never sees a provider API key, exactly as with the Azure +//! Workload Identity path. //! -//! ## Resolution precedence -//! -//! 1. `modelPreference.primary.provider`, when it parses to a known -//! tag — an explicit route preference wins over the policy-level -//! default. -//! 2. `spec.provider` (policy-level default). -//! 3. `azure-openai` (absent / unknown tags — matches the pre-slice -//! behaviour where provider tags were informational-only). +//! `spec.provider` is the only routing selector. The pre-existing +//! `modelPreference.primary.provider` tag stays informational (it +//! drove no routing before this slice), so adding provider routing +//! doesn't retroactively reroute CRs that only set a model +//! preference. Absent/empty `spec.provider` ⇒ Azure; an unrecognised +//! tag warns and falls back to Azure. //! //! `bedrock` is recognised but not yet implemented: a policy that //! declares it gets an explicit 501-style error instead of a silent //! reroute to Azure — declared intent must never be silently ignored. -//! -//! ## Failure semantics -//! -//! A resolvable provider with missing router-side configuration +//! Likewise a resolvable provider with missing router-side config //! (no `ANTHROPIC_API_KEY`, no `OLLAMA_ENDPOINT`) fails the request -//! closed with a specific, operator-actionable error. Falling back to -//! the Azure upstream would send prompts to a provider the operator -//! didn't select. +//! closed rather than falling back to Azure. use crate::config::Config; @@ -130,37 +121,20 @@ impl ProviderTarget { } } -/// Resolve the effective provider for a request. +/// Resolve the effective provider for a request from `spec.provider`. /// -/// `model_pref_tag` is `modelPreference.primary.provider` (when a -/// policy carries a model preference), `policy_tag` is the top-level -/// `spec.provider`. See module docs for precedence. Unknown tags log -/// at WARN and fall through to the next precedence level. -pub fn resolve( - policy_tag: Option<&str>, - model_pref_tag: Option<&str>, - config: &Config, -) -> Result { - let kind = effective_kind(policy_tag, model_pref_tag)?; +/// `spec.provider` is the sole routing selector. The pre-existing +/// `modelPreference.primary.provider` tag stays informational (it +/// drove nothing before this slice) so an unchanged CR that only set +/// a model preference keeps its Azure upstream — routing is opt-in +/// via the new field. Absent/empty ⇒ Azure; an unrecognised tag +/// warns and falls back to Azure; `bedrock` is a hard error. +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>, - model_pref_tag: Option<&str>, -) -> Result { - if let Some(tag) = model_pref_tag.filter(|t| !t.trim().is_empty()) { - match parse_tag(tag)? { - Some(kind) => return Ok(kind), - None => { - tracing::warn!( - tag, - "InferencePolicy modelPreference.primary.provider tag not recognised — \ - falling back to spec.provider / default" - ); - } - } - } +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), @@ -264,27 +238,31 @@ mod tests { } #[test] - fn no_tags_resolves_to_azure() { - let t = resolve(None, None, &cfg(None, None)).unwrap(); - assert_eq!(t, ProviderTarget::AzureOpenAI); + 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 model_pref_tag_wins_over_policy_tag() { - let t = resolve( - Some("anthropic"), - Some("azure-openai"), - &cfg(Some("sk-x"), None), - ) - .unwrap(); - assert_eq!(t, ProviderTarget::AzureOpenAI); + 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 unknown_model_pref_tag_falls_back_to_policy_tag() { - let t = resolve(Some("anthropic"), Some("gemini"), &cfg(Some("sk-x"), None)).unwrap(); + fn anthropic_provider_resolves_with_key() { assert_eq!( - t, + resolve(Some("anthropic"), &cfg(Some("sk-x"), None)).unwrap(), ProviderTarget::Anthropic { endpoint: "https://api.anthropic.com".into(), api_key: "sk-x".into(), @@ -295,7 +273,7 @@ mod tests { #[test] fn anthropic_without_key_fails_closed() { assert!(matches!( - resolve(Some("anthropic"), None, &cfg(None, None)), + resolve(Some("anthropic"), &cfg(None, None)), Err(ProviderError::MissingCredential { provider: "anthropic", .. @@ -306,7 +284,7 @@ mod tests { #[test] fn ollama_without_endpoint_fails_closed() { assert!(matches!( - resolve(Some("ollama"), None, &cfg(None, None)), + resolve(Some("ollama"), &cfg(None, None)), Err(ProviderError::MissingEndpoint { provider: "ollama", .. @@ -316,14 +294,12 @@ mod tests { #[test] fn ollama_with_endpoint_resolves() { - let t = resolve( - Some("ollama"), - None, - &cfg(None, Some("http://ollama.ollama.svc:11434")), - ) - .unwrap(); assert_eq!( - t, + resolve( + Some("ollama"), + &cfg(None, Some("http://ollama.ollama.svc:11434")) + ) + .unwrap(), ProviderTarget::Ollama { endpoint: "http://ollama.ollama.svc:11434".into() } @@ -331,13 +307,9 @@ mod tests { } #[test] - fn bedrock_anywhere_is_unimplemented_not_silent() { - assert!(matches!( - resolve(Some("bedrock"), None, &cfg(None, None)), - Err(ProviderError::Unimplemented { .. }) - )); + fn bedrock_is_unimplemented_not_silent() { assert!(matches!( - resolve(None, Some("bedrock"), &cfg(None, None)), + resolve(Some("bedrock"), &cfg(None, None)), Err(ProviderError::Unimplemented { .. }) )); } diff --git a/inference-router/src/routes/chat_completions.rs b/inference-router/src/routes/chat_completions.rs index d2fc8c416..d5d7c477f 100644 --- a/inference-router/src/routes/chat_completions.rs +++ b/inference-router/src/routes/chat_completions.rs @@ -201,6 +201,79 @@ pub(super) async fn enforce_openai_output_guardrails( } } +/// 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}" + ); + let mut resp = errors::openai(status, &msg, 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, @@ -966,6 +1039,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()); @@ -1416,6 +1499,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..34b8f1940 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}"); diff --git a/inference-router/src/routes/mod.rs b/inference-router/src/routes/mod.rs index c8148c1ba..078ce8471 100644 --- a/inference-router/src/routes/mod.rs +++ b/inference-router/src/routes/mod.rs @@ -385,12 +385,7 @@ pub(crate) fn apply_provider_resolution( upstream: &mut UpstreamConfig, policy: &crate::inference_policy_loader::InferencePolicySnapshot, ) -> Result<(), crate::provider::ProviderError> { - let model_pref_tag = policy - .model_preference - .as_ref() - .map(|m| m.primary.provider.as_str()); - let target = - crate::provider::resolve(policy.provider.as_deref(), model_pref_tag, &state.config)?; + let target = crate::provider::resolve(policy.provider.as_deref(), &state.config)?; match target { crate::provider::ProviderTarget::AzureOpenAI => {} crate::provider::ProviderTarget::Anthropic { endpoint, api_key } => { From 0459ae0c977a542171769e77c0d99b445de38f81 Mon Sep 17 00:00:00 2001 From: johnseong Date: Fri, 31 Jul 2026 21:21:14 -0400 Subject: [PATCH 06/18] fix(guardrails): scan non-JSON SSE data frames instead of fast-releasing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review comment on #488: SseGuardState::ingest_line ignored data: payloads that don't parse as JSON, so unscanned stayed 0 and the fast-release path shipped those bytes without a scan — an upstream drift / malformed frame could bypass the scanned-before-delivery contract. Non-JSON data payloads are now added to the scan buffer as raw text. Valid-JSON structural frames (ping / role-only / stop) still release without a scan round-trip; the non-content extraction surface (tool-call args, thinking deltas) remains the documented follow-up. Two regression tests added. --- inference-router/src/guardrails.rs | 94 ++++++++++++++++++++++++++++-- 1 file changed, 88 insertions(+), 6 deletions(-) diff --git a/inference-router/src/guardrails.rs b/inference-router/src/guardrails.rs index 6e654cbc5..bbf959266 100644 --- a/inference-router/src/guardrails.rs +++ b/inference-router/src/guardrails.rs @@ -732,14 +732,29 @@ impl SseGuardState { return; }; let payload = payload.trim_start(); - if payload == "[DONE]" { + if payload.is_empty() || payload == "[DONE]" { return; } - if let Ok(event) = serde_json::from_str::(payload) - && let Some(text) = delta_text_from_event(self.dialect, &event) - { - self.unscanned += text.chars().count(); - self.accumulated.push_str(&text); + match serde_json::from_str::(payload) { + // Known text-bearing delta: scan the extracted text. + 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 / message_start / role-only / stop); no model + // text to scan. Non-`content` fields (tool-call args, + // "thinking" deltas) are the documented extraction- + // surface gap, not handled here. + } + // Non-JSON `data:` payload (upstream drift / malformed + // frame): scan the raw payload rather than fast-releasing + // it — an unrecognised frame must not bypass the scan. + Err(_) => { + self.unscanned += payload.chars().count(); + self.accumulated.push_str(payload); + } } } @@ -1387,6 +1402,73 @@ mod tests { 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 From ca1933419d8cb1db210dadd8b1ece0bdd3275ebd Mon Sep 17 00:00:00 2001 From: johnseong Date: Fri, 31 Jul 2026 21:41:24 -0400 Subject: [PATCH 07/18] docs(inference): trim verbose comment blocks to load-bearing lines MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Condense the narrative doc/inline comments added across this slice (guardrails, provider, chat_completions, anthropic_messages, config, loader, and the controller CRD types) down to the constraints the code can't show. Also corrects the stale spec.provider field doc — it's the sole routing selector now, modelPreference.provider is informational. Helm CRD regenerated for the trimmed schema descriptions; behaviour unchanged, 1960 tests green. --- controller/src/config_hash.rs | 3 +- controller/src/crd_validations.rs | 5 +- controller/src/inference_policy.rs | 74 ++++--------- controller/src/inference_policy_compile.rs | 7 +- controller/src/inference_policy_reconciler.rs | 9 +- controller/src/reconciler/mod.rs | 6 +- .../kars/templates/crd-inferencepolicy.yaml | 44 +++----- inference-router/src/config.rs | 18 ++-- inference-router/src/guardrails.rs | 101 ++++++------------ .../src/inference_policy_loader.rs | 16 +-- inference-router/src/provider.rs | 38 ++----- .../src/routes/anthropic_messages.rs | 16 +-- .../src/routes/chat_completions.rs | 50 +++------ inference-router/src/routes/mod.rs | 13 +-- 14 files changed, 119 insertions(+), 281 deletions(-) diff --git a/controller/src/config_hash.rs b/controller/src/config_hash.rs index f49dc7d82..b7ac69f1e 100644 --- a/controller/src/config_hash.rs +++ b/controller/src/config_hash.rs @@ -38,8 +38,7 @@ 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 (credentials deliberately excluded — - // this list never hashes secret material, matching the + // Multi-provider endpoints (never secrets — matches the // AZURE_OPENAI_API_KEY precedent). "ANTHROPIC_ENDPOINT", "OLLAMA_ENDPOINT", diff --git a/controller/src/crd_validations.rs b/controller/src/crd_validations.rs index f027c6ca7..7ffa74beb 100644 --- a/controller/src/crd_validations.rs +++ b/controller/src/crd_validations.rs @@ -282,9 +282,8 @@ pub fn inference_policy_validations() -> Vec { reason: Some("FieldValueInvalid".into()), ..ValidationRule::default() }, - // Guardrail pipeline shape: an explicitly-empty list is an - // authoring mistake (delete the field to mean "no pipeline"), - // and a small cap keeps per-request scan fan-out bounded. + // 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()), diff --git a/controller/src/inference_policy.rs b/controller/src/inference_policy.rs index a28056093..67a3ae866 100644 --- a/controller/src/inference_policy.rs +++ b/controller/src/inference_policy.rs @@ -54,25 +54,12 @@ use serde::{Deserialize, Serialize}; use crate::mcp_server::LocalObjectRef; -/// Inference provider selector — the typed form of the kebab-case -/// provider tags this CRD has always documented on [`ModelRef`] -/// (`azure-openai` / `anthropic` / `bedrock` / `ollama`). -/// -/// Serialized with explicit kebab-case renames (NOT `rename_all`) so -/// the wire tags match the existing free-form `ModelRef.provider` -/// strings byte-for-byte — a CR that says `provider: anthropic` under -/// `modelPreference` today can move to the typed field without a -/// migration. -/// -/// Router-side consumption (Slice: multi-provider): `azure-openai` -/// keeps the env-configured Foundry/AOAI upstream (back-compat -/// default); `anthropic` targets the Anthropic Messages API -/// (`ANTHROPIC_ENDPOINT`, key from the router-side secret mount — -/// never visible to the agent process); `ollama` targets an -/// OpenAI-compatible Ollama server (`OLLAMA_ENDPOINT`, no auth). -/// `bedrock` is accepted by the schema for forward-compat but the -/// router rejects it with a clear 501 until the Bedrock client lands -/// — declaring it here keeps the CRD stable across that slice. +/// 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). @@ -93,9 +80,7 @@ pub enum InferenceProvider { } impl InferenceProvider { - /// The kebab-case wire tag — same string serde emits. Exposed for - /// log lines and compile-step JSON so call sites never hand-roll - /// the mapping. + /// The kebab-case wire tag serde emits. #[must_use] pub fn as_tag(&self) -> &'static str { match self { @@ -155,25 +140,18 @@ pub struct InferencePolicySpec { /// [`Self::bundle_ref`]. pub model_preference: Option, - /// Default inference provider for call sites this policy governs. - /// Optional — absent ⇒ the router keeps its env-configured Azure - /// OpenAI / Foundry upstream (back-compat). When set to a - /// non-Azure provider the router swaps the upstream base URL and - /// auth scheme accordingly; provider credentials stay inside the - /// router sidecar (secret mount / env), never in the agent - /// process. `modelPreference.primary.provider`, when it names a - /// recognised tag, takes precedence over this field so a fallback - /// chain can pin its own route. Mutually exclusive with - /// [`Self::bundle_ref`]. + /// 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 - /// inference call (request pre-flight and response — buffered and - /// streaming). Optional — absent ⇒ only the Phase 1 substrate - /// (Foundry guardrail annotations + `contentSafety` floors) - /// applies. Stages run in declaration order; the first stage that - /// flags content blocks the call. Mutually exclusive with - /// [`Self::bundle_ref`]. + /// 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 @@ -302,12 +280,9 @@ pub struct ModelRef { pub deployment: String, } -/// A single stage of the router-side guardrail pipeline. The router -/// materialises each stage into a scanner (network client + policy) -/// at policy load time; a stage whose backend is not configured on -/// the router (e.g. missing moderation API key) fails the *request* -/// closed with an explicit error rather than silently skipping — a -/// declared guardrail that cannot run must never be an open gate. +/// 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 { @@ -318,15 +293,12 @@ pub struct GuardrailStage { pub apply_to: Option, } -/// Guardrail backend. One variant today; Bedrock Guardrails and -/// Model Armor are declared roadmap follow-ups and will extend this -/// enum (adding a variant is a non-breaking CRD change). +/// 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 (`omni-moderation-latest`). Router-side - /// key via `OPENAI_MODERATION_API_KEY` (falls back to - /// `OPENAI_API_KEY`); endpoint override via - /// `OPENAI_MODERATION_ENDPOINT` for Azure-hosted equivalents. + /// OpenAI Moderation API; router-side key via + /// `OPENAI_MODERATION_API_KEY` (falls back to `OPENAI_API_KEY`). #[serde(rename = "openai-moderation")] OpenAIModeration, } diff --git a/controller/src/inference_policy_compile.rs b/controller/src/inference_policy_compile.rs index 28310c007..ae6c630c0 100644 --- a/controller/src/inference_policy_compile.rs +++ b/controller/src/inference_policy_compile.rs @@ -120,11 +120,8 @@ pub fn compile_to_profile(spec: &InferencePolicySpec) -> Value { }) }); - // Provider + guardrails travel as the same kebab-case wire tags - // the CRD serde emits (`InferenceProvider::as_tag` / - // `GuardrailProvider` renames) so the router-side loader parses - // one vocabulary for both the typed field and the free-form - // `modelPreference.*.provider` strings. + // 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| { diff --git a/controller/src/inference_policy_reconciler.rs b/controller/src/inference_policy_reconciler.rs index fda976b30..d7a9eab67 100644 --- a/controller/src/inference_policy_reconciler.rs +++ b/controller/src/inference_policy_reconciler.rs @@ -523,12 +523,9 @@ fn merge_bundle_with_selector( token_budget, content_safety, model_preference, - // The signed-bundle canonical format - // (`policy_canonical::inference`) does not carry - // `provider`/`guardrails` yet — extending that wire contract - // is a coordinated change with the bundle tooling. Until - // then, bundle-sourced policies keep the router defaults on - // these axes. + // 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(), diff --git a/controller/src/reconciler/mod.rs b/controller/src/reconciler/mod.rs index a288519ed..60292d601 100644 --- a/controller/src/reconciler/mod.rs +++ b/controller/src/reconciler/mod.rs @@ -1948,10 +1948,8 @@ async fn reconcile(sandbox: Arc, ctx: Arc) -> Result, /// OpenAI-compatible Ollama endpoint (e.g. @@ -106,12 +103,9 @@ pub struct Config { /// `OPENAI_MODERATION_ENDPOINT`. pub openai_moderation_endpoint: String, - /// API key for the OpenAI Moderation guardrail — - /// `OPENAI_MODERATION_API_KEY` env var (falls back to - /// `OPENAI_API_KEY`, then the secret mounts - /// `/etc/kars/secrets/openai-moderation-api-key` / - /// `/run/secrets/openai-moderation-api-key`). `None` ⇒ policies - /// declaring an `openai-moderation` guardrail stage fail closed. + /// 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 diff --git a/inference-router/src/guardrails.rs b/inference-router/src/guardrails.rs index bbf959266..1df0e9a47 100644 --- a/inference-router/src/guardrails.rs +++ b/inference-router/src/guardrails.rs @@ -1,45 +1,22 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -//! Pluggable guardrail pipeline (multi-cloud guardrails slice). +//! Pluggable guardrail pipeline for `InferencePolicy.spec.guardrails[]`. //! -//! An `InferencePolicy` can declare an ordered list of guardrail -//! stages (`spec.guardrails[]`) that the router runs around every -//! inference call it governs — on the request text before the -//! upstream forward, and on the response text both buffered and -//! streaming. The first backend today is the OpenAI Moderation API; -//! Bedrock Guardrails and Model Armor extend the same [`Guardrail`] -//! trait in follow-up slices. +//! 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 contract +//! 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. //! -//! A *declared* guardrail that cannot run must never become an open -//! gate: -//! -//! - A stage whose backend the router does not recognise, or whose -//! credential/endpoint is missing, fails pipeline construction — -//! the handler rejects the request with an operator-actionable -//! error before any prompt bytes leave the pod. -//! - A scan that errors at runtime (transport failure, non-2xx, -//! unparseable verdict) blocks the request with -//! `guardrail_unavailable` rather than passing unscanned content. -//! -//! ## Streaming semantics (hold-and-release) -//! -//! SSE responses are guarded with a hold-and-release window: chunks -//! are buffered until the accumulated new text reaches -//! [`STREAM_SCAN_THRESHOLD_CHARS`] (or the stream ends), the -//! accumulated text is scanned, and only then is the held window -//! released to the client. No model output is ever delivered before -//! some scan has covered it. On a flagged scan, the client receives a -//! structured SSE error frame + `data: [DONE]` and the upstream -//! stream is dropped. The cost is scan-sized delivery granularity -//! (one moderation round-trip per window), which is the standard -//! trade-off for streaming guardrails. -//! -//! Text longer than [`MAX_SCAN_CHARS`] is scanned in successive -//! windows (see [`scan_windows`]), never truncated, so content can't -//! be hidden past the per-call cap. +//! 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. use std::sync::Arc; @@ -119,11 +96,9 @@ impl Direction { } } -/// One stage as it travels through the compiled policy JSON -/// (`{"provider": "...", "applyTo": "..." | null}`). Parsed liberally -/// by the loader; strictness (unknown backend ⇒ fail closed) applies -/// at pipeline *construction*, where a request is available to -/// reject. +/// 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, @@ -568,13 +543,10 @@ pub fn extract_anthropic_output_text(body: &serde_json::Value) -> String { out.join("\n") } -/// Extract scan text from a body via `extract`, falling back to the -/// raw (lossy-UTF-8) bytes when the body is not JSON. A declared -/// guardrail must never be skipped because a body failed to parse — -/// the raw fallback keeps unknown/malformed shapes covered instead of -/// letting them through unscanned. (A body that parses but yields no -/// extractable text — e.g. tool-call-only responses — is a deliberate -/// pass: the extractors define the scannable surface.) +/// 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) { @@ -624,11 +596,8 @@ fn delta_text_from_event(dialect: StreamDialect, event: &serde_json::Value) -> O } } -/// The client-facing SSE error frame emitted when a stream is cut by -/// a guardrail. OpenAI-style error object works for both dialects' -/// SDK error paths and is what the existing content-safety stream cut -/// emits too. Public so buffered-to-SSE conversion paths (e.g. the -/// Responses-API recovery branch) can emit the same frame shape. +/// 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!( @@ -736,21 +705,16 @@ impl SseGuardState { return; } match serde_json::from_str::(payload) { - // Known text-bearing delta: scan the extracted text. 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 / message_start / role-only / stop); no model - // text to scan. Non-`content` fields (tool-call args, - // "thinking" deltas) are the documented extraction- - // surface gap, not handled here. + // (ping / role-only / stop) — nothing to scan. } - // Non-JSON `data:` payload (upstream drift / malformed - // frame): scan the raw payload rather than fast-releasing - // it — an unrecognised frame must not bypass the 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); @@ -761,9 +725,8 @@ impl SseGuardState { async fn on_chunk(&mut self, chunk: Bytes) -> SseGuardStep { self.ingest_text(&chunk); self.held.push(chunk); - // A buffered partial line's bytes are already in `held` but - // its delta text has not been counted or scanned — never - // release while one is pending, or that text ships unscanned. + // A pending partial line's bytes are in `held` but its text is + // uncounted — never release until the line completes. if !self.line_carry.is_empty() { return SseGuardStep::Release(Vec::new()); } @@ -777,8 +740,7 @@ impl SseGuardState { } async fn on_end(&mut self) -> SseGuardStep { - // Flush a trailing unterminated line (truncated upstream) so - // its text is scanned before the held bytes are released. + // 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); @@ -805,10 +767,9 @@ impl SseGuardState { } } - /// Retain `MAX_SCAN_CHARS - threshold` chars of scanned context - /// after a clean scan: bounds per-connection memory and keeps the - /// next scan (context + up to `threshold` new chars) within one - /// window, with overlap for cross-boundary detection. + /// 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 { diff --git a/inference-router/src/inference_policy_loader.rs b/inference-router/src/inference_policy_loader.rs index 1e8b0ceed..c04b27398 100644 --- a/inference-router/src/inference_policy_loader.rs +++ b/inference-router/src/inference_policy_loader.rs @@ -165,20 +165,12 @@ pub struct LoadedInferencePolicy { /// back to the env-driven default deployment (back-compat). pub model_preference: Option, - /// `spec.provider` — multi-provider slice. Raw kebab-case tag - /// (`azure-openai` / `anthropic` / `ollama` / `bedrock`); - /// interpretation (including the unimplemented-provider - /// fail-closed path) happens per request in - /// [`crate::provider::resolve`] so a policy naming a provider - /// this build can't serve degrades that *request*, not the whole - /// policy load. `None` ⇒ env-driven Azure upstream (back-compat). + /// `spec.provider` — raw kebab-case tag, resolved per request by + /// [`crate::provider::resolve`]. `None` ⇒ Azure (back-compat). pub provider: Option, - /// `spec.guardrails[]` — pluggable guardrail pipeline stages. - /// Empty when the CR omits the block. Stage validity (known - /// backend, credential present) is checked at pipeline build - /// time in [`crate::guardrails::GuardrailPipeline::from_stages`], - /// where a request exists to fail closed. + /// `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 diff --git a/inference-router/src/provider.rs b/inference-router/src/provider.rs index fa5f13d39..06ea713df 100644 --- a/inference-router/src/provider.rs +++ b/inference-router/src/provider.rs @@ -3,26 +3,16 @@ //! Multi-provider upstream resolution. //! -//! Maps `InferencePolicy.spec.provider` (`azure-openai` / `anthropic` -//! / `ollama` / `bedrock`) onto a concrete upstream target: base URL -//! shape + auth scheme. Credentials and endpoints come exclusively -//! from the router's own environment / secret mounts — the agent -//! process never sees a provider API key, exactly as with the Azure -//! Workload Identity path. +//! 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 only routing selector. The pre-existing -//! `modelPreference.primary.provider` tag stays informational (it -//! drove no routing before this slice), so adding provider routing -//! doesn't retroactively reroute CRs that only set a model -//! preference. Absent/empty `spec.provider` ⇒ Azure; an unrecognised -//! tag warns and falls back to Azure. -//! -//! `bedrock` is recognised but not yet implemented: a policy that -//! declares it gets an explicit 501-style error instead of a silent -//! reroute to Azure — declared intent must never be silently ignored. -//! Likewise a resolvable provider with missing router-side config -//! (no `ANTHROPIC_API_KEY`, no `OLLAMA_ENDPOINT`) fails the request -//! closed rather than falling back to Azure. +//! `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; @@ -121,14 +111,8 @@ impl ProviderTarget { } } -/// Resolve the effective provider for a request from `spec.provider`. -/// -/// `spec.provider` is the sole routing selector. The pre-existing -/// `modelPreference.primary.provider` tag stays informational (it -/// drove nothing before this slice) so an unchanged CR that only set -/// a model preference keeps its Azure upstream — routing is opt-in -/// via the new field. Absent/empty ⇒ Azure; an unrecognised tag -/// warns and falls back to Azure; `bedrock` is a hard error. +/// 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) diff --git a/inference-router/src/routes/anthropic_messages.rs b/inference-router/src/routes/anthropic_messages.rs index b90bffb59..20e6c562d 100644 --- a/inference-router/src/routes/anthropic_messages.rs +++ b/inference-router/src/routes/anthropic_messages.rs @@ -286,8 +286,7 @@ pub(super) async fn anthropic_messages( // Slice 2d.1: honour `InferencePolicy.modelPreference.primary.deployment`. crate::routes::apply_model_preference_override(&mut upstream, &policy); - // Multi-provider slice: retarget at the policy-selected provider. - // Fails closed — see `routes::apply_provider_resolution`. + // 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", @@ -305,8 +304,7 @@ pub(super) async fn anthropic_messages( return deny_response(status, &e.to_string(), "api_error"); } - // Guardrail pipeline — build fails closed on declared-but- - // unbuildable stages, input scan runs before any upstream forward. + // Guardrail pipeline; a declared-but-unbuildable stage blocks. let guardrail_pipeline = match super::chat_completions::build_guardrail_pipeline(&state, &policy) { Ok(p) => p, @@ -360,11 +358,8 @@ pub(super) async fn anthropic_messages( } } - // Native Anthropic Messages pass-through — either the policy - // selected `provider: anthropic` (upstream is api.anthropic.com - // with the router-held API key) or the endpoint is GitHub Copilot - // (native /v1/messages). No translation: streaming, tool_use and - // multi-modal content flow through unchanged. + // 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) { @@ -544,8 +539,7 @@ async fn forward_anthropic_passthrough( .await { Ok((status, resp_headers, stream)) => { - // Guardrail output scan (streaming, Anthropic event - // dialect): hold-and-release — see guardrails.rs. + // Streaming output scan (Anthropic event dialect). let guarded = match guardrail_pipeline .as_ref() .filter(|p| p.covers(Direction::Output)) diff --git a/inference-router/src/routes/chat_completions.rs b/inference-router/src/routes/chat_completions.rs index d5d7c477f..794f44bf2 100644 --- a/inference-router/src/routes/chat_completions.rs +++ b/inference-router/src/routes/chat_completions.rs @@ -119,20 +119,16 @@ pub(super) fn build_guardrail_pipeline( .map(|p| Some(Arc::new(p))) } -/// A blocked output scan: either a confirmed content violation or a -/// fail-closed guardrail failure. Kept as data (not a ready-made -/// `Response`) so each transport picks the accurate wire shape — -/// buffered paths map to 403/502/503 HTTP responses, SSE paths to -/// the matching `guardrails::{violation,error}_sse_frame`. +/// 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 — a backend - /// outage must surface as `guardrail_unavailable`, never be - /// mislabelled `content_policy_violation`. + /// 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), @@ -425,9 +421,7 @@ pub(super) async fn chat_completions( // Slice 2d.1: honour `InferencePolicy.modelPreference.primary.deployment`. crate::routes::apply_model_preference_override(&mut upstream, &policy); - // Multi-provider slice: retarget the upstream when the policy - // selects a non-Azure provider. Fails closed on unimplemented / - // unconfigured providers — never a silent Azure fallback. + // 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", @@ -441,13 +435,8 @@ pub(super) async fn chat_completions( return provider_error_response(&e); } - // The Anthropic upstream speaks the Messages API, not - // chat-completions. Anthropic-native runtimes (Claude Agent SDK) - // already target the router's `/anthropic/v1/messages` surface, - // which forwards natively; an OpenAI-shaped client under an - // Anthropic-provider policy gets an explicit 501 (mirrors the - // GitHub-Models 501s for Foundry-only routes) instead of a - // confusing upstream 404. + // Anthropic serves the Messages API — an OpenAI-shaped request + // here gets an explicit 501 pointing at /anthropic/v1/messages. if upstream.provider == ProviderKind::Anthropic { return errors::openai( StatusCode::NOT_IMPLEMENTED, @@ -459,9 +448,7 @@ pub(super) async fn chat_completions( .into_response(); } - // Guardrail pipeline (multi-cloud guardrails slice). Built per - // request from the policy snapshot; a declared stage that cannot - // be materialised blocks the request. + // Guardrail pipeline; a declared-but-unbuildable stage blocks. let guardrail_pipeline = match build_guardrail_pipeline(&state, &policy) { Ok(p) => p, Err(e) => { @@ -610,13 +597,9 @@ pub(super) async fn chat_completions( { budget.record_usage(&sandbox_owned, total).await; } - // Guardrail output scan — the Responses-API - // recovery path must not bypass a declared - // pipeline. This branch already committed to - // text/event-stream, so the block surfaces as - // the outcome-accurate SSE frame (violation - // vs guardrail_unavailable/misconfigured — - // never a mislabelled content violation). + // 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, @@ -937,10 +920,7 @@ pub(super) async fn chat_completions( } chunk }); - // Guardrail output scan (streaming): hold-and-release - // windows — no model text reaches the client before a - // scan has covered it. Skipped entirely (zero cost) - // when the policy declares no output stages. + // Streaming output scan; skipped when no output stages. let guarded: futures::stream::BoxStream<'static, Result> = match guardrail_pipeline .as_ref() @@ -1178,11 +1158,7 @@ pub(super) async fn chat_completions( return resp; } - // Guardrail output scan (buffered) — runs - // beside the contentSafety floor: the floor - // polices upstream-native annotations, the - // pipeline runs the policy's external - // backends (e.g. OpenAI Moderation). + // Buffered output scan (beside the contentSafety floor). if let Err(block) = enforce_openai_output_guardrails( guardrail_pipeline.as_ref(), &resp_body, diff --git a/inference-router/src/routes/mod.rs b/inference-router/src/routes/mod.rs index 078ce8471..6571567a9 100644 --- a/inference-router/src/routes/mod.rs +++ b/inference-router/src/routes/mod.rs @@ -370,16 +370,9 @@ impl AppState { } } -/// Multi-provider slice — retarget `upstream` at the provider the -/// loaded `InferencePolicy` selects. No-op (Azure) when the policy -/// carries no provider opinion, keeping the historic behaviour for -/// every sandbox without a policy. -/// -/// Fails closed: a policy that names a provider the router cannot -/// serve (unimplemented `bedrock`, or missing endpoint/credential -/// config) yields an error the handler must surface — silently -/// falling back to Azure would ship prompts to a provider the -/// operator didn't select. +/// 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, From 650034138f66faf300d4c0ed6034dccc9f740653 Mon Sep 17 00:00:00 2001 From: johnseong Date: Mon, 24 Aug 2026 14:35:36 -0400 Subject: [PATCH 08/18] fix(guardrails): close bypass through Foundry proxy inference routes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The fail-closed route guard covered /v1/completions, /v1/responses, embeddings, and image generation, but not the Foundry proxy families that serve model output: /agents*, /openai/responses*, and /openai/conversations*. An agent blocked by a guardrail on /v1/chat/completions could rerun the same inference through /openai/responses and receive an unscanned response. foundry_proxy now runs guard_unenforced_route first for these path families (whole-segment match, no lookalike false positives) with the same taxonomy as the sibling routes: 403 guardrail_route_unsupported when guardrails are declared, 501 provider_unimplemented for a non-Azure provider, decision headers + route_enforcement_gap audit line included. Non-inference Foundry surfaces (memory stores, files, vector stores, evaluations, ...) are management/storage APIs, not inference channels, and stay unguarded; enforcement scope is now documented in docs/api/crd-reference.md and the CHANGELOG. Tests: unit tests for the path classifier; new tests/foundry_route_guard.rs exercises the real Router::merge wiring from main.rs — guardrail policy 403s all three families, anthropic provider 501s, and control tests prove unaffected requests still reach the proxy body. --- CHANGELOG.md | 12 +- docs/api/crd-reference.md | 2 +- inference-router/src/routes/inference.rs | 80 +++++- inference-router/tests/foundry_route_guard.rs | 247 ++++++++++++++++++ 4 files changed, 335 insertions(+), 6 deletions(-) create mode 100644 inference-router/tests/foundry_route_guard.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 0bddf172e..3e25252ca 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -58,11 +58,15 @@ never sees a provider key. - `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. The sibling inference routes + Anthropic Messages routes only. Every other inference-bearing route (`/v1/completions`, `/v1/responses`, `/v1/embeddings`, image - generation) don't implement either and **fail 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. + 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. Non-inference + Foundry surfaces (memory stores, files, vector stores, …) are + management/storage APIs and stay unguarded. **Inference router — pluggable guardrail pipeline** diff --git a/docs/api/crd-reference.md b/docs/api/crd-reference.md index a2b77ca1d..9eabb90a3 100644 --- a/docs/api/crd-reference.md +++ b/docs/api/crd-reference.md @@ -494,7 +494,7 @@ spec: | `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` / `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`). The sibling inference routes — `/v1/completions`, `/v1/responses`, `/v1/embeddings`, and image generation — do **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. 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. +> **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 remaining Foundry proxy surfaces (memory stores, knowledge bases, evaluations, files, vector stores, containers, and other management/storage APIs) are not general-purpose inference channels and are **not** guarded. 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/inference-router/src/routes/inference.rs b/inference-router/src/routes/inference.rs index 34b8f1940..517bb699a 100644 --- a/inference-router/src/routes/inference.rs +++ b/inference-router/src/routes/inference.rs @@ -691,6 +691,30 @@ async fn list_deployments(State(state): State) -> impl IntoResponse { } } +/// Foundry proxy path families that serve model-generated inference +/// output: agent runs (`/agents*`), the Responses API +/// (`/openai/responses*`), and conversations, which store and return +/// Responses output (`/openai/conversations*`). [`foundry_proxy`] +/// implements neither provider routing nor the guardrail pipeline, so +/// requests to these families 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. The remaining Foundry surfaces +/// (memory stores, knowledge bases, evaluations, files, vector +/// stores, …) are management/storage APIs, not general-purpose +/// inference channels; they stay unguarded — enforcement scope is +/// documented in docs/api/crd-reference.md. +fn inference_bearing_foundry_route(path: &str) -> Option<&'static str> { + let p = path.strip_prefix('/').unwrap_or(path); + ["agents", "openai/responses", "openai/conversations"] + .into_iter() + .find(|family| { + p == *family + || p.strip_prefix(family) + .is_some_and(|rest| rest.starts_with('/')) + }) +} + /// 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/* @@ -703,6 +727,17 @@ async fn foundry_proxy( ) -> impl IntoResponse { let sandbox_name = resolve_sandbox_name(&headers); + // Fail-closed guard for the inference-bearing families this proxy + // serves — must run before anything else so a policy that selects + // a non-Azure provider or declares guardrails cannot be bypassed + // through these routes (see `guard_unenforced_route`). + if let Some(route) = inference_bearing_foundry_route(uri.path()) + && 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 @@ -939,7 +974,50 @@ async fn foundry_proxy( #[cfg(test)] mod tests { - use super::strip_project_prefix; + use super::{inference_bearing_foundry_route, strip_project_prefix}; + + #[test] + fn classifies_inference_bearing_foundry_paths() { + assert_eq!(inference_bearing_foundry_route("/agents"), Some("agents")); + assert_eq!( + inference_bearing_foundry_route("/agents/a1/runs"), + Some("agents") + ); + assert_eq!( + inference_bearing_foundry_route("/openai/responses"), + Some("openai/responses") + ); + assert_eq!( + inference_bearing_foundry_route("/openai/responses/resp_123"), + Some("openai/responses") + ); + assert_eq!( + inference_bearing_foundry_route("/openai/conversations"), + Some("openai/conversations") + ); + assert_eq!( + inference_bearing_foundry_route("/openai/conversations/c1/items"), + Some("openai/conversations") + ); + } + + #[test] + fn leaves_non_inference_foundry_paths_unclassified() { + for path in [ + "/memory_stores", + "/knowledgebases/kb1/queries", + "/openai/files", + "/openai/vector_stores/vs1", + "/openai/containers/c1/files/f1/content", + "/evaluations", + // Prefix must match a whole path segment — no false + // positives on lookalike names. + "/agentsmith", + "/openai/responsesx", + ] { + assert_eq!(inference_bearing_foundry_route(path), None, "path: {path}"); + } + } #[test] fn strips_foundry_project_prefix() { diff --git a/inference-router/tests/foundry_route_guard.rs b/inference-router/tests/foundry_route_guard.rs new file mode 100644 index 000000000..300f6d3df --- /dev/null +++ b/inference-router/tests/foundry_route_guard.rs @@ -0,0 +1,247 @@ +// 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); + + let (status, body) = send(&app, "POST", "/openai/files").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}" + ); +} From 7a1e22688cc188639323bf6565a23efd60a4a73a Mon Sep 17 00:00:00 2001 From: johnseong Date: Tue, 25 Aug 2026 10:59:47 -0400 Subject: [PATCH 09/18] refactor(controller): extract reconciler pod-spec helpers into submodule Move build_pod_security_context, isolation_scheduling, and build_egress_guard_command (plus their tests) out of reconciler/mod.rs into reconciler/pod_spec.rs, re-exported via pub(crate) use. Brings reconciler/mod.rs back under the 3700-line CI cap. Pure move, no behavior change. --- controller/src/reconciler/mod.rs | 220 +------------------------ controller/src/reconciler/pod_spec.rs | 226 ++++++++++++++++++++++++++ 2 files changed, 231 insertions(+), 215 deletions(-) create mode 100644 controller/src/reconciler/pod_spec.rs diff --git a/controller/src/reconciler/mod.rs b/controller/src/reconciler/mod.rs index 60292d601..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}")] 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}" + ); + } + } +} From 720456b306e9a38ed3fad8f987376b955e2e64ec Mon Sep 17 00:00:00 2001 From: johnseong Date: Tue, 25 Aug 2026 11:00:04 -0400 Subject: [PATCH 10/18] fix(guardrails): fail-closed pipeline hardening + decompose module Correctness: - from_compiled_json now returns Result; a present-but-malformed guardrails block poisons the policy so every request fails closed instead of silently degrading to "no guardrails" (which also disarmed the sibling route-gap guard). The loader installs an unbuildable sentinel stage on parse error. - SSE moderation reconstructs multi-byte UTF-8 code points split across chunk boundaries, so the scanner sees the same text the client receives (no U+FFFD divergence). - Streaming delta extraction scans every choices[] entry, not just choices[0]. Refactor: - Split guardrails.rs (>1600 lines) into guardrails/{mod,backend, stream,tests}.rs, each under the 800-line CI cap. Public API is unchanged (re-exported); a few items raised to pub(crate) so the test module can construct and inspect them. Regression tests for the UTF-8 split (flagged and benign), multi-choice extraction, and malformed-config fail-closed. --- inference-router/src/guardrails.rs | 1490 ----------------- inference-router/src/guardrails/backend.rs | 372 ++++ inference-router/src/guardrails/mod.rs | 225 +++ inference-router/src/guardrails/stream.rs | 406 +++++ inference-router/src/guardrails/tests.rs | 698 ++++++++ .../src/inference_policy_loader.rs | 24 +- 6 files changed, 1723 insertions(+), 1492 deletions(-) delete mode 100644 inference-router/src/guardrails.rs create mode 100644 inference-router/src/guardrails/backend.rs create mode 100644 inference-router/src/guardrails/mod.rs create mode 100644 inference-router/src/guardrails/stream.rs create mode 100644 inference-router/src/guardrails/tests.rs diff --git a/inference-router/src/guardrails.rs b/inference-router/src/guardrails.rs deleted file mode 100644 index 1df0e9a47..000000000 --- a/inference-router/src/guardrails.rs +++ /dev/null @@ -1,1490 +0,0 @@ -// 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. - -use std::sync::Arc; - -use async_trait::async_trait; -use bytes::Bytes; -use futures::stream::BoxStream; -use futures::stream::StreamExt; -use reqwest::Client; - -use crate::config::Config; -use crate::metrics; - -/// 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 (array | null | absent). - /// Entries without a string `provider` are dropped with a WARN — - /// they cannot be built into anything enforceable and the - /// controller schema rejects them at admission anyway. - #[must_use] - pub fn from_compiled_json(v: &serde_json::Value) -> Vec { - let Some(arr) = v.as_array() else { - return Vec::new(); - }; - arr.iter() - .filter_map(|stage| { - let Some(provider) = stage.get("provider").and_then(|p| p.as_str()) else { - tracing::warn!( - stage = %stage, - "guardrail stage missing string 'provider' — dropped" - ); - return None; - }; - Some(Self { - provider: provider.to_string(), - apply_to: ApplyTo::parse(stage.get("applyTo").and_then(|a| a.as_str())), - }) - }) - .collect() - } -} - -/// 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; -} - -// ─── 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 ──────────────────────────────────────────────────────────────── - -struct BuiltStage { - apply_to: ApplyTo, - guard: Arc, -} - -/// Ordered guardrail stages materialised from a policy snapshot. -/// Cheap to build per request (clones a shared `reqwest::Client`). -pub struct GuardrailPipeline { - 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. -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(), - } -} - -// ─── 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] -fn delta_text_from_event(dialect: StreamDialect, event: &serde_json::Value) -> Option { - match dialect { - StreamDialect::OpenAiChat => event - .get("choices") - .and_then(|c| c.as_array()) - .and_then(|c| c.first()) - .and_then(|c| c.get("delta")) - .and_then(|d| d.get("content")) - .and_then(|t| t.as_str()) - .filter(|t| !t.is_empty()) - .map(str::to_string), - 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`]. -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 `data:` lines split across chunk boundaries. - line_carry: String, - /// All delta text accumulated so far (scan context). - accumulated: String, - /// Chars of `accumulated` not yet covered by a scan. - unscanned: usize, -} - -/// What the state machine wants the adaptor to emit next. -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 { - fn new(pipeline: Arc, dialect: StreamDialect, threshold: usize) -> Self { - Self { - pipeline, - dialect, - threshold, - held: Vec::new(), - line_carry: String::new(), - accumulated: String::new(), - unscanned: 0, - } - } - - fn ingest_text(&mut self, chunk: &[u8]) { - self.line_carry.push_str(&String::from_utf8_lossy(chunk)); - 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); - } - } - } - - async fn on_chunk(&mut self, chunk: Bytes) -> SseGuardStep { - self.ingest_text(&chunk); - self.held.push(chunk); - // A pending partial line's bytes are in `held` but its text is - // uncounted — never release until the line completes. - if !self.line_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 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) -} - -#[cfg(test)] -mod tests { - use super::*; - use std::sync::atomic::{AtomicUsize, Ordering}; - - // ---- 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 }, - { "applyTo": "input" } // dropped: no provider - ]); - let stages = GuardrailStageCfg::from_compiled_json(&v); - 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_handles_null_and_absent() { - assert!(GuardrailStageCfg::from_compiled_json(&serde_json::Value::Null).is_empty()); - assert!(GuardrailStageCfg::from_compiled_json(&serde_json::json!({})).is_empty()); - } - - // ---- 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 - ); - } - - // ---- pipeline + streaming with a fake backend ---- - - /// Test backend: flags any text containing the marker. Counts - /// scans so tests can assert hold-and-release windowing. - 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"); - } -} 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 c04b27398..bf7b1efbb 100644 --- a/inference-router/src/inference_policy_loader.rs +++ b/inference-router/src/inference_policy_loader.rs @@ -344,9 +344,29 @@ pub fn load_inference_policy_from_dir( .and_then(|p| p.as_str()) .filter(|p| !p.trim().is_empty()) .map(str::to_string); - let guardrails = crate::guardrails::GuardrailStageCfg::from_compiled_json( + // `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. From be5cb8beec320ab9c1ad94dc9befba86b832e8d4 Mon Sep 17 00:00:00 2001 From: johnseong Date: Tue, 25 Aug 2026 11:00:19 -0400 Subject: [PATCH 11/18] fix(inference): route fail-closed enforcement, Anthropic streaming usage, unified error codes Security / fail-closed: - foundry_proxy percent-decodes each path segment and rejects (400 invalid_path) any dot / empty / encoded-slash segment before classification or forwarding, closing a traversal bypass where /openai/files/../responses normalized upstream into a guarded inference route. The guard is now default-deny: everything is guarded except an explicit exempt set of management/storage APIs. - Buffered output guardrail enforcement is 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. Fixes: - Anthropic streaming passthrough records token usage (input from message_start, output from the latest message_delta), closing a budget-accounting gap; multi-byte UTF-8 is carried across chunk boundaries like the guardrail stream. - Provider/guardrail denials carry a stable error.code across chat-completions, the sibling routes, and the Anthropic route via a shared errors::openai_coded helper. CI: - Box the output-guardrail error response (clippy result_large_err under -D warnings) and collapse an else-if block. Tests: traversal matrix + reqwest::Url normalization premise, default-deny lookalikes, non-JSON buffered output scan, Anthropic streaming usage. --- inference-router/src/errors.rs | 26 ++ .../src/routes/anthropic_messages.rs | 189 ++++++++++++++ .../src/routes/chat_completions.rs | 49 ++-- inference-router/src/routes/inference.rs | 219 +++++++++++++---- inference-router/src/routes/mod.rs | 30 ++- .../tests/chat_output_guardrail_nonjson.rs | 232 ++++++++++++++++++ inference-router/tests/foundry_route_guard.rs | 127 +++++++++- 7 files changed, 787 insertions(+), 85 deletions(-) create mode 100644 inference-router/tests/chat_output_guardrail_nonjson.rs 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/routes/anthropic_messages.rs b/inference-router/src/routes/anthropic_messages.rs index 20e6c562d..4d56242f1 100644 --- a/inference-router/src/routes/anthropic_messages.rs +++ b/inference-router/src/routes/anthropic_messages.rs @@ -50,12 +50,16 @@ fn is_hop_by_hop(name: &str) -> bool { } 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. ( status, Json(json!({ "type": "error", "error": { "type": code, + "code": code, "message": message, } })), @@ -494,6 +498,136 @@ 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; + } + } + _ => {} + } +} + +/// Best-effort token accounting for a streamed Anthropic passthrough. +/// Taps the SSE stream for `input_tokens`/`output_tokens` and records +/// the total with the budget tracker at end-of-stream, so streamed +/// Anthropic inference feeds the same daily/monthly accumulation as the +/// buffered path. Every byte passes through unchanged; usage is +/// recorded once (also on an upstream error, with whatever was seen). +fn tap_anthropic_stream_usage( + stream: futures::stream::BoxStream<'static, Result>, + budget: crate::budget::TokenBudgetTracker, + sandbox: String, +) -> futures::stream::BoxStream<'static, Result> +where + E: Send + 'static, +{ + struct Ctx { + inner: futures::stream::BoxStream<'static, Result>, + budget: crate::budget::TokenBudgetTracker, + sandbox: String, + // 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, + input: u64, + output: u64, + recorded: bool, + } + async fn record(ctx: &mut Ctx) { + if ctx.recorded { + return; + } + ctx.recorded = true; + let total = ctx.input + ctx.output; + if total > 0 { + ctx.budget.record_usage(&ctx.sandbox, total).await; + } + } + 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) { + update_anthropic_usage(&ev, &mut ctx.input, &mut ctx.output); + } + } + } + let ctx = Ctx { + inner: stream, + budget, + sandbox, + byte_carry: Vec::new(), + line_carry: String::new(), + input: 0, + output: 0, + recorded: false, + }; + futures::stream::unfold(ctx, |mut ctx| async move { + match ctx.inner.next().await { + Some(Ok(chunk)) => { + // Decode only complete code points; carry a trailing + // incomplete UTF-8 sequence to the next 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)) => { + record(&mut ctx).await; + 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); + record(&mut ctx).await; + None + } + } + }) + .boxed() +} + /// Native passthrough for Copilot's Anthropic Messages API. /// /// No translation: forwards body verbatim to `{copilot_endpoint}/v1/messages`, @@ -539,6 +673,13 @@ async fn forward_anthropic_passthrough( .await { Ok((status, resp_headers, stream)) => { + // Tap usage BEFORE guarding so input tokens are still + // recorded if the guard later cuts the stream. + let stream = tap_anthropic_stream_usage( + stream, + state.budget.clone(), + sandbox_name.to_string(), + ); // Streaming output scan (Anthropic event dialect). let guarded = match guardrail_pipeline .as_ref() @@ -781,4 +922,52 @@ 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(); + let tapped = tap_anthropic_stream_usage( + futures::stream::iter(chunks).boxed(), + budget.clone(), + "sbx".into(), + ); + // Drain the stream (client role). + let _drained: Vec<_> = tapped.collect().await; + let (used, _) = budget.get_usage("sbx").await; + assert_eq!(used, 42, "30 input + 12 output recorded once at stream end"); + } } diff --git a/inference-router/src/routes/chat_completions.rs b/inference-router/src/routes/chat_completions.rs index 794f44bf2..f8006458a 100644 --- a/inference-router/src/routes/chat_completions.rs +++ b/inference-router/src/routes/chat_completions.rs @@ -184,16 +184,20 @@ pub(super) async fn scan_openai_output_guardrails( /// 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<(), axum::response::Response> { +) -> Result<(), Box> { match scan_openai_output_guardrails(pipeline, resp_body, sandbox_name, policy_digest).await { None => Ok(()), - Some(OutputGuardrailBlock::Violation(v)) => Err(guardrail_violation_response(&v)), - Some(OutputGuardrailBlock::Error(e)) => Err(guardrail_error_response(&e)), + Some(OutputGuardrailBlock::Violation(v)) => Err(Box::new(guardrail_violation_response(&v))), + Some(OutputGuardrailBlock::Error(e)) => Err(Box::new(guardrail_error_response(&e))), } } @@ -265,7 +269,9 @@ pub(super) async fn guard_unenforced_route( route, "{msg}" ); - let mut resp = errors::openai(status, &msg, code).into_response(); + // 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) } @@ -673,7 +679,7 @@ pub(super) async fn chat_completions( ) .await { - return block; + return *block; } let mut response = (resp_status, Body::from(chat_body)).into_response(); if let Some(ct) = resp_hdrs.get("content-type") { @@ -784,7 +790,7 @@ pub(super) async fn chat_completions( ) .await { - return block; + return *block; } // Wrap as SSE so the streaming client can parse it let sse = format!( @@ -1027,7 +1033,7 @@ pub(super) async fn chat_completions( ) .await { - return block; + return *block; } let mut response = (resp_status, Body::from(chat_body)).into_response(); if let Some(ct) = resp_hdrs.get("content-type") { @@ -1068,6 +1074,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 @@ -1158,18 +1181,6 @@ pub(super) async fn chat_completions( return resp; } - // Buffered output scan (beside the contentSafety floor). - if let Err(block) = enforce_openai_output_guardrails( - guardrail_pipeline.as_ref(), - &resp_body, - sandbox_name, - &policy.digest, - ) - .await - { - return block; - } - // AGT output pipeline: redact → scan → policy check (blocking) let response_text = body_json .get("choices") diff --git a/inference-router/src/routes/inference.rs b/inference-router/src/routes/inference.rs index 517bb699a..baa5e71bd 100644 --- a/inference-router/src/routes/inference.rs +++ b/inference-router/src/routes/inference.rs @@ -691,28 +691,94 @@ async fn list_deployments(State(state): State) -> impl IntoResponse { } } -/// Foundry proxy path families that serve model-generated inference -/// output: agent runs (`/agents*`), the Responses API -/// (`/openai/responses*`), and conversations, which store and return -/// Responses output (`/openai/conversations*`). [`foundry_proxy`] -/// implements neither provider routing nor the guardrail pipeline, so -/// requests to these families 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. The remaining Foundry surfaces -/// (memory stores, knowledge bases, evaluations, files, vector -/// stores, …) are management/storage APIs, not general-purpose -/// inference channels; they stay unguarded — enforcement scope is -/// documented in docs/api/crd-reference.md. -fn inference_bearing_foundry_route(path: &str) -> Option<&'static str> { - let p = path.strip_prefix('/').unwrap_or(path); - ["agents", "openai/responses", "openai/conversations"] - .into_iter() - .find(|family| { - p == *family - || p.strip_prefix(family) - .is_some_and(|rest| rest.starts_with('/')) - }) +/// 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), empty segments (`//`, trailing `/`), and decoded +/// slashes smuggled inside a segment (`..%2f`). [`foundry_proxy`] +/// concatenates the raw path into the upstream URL, where URL parsing +/// resolves dot segments, so classification and forwarding would +/// otherwise disagree. Foundry API paths never legitimately contain +/// any of these forms. +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() || 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. @@ -727,11 +793,34 @@ async fn foundry_proxy( ) -> impl IntoResponse { let sandbox_name = resolve_sandbox_name(&headers); - // Fail-closed guard for the inference-bearing families this proxy - // serves — must run before anything else so a policy that selects - // a non-Azure provider or declares guardrails cannot be bypassed - // through these routes (see `guard_unenforced_route`). - if let Some(route) = inference_bearing_foundry_route(uri.path()) + // 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. + // Reject every form that could normalize away before anything + // else runs. + 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 dot, empty, or encoded-slash segments" + ); + return errors::openai_coded( + StatusCode::BAD_REQUEST, + "path contains dot, empty, or encoded-slash segments", + "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 { @@ -974,51 +1063,89 @@ async fn foundry_proxy( #[cfg(test)] mod tests { - use super::{inference_bearing_foundry_route, 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!(inference_bearing_foundry_route("/agents"), Some("agents")); + assert_eq!(classify("/agents"), Some("agents")); + assert_eq!(classify("/agents/a1/runs"), Some("agents")); + assert_eq!(classify("/openai/responses"), Some("openai/responses")); assert_eq!( - inference_bearing_foundry_route("/agents/a1/runs"), - Some("agents") - ); - assert_eq!( - inference_bearing_foundry_route("/openai/responses"), + classify("/openai/responses/resp_123"), Some("openai/responses") ); assert_eq!( - inference_bearing_foundry_route("/openai/responses/resp_123"), - Some("openai/responses") - ); - assert_eq!( - inference_bearing_foundry_route("/openai/conversations"), + classify("/openai/conversations"), Some("openai/conversations") ); assert_eq!( - inference_bearing_foundry_route("/openai/conversations/c1/items"), + classify("/openai/conversations/c1/items"), Some("openai/conversations") ); } #[test] - fn leaves_non_inference_foundry_paths_unclassified() { + 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", - // Prefix must match a whole path segment — no false - // positives on lookalike names. - "/agentsmith", - "/openai/responsesx", + "/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//responses", // empty segment + "/openai/responses/", // trailing empty segment + "/openai/files/%zz", // malformed escape + "/openai/files/%2", // truncated escape ] { - assert_eq!(inference_bearing_foundry_route(path), None, "path: {path}"); + assert_eq!(decoded_path_segments(path), None, "path: {path}"); } } + #[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() { assert_eq!( diff --git a/inference-router/src/routes/mod.rs b/inference-router/src/routes/mod.rs index 6571567a9..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()); 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/foundry_route_guard.rs b/inference-router/tests/foundry_route_guard.rs index 300f6d3df..749205dc2 100644 --- a/inference-router/tests/foundry_route_guard.rs +++ b/inference-router/tests/foundry_route_guard.rs @@ -237,11 +237,130 @@ async fn guardrail_policy_leaves_non_inference_foundry_routes_unguarded() { install_policy(&state, None, true).await; let app = app(state); - let (status, body) = send(&app, "POST", "/openai/files").await; - assert_eq!(status, StatusCode::NOT_IMPLEMENTED, "{body}"); + 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 / empty-segment 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. +#[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", + // Empty segment (double slash) through a wildcard. + "/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}" + ); + } +} + +/// 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("unsupported_for_provider"), - "expected the proxy's own github-models 501 (proof the guard let it through): {body}" + Some("guardrail_route_unsupported"), + "{body}" ); } From 708b2401d02afadcc9fc9a311a42cabdcbf2b74e Mon Sep 17 00:00:00 2001 From: johnseong Date: Tue, 25 Aug 2026 11:00:19 -0400 Subject: [PATCH 12/18] docs(inference): security audit + enforcement-scope updates - Add docs/security-audits/2026-08-25-multi-provider-guardrails.md (T1/T2/T3 triage + sign-offs) required by the security-audit CI gate for capability-path changes. - Document the default-deny Foundry proxy guard and path-rejection behavior in the CRD reference and CHANGELOG. --- CHANGELOG.md | 9 +- docs/api/crd-reference.md | 2 +- .../2026-08-25-multi-provider-guardrails.md | 112 ++++++++++++++++++ 3 files changed, 119 insertions(+), 4 deletions(-) create mode 100644 docs/security-audits/2026-08-25-multi-provider-guardrails.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e25252ca..14c35b007 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -64,9 +64,12 @@ never sees a provider key. `/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. Non-inference - Foundry surfaces (memory stores, files, vector stores, …) are - management/storage APIs and stay unguarded. + 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 with dot/empty/percent-encoded-dot segments + are rejected (400) so a request cannot normalize into a different + upstream route than the one classified. **Inference router — pluggable guardrail pipeline** diff --git a/docs/api/crd-reference.md b/docs/api/crd-reference.md index 9eabb90a3..ad15a9a7b 100644 --- a/docs/api/crd-reference.md +++ b/docs/api/crd-reference.md @@ -494,7 +494,7 @@ spec: | `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` / `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 remaining Foundry proxy surfaces (memory stores, knowledge bases, evaluations, files, vector stores, containers, and other management/storage APIs) are not general-purpose inference channels and are **not** guarded. 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. +> **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 containing dot segments, empty segments, or percent-encoded dot/slash forms 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/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..ca575bd59 --- /dev/null +++ b/docs/security-audits/2026-08-25-multi-provider-guardrails.md @@ -0,0 +1,112 @@ +# 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; the + existing blocklist/egress guard still applies. `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 and rejects (400 `invalid_path`) any dot / empty / + encoded-slash segment 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 From db03955116b5473e7c1b4fc3db117ac0c1efe91e Mon Sep 17 00:00:00 2001 From: johnseong Date: Wed, 26 Aug 2026 21:44:15 -0400 Subject: [PATCH 13/18] fix(inference): Anthropic streaming usage across a mid-stream cut, coded errors + decision headers Streaming usage (MEDIUM): the usage tap sat inside guard_sse_stream, so a mid-stream guardrail cut dropped it before its terminal ran and budget stayed zero despite consumed tokens. Split into an observer (inside the guard, updates shared usage) and a finalizer (outside the guard, records once on any terminal: cut, error, or clean end). Error contract (MEDIUM + cleanup): buffered Anthropic denials now carry the stable machine code (guardrail_blocked / guardrail_misconfigured / guardrail_unavailable / provider_*) in error.code while keeping the Anthropic-native error.type, via a new deny_policy helper that also attaches the x-kars-decision* headers. The chat provider:anthropic 501 now uses the coded helper + decision headers too. Tests: mid-stream-cut accounting (input billed once through a real guard cut), buffered coded-error shape, and decision-header presence only on policy denials. --- .../src/routes/anthropic_messages.rs | 369 +++++++++++++++--- .../src/routes/chat_completions.rs | 16 +- 2 files changed, 317 insertions(+), 68 deletions(-) diff --git a/inference-router/src/routes/anthropic_messages.rs b/inference-router/src/routes/anthropic_messages.rs index 4d56242f1..cdf37c4df 100644 --- a/inference-router/src/routes/anthropic_messages.rs +++ b/inference-router/src/routes/anthropic_messages.rs @@ -25,7 +25,7 @@ use futures::stream::StreamExt; use serde_json::{Value, json}; use super::AppState; -use crate::guardrails::{self, Direction, GuardrailPipeline}; +use crate::guardrails::{self, Direction, GuardrailError, GuardrailPipeline}; use crate::provider::{ProviderError, ProviderKind}; use crate::proxy; use std::sync::Arc; @@ -52,19 +52,73 @@ fn is_hop_by_hop(name: &str) -> bool { 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. - ( + // `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. @@ -305,7 +359,7 @@ pub(super) async fn anthropic_messages( ProviderError::Unimplemented { .. } => StatusCode::NOT_IMPLEMENTED, _ => StatusCode::SERVICE_UNAVAILABLE, }; - return deny_response(status, &e.to_string(), "api_error"); + return deny_policy(status, &e.to_string(), "api_error", provider_error_code(&e)); } // Guardrail pipeline; a declared-but-unbuildable stage blocks. @@ -322,7 +376,12 @@ pub(super) async fn anthropic_messages( error = %e, "guardrail pipeline could not be built (anthropic route) — failing closed" ); - return deny_response(StatusCode::SERVICE_UNAVAILABLE, &e.to_string(), "api_error"); + return deny_policy( + guardrail_error_status(&e), + &e.to_string(), + "api_error", + e.code(), + ); } }; if let Some(ref p) = guardrail_pipeline @@ -341,10 +400,11 @@ pub(super) async fn anthropic_messages( categories = ?v.categories, "guardrail pipeline blocked request (anthropic route)" ); - return deny_response( + return deny_policy( StatusCode::FORBIDDEN, &v.message(), "content_policy_violation", + v.code(), ); } Err(e) => { @@ -357,7 +417,12 @@ pub(super) async fn anthropic_messages( error = %e, "guardrail pipeline unavailable (anthropic route) — failing closed" ); - return deny_response(StatusCode::BAD_GATEWAY, &e.to_string(), "api_error"); + return deny_policy( + guardrail_error_status(&e), + &e.to_string(), + "api_error", + e.code(), + ); } } } @@ -464,10 +529,11 @@ pub(super) async fn anthropic_messages( categories = ?v.categories, "guardrail pipeline blocked translated response (anthropic route)" ); - return deny_response( + return deny_policy( StatusCode::FORBIDDEN, &v.message(), "content_policy_violation", + v.code(), ); } Err(e) => { @@ -480,7 +546,12 @@ pub(super) async fn anthropic_messages( error = %e, "guardrail pipeline unavailable (anthropic route) — failing closed" ); - return deny_response(StatusCode::BAD_GATEWAY, &e.to_string(), "api_error"); + return deny_policy( + guardrail_error_status(&e), + &e.to_string(), + "api_error", + e.code(), + ); } } } @@ -527,42 +598,41 @@ fn update_anthropic_usage(ev: &Value, input: &mut u64, output: &mut u64) { } } -/// Best-effort token accounting for a streamed Anthropic passthrough. -/// Taps the SSE stream for `input_tokens`/`output_tokens` and records -/// the total with the budget tracker at end-of-stream, so streamed -/// Anthropic inference feeds the same daily/monthly accumulation as the -/// buffered path. Every byte passes through unchanged; usage is -/// recorded once (also on an upstream error, with whatever was seen). -fn tap_anthropic_stream_usage( +/// 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>, - budget: crate::budget::TokenBudgetTracker, - sandbox: String, + usage: Arc>, ) -> futures::stream::BoxStream<'static, Result> where E: Send + 'static, { struct Ctx { inner: futures::stream::BoxStream<'static, Result>, - budget: crate::budget::TokenBudgetTracker, - sandbox: String, + 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, - input: u64, - output: u64, - recorded: bool, - } - async fn record(ctx: &mut Ctx) { - if ctx.recorded { - return; - } - ctx.recorded = true; - let total = ctx.input + ctx.output; - if total > 0 { - ctx.budget.record_usage(&ctx.sandbox, total).await; - } } fn scan_lines(ctx: &mut Ctx, text: &str) { for line in text.lines() { @@ -574,25 +644,21 @@ where continue; } if let Ok(ev) = serde_json::from_str::(payload) { - update_anthropic_usage(&ev, &mut ctx.input, &mut ctx.output); + let mut u = ctx.usage.lock().unwrap(); + let StreamUsage { input, output } = &mut *u; + update_anthropic_usage(&ev, input, output); } } } let ctx = Ctx { inner: stream, - budget, - sandbox, + usage, byte_carry: Vec::new(), line_carry: String::new(), - input: 0, - output: 0, - recorded: false, }; futures::stream::unfold(ctx, |mut ctx| async move { match ctx.inner.next().await { Some(Ok(chunk)) => { - // Decode only complete code points; carry a trailing - // incomplete UTF-8 sequence to the next chunk. let mut bytes = std::mem::take(&mut ctx.byte_carry); bytes.extend_from_slice(&chunk); let decodable = match std::str::from_utf8(&bytes) { @@ -609,10 +675,7 @@ where } Some((Ok(chunk), ctx)) } - Some(Err(e)) => { - record(&mut ctx).await; - Some((Err(e), ctx)) - } + Some(Err(e)) => Some((Err(e), ctx)), None => { if !ctx.byte_carry.is_empty() { let tail = std::mem::take(&mut ctx.byte_carry); @@ -620,6 +683,61 @@ where } 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 } @@ -673,28 +791,32 @@ async fn forward_anthropic_passthrough( .await { Ok((status, resp_headers, stream)) => { - // Tap usage BEFORE guarding so input tokens are still - // recorded if the guard later cuts the stream. - let stream = tap_anthropic_stream_usage( - stream, - state.budget.clone(), - sandbox_name.to_string(), - ); + // 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( - stream, + observed, p.clone(), guardrails::StreamDialect::AnthropicMessages, sandbox_name.to_string(), policy_digest.clone(), ), - None => stream, + None => observed, }; - let body = Body::from_stream(guarded.map(|c| c.map_err(std::io::Error::other))); + 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() { @@ -960,14 +1082,135 @@ mod tests { ]; let chunks: Vec> = frames.iter().map(|f| Ok(Bytes::from(*f))).collect(); - let tapped = tap_anthropic_stream_usage( - futures::stream::iter(chunks).boxed(), - budget.clone(), + // 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(), ); - // Drain the stream (client role). - let _drained: Vec<_> = tapped.collect().await; + 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, 42, "30 input + 12 output recorded once at stream end"); + 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 f8006458a..62d539a9a 100644 --- a/inference-router/src/routes/chat_completions.rs +++ b/inference-router/src/routes/chat_completions.rs @@ -31,7 +31,7 @@ use std::sync::Arc; /// 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, @@ -443,15 +443,21 @@ pub(super) async fn chat_completions( // 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 { - return errors::openai( - StatusCode::NOT_IMPLEMENTED, - "InferencePolicy selects provider 'anthropic', which serves the 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)", + (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. From b8314825de0a187049c6c1d2a2ad004f9bb976df Mon Sep 17 00:00:00 2001 From: johnseong Date: Wed, 26 Aug 2026 21:44:15 -0400 Subject: [PATCH 14/18] fix(inference): collapse benign trailing/double slashes instead of rejecting COMPATIBILITY: decoded_path_segments rejected empty segments outright, so historically proxied management paths (POST /openai/files/f1/, /memory_stores/s1/, /openai/files//f1) returned 400 even with no policy configured. Collapse benign empty segments to canonical form while still rejecting dot segments (., .., percent-encoded), smuggled slashes/backslashes, and malformed escapes. Adds no-policy compat tests across guardrails on/off. --- inference-router/src/routes/inference.rs | 47 +++++++++++++++---- inference-router/tests/foundry_route_guard.rs | 45 ++++++++++++++++-- 2 files changed, 79 insertions(+), 13 deletions(-) diff --git a/inference-router/src/routes/inference.rs b/inference-router/src/routes/inference.rs index baa5e71bd..ba4ef4200 100644 --- a/inference-router/src/routes/inference.rs +++ b/inference-router/src/routes/inference.rs @@ -715,17 +715,24 @@ fn percent_decode_segment(seg: &str) -> Option { /// 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), empty segments (`//`, trailing `/`), and decoded -/// slashes smuggled inside a segment (`..%2f`). [`foundry_proxy`] -/// concatenates the raw path into the upstream URL, where URL parsing -/// resolves dot segments, so classification and forwarding would -/// otherwise disagree. Foundry API paths never legitimately contain -/// any of these forms. +/// 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() || seg == "." || seg == ".." || seg.contains('/') || seg.contains('\\') { + if seg.is_empty() { + continue; // collapse `//` and trailing `/` + } + if seg == "." || seg == ".." || seg.contains('/') || seg.contains('\\') { return None; } segments.push(seg); @@ -1126,8 +1133,6 @@ mod tests { "/memory_stores/../openai/responses", "/openai/files/..%2fresponses", // decoded slash inside a segment "/openai/files/..%5cresponses", // decoded backslash - "/openai//responses", // empty segment - "/openai/responses/", // trailing empty segment "/openai/files/%zz", // malformed escape "/openai/files/%2", // truncated escape ] { @@ -1135,6 +1140,30 @@ mod tests { } } + #[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 diff --git a/inference-router/tests/foundry_route_guard.rs b/inference-router/tests/foundry_route_guard.rs index 749205dc2..41705e653 100644 --- a/inference-router/tests/foundry_route_guard.rs +++ b/inference-router/tests/foundry_route_guard.rs @@ -266,10 +266,12 @@ fn url_parsing_normalizes_dot_segments() { } } -/// A dot-segment / encoded-dot / empty-segment 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. +/// 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(); @@ -346,6 +348,41 @@ async fn slash_variants_of_guarded_families_never_reach_proxy_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 From 6461396993f4ca6d988d0de854c5651cf1e0e383 Mon Sep 17 00:00:00 2001 From: johnseong Date: Wed, 26 Aug 2026 21:44:15 -0400 Subject: [PATCH 15/18] docs(security): correct audit sign-off and provider egress claim - Remove the maintainer sign-off added by the contributor; leave only the author sign-off. The independent reviewer adds theirs via a maintainer-owned commit. - Correct the T1 claim: the Anthropic/Ollama provider forward paths do not invoke the blocklist (is_blocked); destinations are operator-controlled but not blocklist-checked. --- .../2026-08-25-multi-provider-guardrails.md | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/docs/security-audits/2026-08-25-multi-provider-guardrails.md b/docs/security-audits/2026-08-25-multi-provider-guardrails.md index ca575bd59..53a59b613 100644 --- a/docs/security-audits/2026-08-25-multi-provider-guardrails.md +++ b/docs/security-audits/2026-08-25-multi-provider-guardrails.md @@ -23,9 +23,11 @@ fail-closed/accounting fixes. - **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; the - existing blocklist/egress guard still applies. `provider::resolve` fails closed - (501 unimplemented / 503 unconfigured) rather than defaulting to an + (`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 @@ -109,4 +111,3 @@ credential surface is router-side-only, bounded, and default-deny; the identifie bypass and functional gaps are fixed and regression-tested. Signed-off-by: John Seong -Signed-off-by: Pal Lakatos-Toth From 4e20aa5d8ad6597761811bd38eae38b27e0b8a82 Mon Sep 17 00:00:00 2001 From: johnseong Date: Thu, 27 Aug 2026 17:46:00 -0400 Subject: [PATCH 16/18] fix(inference): native buffered Anthropic guardrail denials carry coded errors + decision headers forward_anthropic_passthrough's native buffered output branches still used the legacy deny_response (content_policy_violation / generic api_error with a hard-coded 502, no decision headers), so an identical request returned the stable codes when streamed but legacy/generic codes when buffered. Route both branches through deny_policy: v.code() for violations, guardrail_error_status(&e) + e.code() for errors, matching the translated and streaming paths. This was the last guardrail/provider denial site not on the coded contract. Adds route-level tests that drive the real native buffered /v1/messages handler (wiremock Anthropic upstream + moderation): a flagged output gives 403 guardrail_blocked and a moderation outage gives 502 guardrail_unavailable, both with x-kars-decision headers. --- .../src/routes/anthropic_messages.rs | 8 +- .../tests/anthropic_buffered_guardrail.rs | 265 ++++++++++++++++++ 2 files changed, 270 insertions(+), 3 deletions(-) create mode 100644 inference-router/tests/anthropic_buffered_guardrail.rs diff --git a/inference-router/src/routes/anthropic_messages.rs b/inference-router/src/routes/anthropic_messages.rs index cdf37c4df..0e939f24c 100644 --- a/inference-router/src/routes/anthropic_messages.rs +++ b/inference-router/src/routes/anthropic_messages.rs @@ -902,10 +902,11 @@ async fn forward_anthropic_passthrough( categories = ?v.categories, "guardrail pipeline blocked buffered response (anthropic route)" ); - return deny_response( + return deny_policy( StatusCode::FORBIDDEN, &v.message(), "content_policy_violation", + v.code(), ); } Err(e) => { @@ -918,10 +919,11 @@ async fn forward_anthropic_passthrough( error = %e, "guardrail pipeline unavailable (anthropic route) — failing closed" ); - return deny_response( - StatusCode::BAD_GATEWAY, + return deny_policy( + guardrail_error_status(&e), &e.to_string(), "api_error", + e.code(), ); } } 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" + ); +} From ab90ffba828b2e1f95ff278c3bc09440f6922525 Mon Sep 17 00:00:00 2001 From: johnseong Date: Thu, 27 Aug 2026 17:46:00 -0400 Subject: [PATCH 17/18] docs(inference): correct path-canonicalization wording (empty segments collapsed) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit decoded_path_segments now collapses benign empty segments (//, trailing /) rather than rejecting them, but five descriptions still said empty segments are rejected. Align all of them — the audit-log message and client 400 text (inference.rs), the security audit doc, the CRD reference, the CHANGELOG, and a test comment — to state that empty segments are canonicalized while dot (literal/percent-encoded), encoded slash/backslash, and malformed escapes are rejected. --- CHANGELOG.md | 6 ++++-- docs/api/crd-reference.md | 2 +- .../2026-08-25-multi-provider-guardrails.md | 6 ++++-- inference-router/src/routes/inference.rs | 8 ++++---- inference-router/tests/foundry_route_guard.rs | 3 ++- 5 files changed, 15 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 14c35b007..b62f8e810 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -67,8 +67,10 @@ never sees a provider key. 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 with dot/empty/percent-encoded-dot segments - are rejected (400) so a request cannot normalize into a different + 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** diff --git a/docs/api/crd-reference.md b/docs/api/crd-reference.md index ad15a9a7b..f2009f861 100644 --- a/docs/api/crd-reference.md +++ b/docs/api/crd-reference.md @@ -494,7 +494,7 @@ spec: | `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` / `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 containing dot segments, empty segments, or percent-encoded dot/slash forms 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. +> **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/security-audits/2026-08-25-multi-provider-guardrails.md b/docs/security-audits/2026-08-25-multi-provider-guardrails.md index 53a59b613..ca5a725a5 100644 --- a/docs/security-audits/2026-08-25-multi-provider-guardrails.md +++ b/docs/security-audits/2026-08-25-multi-provider-guardrails.md @@ -55,8 +55,10 @@ fail-closed/accounting fixes. 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 and rejects (400 `invalid_path`) any dot / empty / - encoded-slash segment before classification or forwarding, closing a traversal + 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 diff --git a/inference-router/src/routes/inference.rs b/inference-router/src/routes/inference.rs index ba4ef4200..d0979b725 100644 --- a/inference-router/src/routes/inference.rs +++ b/inference-router/src/routes/inference.rs @@ -804,8 +804,8 @@ async fn foundry_proxy( // URL parsing resolves `.`/`..` segments (including // percent-encoded `%2e` forms), a crafted path could therefore // reach a different upstream route than the one classified here. - // Reject every form that could normalize away before anything - // else runs. + // 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", @@ -813,11 +813,11 @@ async fn foundry_proxy( path = %uri.path(), decision = "deny", gate = "path_canonicalization", - "Foundry proxy path contains dot, empty, or encoded-slash segments" + "Foundry proxy path contains a dot segment, encoded slash/backslash, or malformed escape" ); return errors::openai_coded( StatusCode::BAD_REQUEST, - "path contains dot, empty, or encoded-slash segments", + "path contains a dot segment, encoded slash/backslash, or malformed escape", "invalid_path", "invalid_path", ) diff --git a/inference-router/tests/foundry_route_guard.rs b/inference-router/tests/foundry_route_guard.rs index 41705e653..e8e3eab84 100644 --- a/inference-router/tests/foundry_route_guard.rs +++ b/inference-router/tests/foundry_route_guard.rs @@ -295,7 +295,8 @@ async fn traversal_paths_are_rejected_not_forwarded() { "/openai/files/./responses", // Encoded slash smuggled inside one segment. "/openai/files/..%2fresponses", - // Empty segment (double slash) through a wildcard. + // Double slash + dot segment: rejected on the `..` (the `//` + // itself is collapsed, not rejected). "/openai/files//../responses", // Malformed percent escape. "/openai/files/%zz/responses", From f3200672f021b808975dc8239723a8975971dfe0 Mon Sep 17 00:00:00 2001 From: Pal Lakatos-Toth Date: Fri, 28 Aug 2026 10:20:50 +0200 Subject: [PATCH 18/18] docs(security): add independent guardrails audit sign-off Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: ee78c86b-0001-4d24-b829-7c75b59c9316 --- docs/security-audits/2026-08-25-multi-provider-guardrails.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/security-audits/2026-08-25-multi-provider-guardrails.md b/docs/security-audits/2026-08-25-multi-provider-guardrails.md index ca5a725a5..ebce1cd36 100644 --- a/docs/security-audits/2026-08-25-multi-provider-guardrails.md +++ b/docs/security-audits/2026-08-25-multi-provider-guardrails.md @@ -113,3 +113,4 @@ credential surface is router-side-only, bounded, and default-deny; the identifie bypass and functional gaps are fixed and regression-tested. Signed-off-by: John Seong +Signed-off-by: Pal Lakatos-Toth