diff --git a/CHANGELOG.md b/CHANGELOG.md index b389fe637..77d4f9455 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,91 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] — `crd-well-oiled-machine` +### Multi-cloud LLM providers + pluggable guardrails (slice 1: Anthropic, Ollama, OpenAI Moderation) + +First slice of the "multi-cloud LLM providers + native guardrails" +roadmap theme. Credentials stay router-side (secret mount / env on the +sidecar) — the agent process keeps talking to the localhost proxy and +never sees a provider key. + +**`InferencePolicy` CRD (controller)** + +- `spec.provider` — new optional typed enum (`InferenceProvider`: + `azure-openai` | `anthropic` | `ollama` | `bedrock`), serialized with + the same kebab-case tags `ModelRef.provider` strings have always + documented, so existing YAML values stay valid. `bedrock` is + schema-accepted for forward-compat; the router answers 501 until the + Bedrock client lands (declared intent is never silently rerouted). +- `spec.guardrails[]` — new optional ordered guardrail pipeline + (`{provider: openai-moderation, applyTo: input|output|both}`, 1–8 + stages via CEL). Both new fields are mutually exclusive with + `spec.bundleRef` (CEL + reconciler defense-in-depth) and flow through + `compile_to_profile` into the compiled `inference-policy.json` + (`"provider"` / `"guardrails"` keys, null when absent). +- Helm CRD template regenerated (`crd-inferencepolicy.yaml`). +- Reconciler forwards router-only provider config to every sidecar when + present on the controller env: `ANTHROPIC_API_KEY`, + `ANTHROPIC_ENDPOINT`, `OLLAMA_ENDPOINT`, `OPENAI_MODERATION_API_KEY` + (falls back to `OPENAI_API_KEY`), `OPENAI_MODERATION_ENDPOINT`. + Endpoints (not secrets) join `CONFIG_HASH_INPUTS`. + +**Inference router — multi-provider upstreams** + +- New `provider` module: tag parsing + fail-closed resolution. + `spec.provider` is the sole routing selector; the pre-existing + `modelPreference.primary.provider` tag stays informational (it drove + no routing before this slice), so an unchanged CR that only set a + model preference keeps its Azure upstream. Unknown tags warn and fall + back to Azure; missing endpoint/credential → 503; unimplemented + `bedrock` → 501. Never a silent Azure fallback for a selected + non-Azure provider. +- `UpstreamConfig` carries `provider` + router-held `api_key`; + `proxy.rs` gains per-provider URL shapes (Anthropic: path verbatim; + Ollama: OpenAI-compat under `/v1/`) and auth schemes (Anthropic: + `x-api-key` + default `anthropic-version`; Ollama: no credential). + Agent-supplied `x-api-key` headers are stripped as before. +- `provider: anthropic` serves the Anthropic Messages surface natively + (`/anthropic/v1/messages`, `/v1/messages`): streaming, tool use and + multi-modal content pass through. OpenAI-shaped `/v1/chat/completions` + under an Anthropic policy returns an explicit 501 pointing at the + Messages surface (mirrors the GitHub-Models 501 precedent). +- `provider: ollama` serves `/v1/chat/completions` (buffered + SSE) + against `OLLAMA_ENDPOINT` with token metering and budget tracking. +- Routing + guardrails are wired on `/v1/chat/completions` and the + Anthropic Messages routes only. 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 on the governed routes: request pre-flight (input + stages), buffered responses (including all Responses-API recovery + paths), and SSE streams via **hold-and-release** windows + (`GUARDRAIL_STREAM_SCAN_CHARS`, default 1000, clamped to the 16k scan + cap) — no assistant message text is delivered before a scan covers + it, including across partial-line chunk boundaries; flagged streams + are cut with a structured SSE error frame + `data: [DONE]`. Text over + 16k chars is scanned in successive windows, not truncated, so nothing + can be hidden past the cap. Tool-call arguments and provider + extended-"thinking" deltas are not yet scanned (roadmap). +- Fail-closed contract: declared-but-unbuildable stages reject the + request (503 `guardrail_misconfigured`); backend outages block (502 + `guardrail_unavailable`); unparseable bodies are scanned as raw text + rather than skipped. New `kars_guardrail_scans_total` metric + (provider / direction / outcome) and `x-kars-decision*` headers on + every block. + +Tests: compile/round-trip + enum wire-tag pins (controller), provider +resolution truth table, loader parsing, hold-and-release SSE guard +(clean / violation / split-event / scan-error), and wiremock +integration tests against fake Anthropic / Ollama / moderation +upstreams (`inference-router/tests/multi_provider_guardrails.rs`). + ### Upstream alignment — AGT mesh-identity hardening merged [microsoft/agent-governance-toolkit#2719](https://github.com/microsoft/agent-governance-toolkit/pull/2719) diff --git a/controller/src/config_hash.rs b/controller/src/config_hash.rs index 5a99ec205..b7ac69f1e 100644 --- a/controller/src/config_hash.rs +++ b/controller/src/config_hash.rs @@ -38,6 +38,11 @@ use std::sync::LazyLock; /// change and should be called out in the audit trail. pub const CONFIG_HASH_INPUTS: &[&str] = &[ "KARS_DISABLE_ENTRA_AUTH", + // Multi-provider endpoints (never secrets — matches the + // AZURE_OPENAI_API_KEY precedent). + "ANTHROPIC_ENDPOINT", + "OLLAMA_ENDPOINT", + "OPENAI_MODERATION_ENDPOINT", "AZURE_AUTHORITY_HOST", "AZURE_OPENAI_ENDPOINT", "AZURE_SUBSCRIPTION_ID", diff --git a/controller/src/crd_validations.rs b/controller/src/crd_validations.rs index bf121b5cd..7ffa74beb 100644 --- a/controller/src/crd_validations.rs +++ b/controller/src/crd_validations.rs @@ -277,8 +277,16 @@ pub fn inference_policy_validations() -> Vec { let severities = "['Safe','Low','Medium','High']"; vec![ ValidationRule { - rule: "!has(self.bundleRef) || (!has(self.tokenBudget) && !has(self.contentSafety) && !has(self.modelPreference) && !has(self.displayName))".into(), - message: Some("spec.bundleRef is mutually exclusive with spec.tokenBudget, spec.contentSafety, spec.modelPreference, and spec.displayName; the bundle carries those content fields".into()), + rule: "!has(self.bundleRef) || (!has(self.tokenBudget) && !has(self.contentSafety) && !has(self.modelPreference) && !has(self.provider) && !has(self.guardrails) && !has(self.displayName))".into(), + message: Some("spec.bundleRef is mutually exclusive with spec.tokenBudget, spec.contentSafety, spec.modelPreference, spec.provider, spec.guardrails, and spec.displayName; the bundle carries those content fields".into()), + reason: Some("FieldValueInvalid".into()), + ..ValidationRule::default() + }, + // An empty list is an authoring mistake (omit the field for + // no pipeline); the cap bounds per-request scan fan-out. + ValidationRule { + rule: "!has(self.guardrails) || (size(self.guardrails) >= 1 && size(self.guardrails) <= 8)".into(), + message: Some("spec.guardrails, when set, must contain 1-8 stages (omit the field for no pipeline)".into()), reason: Some("FieldValueInvalid".into()), ..ValidationRule::default() }, diff --git a/controller/src/inference_policy.rs b/controller/src/inference_policy.rs index cc0d35910..67a3ae866 100644 --- a/controller/src/inference_policy.rs +++ b/controller/src/inference_policy.rs @@ -54,6 +54,44 @@ use serde::{Deserialize, Serialize}; use crate::mcp_server::LocalObjectRef; +/// Inference provider selector. Explicit kebab-case renames (not +/// `rename_all`) keep the wire tags byte-identical to the free-form +/// `ModelRef.provider` strings this CRD has always documented, so a +/// CR can move to the typed field without a migration. `bedrock` is +/// schema-accepted for forward-compat; the router returns 501 until +/// its client lands. +#[derive(Serialize, Deserialize, Clone, Copy, Debug, Default, PartialEq, Eq, JsonSchema)] +pub enum InferenceProvider { + /// Azure OpenAI / Azure AI Foundry (default — Phase 1 substrate). + #[default] + #[serde(rename = "azure-openai")] + AzureOpenAI, + /// Anthropic Messages API (`api.anthropic.com` or a compatible + /// gateway). + #[serde(rename = "anthropic")] + Anthropic, + /// OpenAI-compatible Ollama server (in-cluster or local). + #[serde(rename = "ollama")] + Ollama, + /// AWS Bedrock — schema-level forward-compat; router support is a + /// follow-up slice. + #[serde(rename = "bedrock")] + AWSBedrock, +} + +impl InferenceProvider { + /// The kebab-case wire tag serde emits. + #[must_use] + pub fn as_tag(&self) -> &'static str { + match self { + Self::AzureOpenAI => "azure-openai", + Self::Anthropic => "anthropic", + Self::Ollama => "ollama", + Self::AWSBedrock => "bedrock", + } + } +} + /// `InferencePolicy.spec` — declares per-sandbox inference-time /// guardrails: token budgets, Content Safety severity floors, model /// preference + fallback chain. @@ -102,6 +140,20 @@ pub struct InferencePolicySpec { /// [`Self::bundle_ref`]. pub model_preference: Option, + /// Inference provider for the call sites this policy governs, and + /// the sole routing selector. Absent ⇒ the env-configured Azure + /// OpenAI / Foundry upstream. A non-Azure provider swaps the base + /// URL and auth scheme; credentials stay in the router sidecar, + /// never in the agent. Mutually exclusive with [`Self::bundle_ref`]. + pub provider: Option, + + /// Ordered guardrail pipeline stages the router runs around each + /// call (input pre-flight + response, buffered and streaming), in + /// declaration order — first flag blocks. Absent ⇒ only the + /// Phase 1 substrate (Foundry annotations + `contentSafety`). + /// Mutually exclusive with [`Self::bundle_ref`]. + pub guardrails: Option>, + /// Optional human-readable label. Mutually exclusive with /// [`Self::bundle_ref`] — when `bundleRef` is set, the label /// comes from the signed bundle. @@ -228,6 +280,45 @@ pub struct ModelRef { pub deployment: String, } +/// A single router-side guardrail stage. A stage whose backend isn't +/// configured (e.g. missing moderation key) fails the request closed, +/// never skips. +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, JsonSchema)] +#[serde(rename_all = "camelCase")] +pub struct GuardrailStage { + /// Guardrail backend. + pub provider: GuardrailProvider, + + /// Which direction(s) this stage scans. Absent ⇒ `both`. + pub apply_to: Option, +} + +/// Guardrail backend. Bedrock Guardrails / Model Armor are roadmap +/// follow-ups that extend this enum. +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, JsonSchema)] +pub enum GuardrailProvider { + /// OpenAI Moderation API; router-side key via + /// `OPENAI_MODERATION_API_KEY` (falls back to `OPENAI_API_KEY`). + #[serde(rename = "openai-moderation")] + OpenAIModeration, +} + +/// Scan direction for a [`GuardrailStage`]. +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, Default, JsonSchema)] +pub enum GuardrailApplyTo { + /// Scan only the request (prompt) text. + #[serde(rename = "input")] + Input, + /// Scan only the response (completion) text — buffered and + /// streaming. + #[serde(rename = "output")] + Output, + /// Scan both directions (default). + #[default] + #[serde(rename = "both")] + Both, +} + /// Status of an `InferencePolicy` reconcile. #[derive(Debug, Serialize, Deserialize, Default, Clone, JsonSchema)] #[serde(rename_all = "camelCase")] diff --git a/controller/src/inference_policy_compile.rs b/controller/src/inference_policy_compile.rs index 548f1def1..ae6c630c0 100644 --- a/controller/src/inference_policy_compile.rs +++ b/controller/src/inference_policy_compile.rs @@ -72,6 +72,8 @@ use crate::inference_policy::InferencePolicySpec; /// "tokenBudget": { "perRequestTokens": ..., "dailyTokens": ..., "monthlyTokens": ... } | null, /// "contentSafety": { "hate": ..., "selfHarm": ..., "sexual": ..., "violence": ..., "requirePromptShields": ... } | null, /// "modelPreference": { "primary": {provider, deployment}, "fallback": [...] } | null, +/// "provider": "azure-openai" | "anthropic" | "ollama" | "bedrock" | null, +/// "guardrails": [ { "provider": "openai-moderation", "applyTo": "input"|"output"|"both"|null }, ... ] | null, /// "displayName": "..." | null /// } /// ``` @@ -118,11 +120,29 @@ pub fn compile_to_profile(spec: &InferencePolicySpec) -> Value { }) }); + // Provider + guardrails emit the same kebab-case wire tags the + // CRD serde uses, so the router parses one vocabulary. + let provider = spec.provider.as_ref().map(|p| json!(p.as_tag())); + + let guardrails = spec.guardrails.as_ref().map(|stages| { + json!( + stages + .iter() + .map(|g| json!({ + "provider": g.provider, + "applyTo": g.apply_to, + })) + .collect::>() + ) + }); + json!({ "appliesTo": applies_to, "tokenBudget": token_budget, "contentSafety": content_safety, "modelPreference": model_preference, + "provider": provider, + "guardrails": guardrails, "displayName": spec.display_name, }) } @@ -195,7 +215,8 @@ pub fn inference_policy_digest(body: &[u8]) -> String { mod tests { use super::*; use crate::inference_policy::{ - ContentSafetyFloor, InferenceAppliesTo, InferencePolicySpec, ModelPreference, ModelRef, + ContentSafetyFloor, GuardrailApplyTo, GuardrailProvider, GuardrailStage, + InferenceAppliesTo, InferencePolicySpec, InferenceProvider, ModelPreference, ModelRef, TokenBudget, }; @@ -230,6 +251,17 @@ mod tests { deployment: "claude-3-5-sonnet".into(), }], }), + provider: Some(InferenceProvider::Anthropic), + guardrails: Some(vec![ + GuardrailStage { + provider: GuardrailProvider::OpenAIModeration, + apply_to: Some(GuardrailApplyTo::Output), + }, + GuardrailStage { + provider: GuardrailProvider::OpenAIModeration, + apply_to: None, + }, + ]), display_name: Some("Prod chat policy".into()), bundle_ref: None, } @@ -243,6 +275,8 @@ mod tests { assert!(profile.get("tokenBudget").unwrap().is_null()); assert!(profile.get("contentSafety").unwrap().is_null()); assert!(profile.get("modelPreference").unwrap().is_null()); + assert!(profile.get("provider").unwrap().is_null()); + assert!(profile.get("guardrails").unwrap().is_null()); assert!(profile.get("appliesTo").unwrap().is_object()); } @@ -271,9 +305,60 @@ mod tests { profile["modelPreference"]["fallback"][0]["provider"], "anthropic" ); + assert_eq!(profile["provider"], "anthropic"); + let stages = profile["guardrails"].as_array().unwrap(); + assert_eq!(stages.len(), 2); + assert_eq!(stages[0]["provider"], "openai-moderation"); + assert_eq!(stages[0]["applyTo"], "output"); + assert_eq!(stages[1]["provider"], "openai-moderation"); + assert!(stages[1]["applyTo"].is_null()); assert_eq!(profile["displayName"], "Prod chat policy"); } + #[test] + fn provider_enum_serializes_to_kebab_case_wire_tags() { + // Wire-contract pin: the typed `spec.provider` must emit the + // same kebab-case tags the free-form `ModelRef.provider` + // strings have always documented, so existing YAML values + // stay valid when operators migrate to the typed field. + for (variant, tag) in [ + (InferenceProvider::AzureOpenAI, "azure-openai"), + (InferenceProvider::Anthropic, "anthropic"), + (InferenceProvider::Ollama, "ollama"), + (InferenceProvider::AWSBedrock, "bedrock"), + ] { + assert_eq!(serde_json::to_value(variant).unwrap(), tag); + assert_eq!(variant.as_tag(), tag); + let parsed: InferenceProvider = serde_json::from_value(serde_json::json!(tag)).unwrap(); + assert_eq!(parsed, variant); + } + } + + #[test] + fn version_hash_changes_when_provider_changes() { + let mut a = full_spec(); + let b = full_spec(); + a.provider = Some(InferenceProvider::Ollama); + assert_ne!( + version_hash(&compile_to_profile(&a)), + version_hash(&compile_to_profile(&b)) + ); + } + + #[test] + fn version_hash_changes_when_guardrails_change() { + let mut a = full_spec(); + let b = full_spec(); + a.guardrails = Some(vec![GuardrailStage { + provider: GuardrailProvider::OpenAIModeration, + apply_to: Some(GuardrailApplyTo::Input), + }]); + assert_ne!( + version_hash(&compile_to_profile(&a)), + version_hash(&compile_to_profile(&b)) + ); + } + #[test] fn compile_is_deterministic() { let spec = full_spec(); diff --git a/controller/src/inference_policy_reconciler.rs b/controller/src/inference_policy_reconciler.rs index 30e64feda..d7a9eab67 100644 --- a/controller/src/inference_policy_reconciler.rs +++ b/controller/src/inference_policy_reconciler.rs @@ -375,6 +375,8 @@ async fn resolve_inference_source( let inline_any = spec.token_budget.is_some() || spec.content_safety.is_some() || spec.model_preference.is_some() + || spec.provider.is_some() + || spec.guardrails.is_some() || spec.display_name.is_some(); let bundle_set = spec.bundle_ref.is_some(); @@ -389,7 +391,8 @@ async fn resolve_inference_source( Some(( "InvalidSpec", "spec.bundleRef is mutually exclusive with spec.tokenBudget, \ - spec.contentSafety, spec.modelPreference, and spec.displayName" + spec.contentSafety, spec.modelPreference, spec.provider, \ + spec.guardrails, and spec.displayName" .into(), )), ); @@ -520,6 +523,11 @@ fn merge_bundle_with_selector( token_budget, content_safety, model_preference, + // The signed-bundle canonical format doesn't carry + // provider/guardrails yet; bundle-sourced policies keep the + // router defaults on these axes. + provider: None, + guardrails: None, display_name: verified.display_name.clone(), bundle_ref: None, } diff --git a/controller/src/reconciler/mod.rs b/controller/src/reconciler/mod.rs index ce919983d..5470b10fd 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,22 @@ async fn reconcile(sandbox: Arc, ctx: Arc) -> Result Result<()> { dev_openai_api_key, dev_provider, dev_copilot_github_token, + anthropic_api_key: std::env::var("ANTHROPIC_API_KEY").unwrap_or_default(), + anthropic_endpoint: std::env::var("ANTHROPIC_ENDPOINT").unwrap_or_default(), + ollama_endpoint: std::env::var("OLLAMA_ENDPOINT").unwrap_or_default(), + openai_moderation_api_key: std::env::var("OPENAI_MODERATION_API_KEY") + .or_else(|_| std::env::var("OPENAI_API_KEY")) + .unwrap_or_default(), + openai_moderation_endpoint: std::env::var("OPENAI_MODERATION_ENDPOINT").unwrap_or_default(), dev_profile, cluster_name: std::env::var("CLUSTER_NAME") .ok() diff --git a/deploy/helm/kars/templates/crd-inferencepolicy.yaml b/deploy/helm/kars/templates/crd-inferencepolicy.yaml index fceb79c29..202ffc618 100644 --- a/deploy/helm/kars/templates/crd-inferencepolicy.yaml +++ b/deploy/helm/kars/templates/crd-inferencepolicy.yaml @@ -181,6 +181,37 @@ spec: comes from the signed bundle. nullable: true type: string + guardrails: + description: |- + Ordered guardrail pipeline stages the router runs around each + call (input pre-flight + response, buffered and streaming), in + declaration order — first flag blocks. Absent ⇒ only the + Phase 1 substrate (Foundry annotations + `contentSafety`). + Mutually exclusive with [`Self::bundle_ref`]. + items: + description: |- + A single router-side guardrail stage. A stage whose backend isn't + configured (e.g. missing moderation key) fails the request closed, + never skips. + properties: + applyTo: + description: Scan direction for a [`GuardrailStage`]. + enum: + - input + - output + - both + nullable: true + type: string + provider: + description: Guardrail backend. + enum: + - openai-moderation + type: string + required: + - provider + type: object + nullable: true + type: array modelPreference: description: |- Model preference + fallback chain. Optional. **Not** a router: @@ -226,6 +257,21 @@ spec: required: - primary type: object + provider: + description: |- + Inference provider selector. Explicit kebab-case renames (not + `rename_all`) keep the wire tags byte-identical to the free-form + `ModelRef.provider` strings this CRD has always documented, so a + CR can move to the typed field without a migration. `bedrock` is + schema-accepted for forward-compat; the router returns 501 until + its client lands. + enum: + - azure-openai + - anthropic + - ollama + - bedrock + nullable: true + type: string tokenBudget: description: |- Token-budget caps. Optional — absent ⇒ no budget enforcement. @@ -262,9 +308,12 @@ spec: - appliesTo type: object x-kubernetes-validations: - - message: spec.bundleRef is mutually exclusive with spec.tokenBudget, spec.contentSafety, spec.modelPreference, and spec.displayName; the bundle carries those content fields + - message: spec.bundleRef is mutually exclusive with spec.tokenBudget, spec.contentSafety, spec.modelPreference, spec.provider, spec.guardrails, and spec.displayName; the bundle carries those content fields + reason: FieldValueInvalid + rule: '!has(self.bundleRef) || (!has(self.tokenBudget) && !has(self.contentSafety) && !has(self.modelPreference) && !has(self.provider) && !has(self.guardrails) && !has(self.displayName))' + - message: spec.guardrails, when set, must contain 1-8 stages (omit the field for no pipeline) reason: FieldValueInvalid - rule: '!has(self.bundleRef) || (!has(self.tokenBudget) && !has(self.contentSafety) && !has(self.modelPreference) && !has(self.displayName))' + rule: '!has(self.guardrails) || (size(self.guardrails) >= 1 && size(self.guardrails) <= 8)' - message: spec.tokenBudget.monthlyTokens must be >= spec.tokenBudget.dailyTokens reason: FieldValueInvalid rule: '!has(self.tokenBudget) || !has(self.tokenBudget.monthlyTokens) || !has(self.tokenBudget.dailyTokens) || self.tokenBudget.monthlyTokens >= self.tokenBudget.dailyTokens' @@ -403,4 +452,3 @@ spec: storage: true subresources: status: {} - diff --git a/docs/api/crd-reference.md b/docs/api/crd-reference.md index 3f7c49029..a2b77ca1d 100644 --- a/docs/api/crd-reference.md +++ b/docs/api/crd-reference.md @@ -458,6 +458,10 @@ spec: sandboxName: my-agent # exact match; empty = any in ns sandboxMatchLabels: {} # AND with sandboxName action: "*" # chat | responses | image | embeddings | * + provider: azure-openai # optional: azure-openai | anthropic | ollama | bedrock + guardrails: # optional: ordered pipeline, first flag blocks + - provider: openai-moderation + applyTo: both # input | output | both (default both) modelPreference: primary: provider: azure-openai # azure-openai | anthropic | gemini | bedrock | ollama @@ -480,13 +484,17 @@ spec: | Field | Notes | |---|---| | `spec.appliesTo` | Required selector — AND of `sandboxName`, `sandboxMatchLabels`, `action`. | -| `spec.modelPreference.primary` | `{provider, deployment}`. `provider` is one of `azure-openai`, `anthropic`, `gemini`, `bedrock`, `ollama`. | -| `spec.modelPreference.fallback[]` | Ordered fallback routes — first healthy wins, deterministically. No load-balancing. | +| `spec.provider` | Optional typed default provider (`azure-openai` \| `anthropic` \| `ollama` \| `bedrock`), and the **only** field that drives provider routing. Absent ⇒ the env-configured Azure OpenAI / Foundry upstream. `anthropic` serves the Anthropic Messages surface (`/anthropic/v1/messages`, streaming + tools pass-through) with the router-held `ANTHROPIC_API_KEY`; `ollama` serves OpenAI-compatible chat completions against `OLLAMA_ENDPOINT` (no credential). `bedrock` is schema-accepted but the router returns 501 until the Bedrock client lands. Routing applies to `/v1/chat/completions` and the Anthropic Messages routes only; the other inference routes (see note below) refuse a non-Azure provider with 501. Credentials/endpoints are router-sidecar config only — never visible to the agent. | +| `spec.guardrails[]` | Optional ordered guardrail pipeline (1–8 stages) the router runs on `/v1/chat/completions` and the Anthropic Messages routes. Stage: `{provider, applyTo}` with `provider: openai-moderation` (Bedrock Guardrails / Model Armor are roadmap follow-ups) and `applyTo: input \| output \| both`. Fail-closed: a declared stage that cannot run (missing key, backend outage) blocks the request. Streaming responses use hold-and-release windows — no assistant message text reaches the client before a scan covers it (tool-call arguments and provider "thinking" deltas are not yet scanned — see note). Text over 16k chars is scanned in successive windows, not truncated. | +| `spec.modelPreference.primary` | `{provider, deployment}`. `provider` is one of `azure-openai`, `anthropic`, `gemini`, `bedrock`, `ollama`. This tag stays **informational** — it does not select the upstream (use `spec.provider`); `modelPreference` drives deployment failover within the resolved provider. | +| `spec.modelPreference.fallback[]` | Ordered fallback routes — first healthy wins, deterministically. No load-balancing. Failover walks deployments on the resolved provider (cross-provider failover is a follow-up). | | `spec.tokenBudget.perRequestTokens` | Per-call hard cap. Inference calls exceeding this are refused **before** the upstream forward. | | `spec.tokenBudget.dailyTokens` / `monthlyTokens` | Accepted and surfaced in status; **aggregate enforcement is not yet wired** — see roadmap below. CEL enforces `monthlyTokens ≥ dailyTokens`. | | `spec.contentSafety` | Per-category severity floors (`Safe` \| `Low` \| `Medium` \| `High`). The router parses Foundry `prompt_filter_results` inline; there is **no** separate Content Safety call. | | `spec.contentSafety.requirePromptShields` | Fail-closed if Prompt Shields are advertised by the deployment but the response lacks the corresponding annotations. | -| `spec.bundleRef` | Signed OCI artifact alternative to inline `tokenBudget` / `contentSafety` / `modelPreference` / `displayName`. `appliesTo` always comes from the CR. | +| `spec.bundleRef` | Signed OCI artifact alternative to inline `tokenBudget` / `contentSafety` / `modelPreference` / `provider` / `guardrails` / `displayName`. `appliesTo` always comes from the CR. | + +> **Provider + guardrail enforcement scope today.** `spec.provider` routing and `spec.guardrails[]` scanning are wired on `/v1/chat/completions` and the Anthropic Messages routes (`/anthropic/v1/messages`, `/v1/messages`). 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/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..3438a106b 100644 --- a/inference-router/src/config.rs +++ b/inference-router/src/config.rs @@ -81,6 +81,49 @@ pub struct Config { /// Captured at config-load time so provider detection is a pure /// function on the `Config` struct (testable without env hacks). pub provider_override: Option, + + /// Anthropic Messages API endpoint used when an `InferencePolicy` + /// selects `provider: anthropic`. Default `https://api.anthropic.com`; + /// override with `ANTHROPIC_ENDPOINT` for gateways. + pub anthropic_endpoint: String, + + /// Anthropic API key — `ANTHROPIC_API_KEY` env or the + /// `anthropic-api-key` secret mount. Router-side only. `None` ⇒ + /// Anthropic policies fail closed. + pub anthropic_api_key: Option, + + /// OpenAI-compatible Ollama endpoint (e.g. + /// `http://ollama.ollama.svc:11434`) used when an + /// `InferencePolicy` selects `provider: ollama`. No default — the + /// operator must opt in via `OLLAMA_ENDPOINT`. + pub ollama_endpoint: Option, + + /// Base endpoint for the OpenAI Moderation guardrail stage. + /// Default `https://api.openai.com`; override with + /// `OPENAI_MODERATION_ENDPOINT`. + pub openai_moderation_endpoint: String, + + /// OpenAI Moderation key — `OPENAI_MODERATION_API_KEY` env (falls + /// back to `OPENAI_API_KEY`, then the `openai-moderation-api-key` + /// secret mount). `None` ⇒ `openai-moderation` stages fail closed. + pub openai_moderation_api_key: Option, + + /// Moderation model (`OPENAI_MODERATION_MODEL`, default + /// `omni-moderation-latest`). + pub openai_moderation_model: String, +} + +/// Read a credential from an env var, falling back to the standard +/// kars secret mounts (`/etc/kars/secrets/` then +/// `/run/secrets/`). Mirrors the admin-token lookup in +/// `routes::AppState::new`. Empty values are treated as unset. +fn secret_from_env_or_mount(env: &str, file: &str) -> Option { + std::env::var(env) + .ok() + .or_else(|| std::fs::read_to_string(format!("/etc/kars/secrets/{file}")).ok()) + .or_else(|| std::fs::read_to_string(format!("/run/secrets/{file}")).ok()) + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) } impl Config { @@ -145,6 +188,38 @@ impl Config { .ok() .filter(|s| !s.is_empty()) .map(|s| s.to_ascii_lowercase()), + + anthropic_endpoint: std::env::var("ANTHROPIC_ENDPOINT") + .ok() + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| "https://api.anthropic.com".into()), + + anthropic_api_key: secret_from_env_or_mount("ANTHROPIC_API_KEY", "anthropic-api-key"), + + ollama_endpoint: std::env::var("OLLAMA_ENDPOINT") + .ok() + .filter(|s| !s.is_empty()), + + openai_moderation_endpoint: std::env::var("OPENAI_MODERATION_ENDPOINT") + .ok() + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| "https://api.openai.com".into()), + + openai_moderation_api_key: secret_from_env_or_mount( + "OPENAI_MODERATION_API_KEY", + "openai-moderation-api-key", + ) + .or_else(|| { + std::env::var("OPENAI_API_KEY") + .ok() + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + }), + + openai_moderation_model: std::env::var("OPENAI_MODERATION_MODEL") + .ok() + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| "omni-moderation-latest".into()), }) } @@ -212,6 +287,12 @@ mod tests { registry_mode: RegistryMode::Local, registry_url: None, provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: None, + ollama_endpoint: None, + openai_moderation_endpoint: "https://api.openai.com".into(), + openai_moderation_api_key: None, + openai_moderation_model: "omni-moderation-latest".into(), } } diff --git a/inference-router/src/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..1df0e9a47 --- /dev/null +++ b/inference-router/src/guardrails.rs @@ -0,0 +1,1490 @@ +// 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/inference_policy_loader.rs b/inference-router/src/inference_policy_loader.rs index 30defe06c..c04b27398 100644 --- a/inference-router/src/inference_policy_loader.rs +++ b/inference-router/src/inference_policy_loader.rs @@ -165,6 +165,14 @@ pub struct LoadedInferencePolicy { /// back to the env-driven default deployment (back-compat). pub model_preference: Option, + /// `spec.provider` — raw kebab-case tag, resolved per request by + /// [`crate::provider::resolve`]. `None` ⇒ Azure (back-compat). + pub provider: Option, + + /// `spec.guardrails[]` — pipeline stages; validity is checked at + /// build time in [`crate::guardrails::GuardrailPipeline::from_stages`]. + pub guardrails: Vec, + /// Whole profile JSON, kept so subsequent sub-slices can pick up /// other axes without a new loader. pub raw: serde_json::Value, @@ -327,6 +335,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 +363,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 +406,8 @@ pub fn load_inference_policy_from_dir( monthly_tokens, content_safety, model_preference, + provider, + guardrails, raw: parsed, }) } @@ -507,6 +532,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 +561,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 +777,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 +800,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 99024ea1a..6890959e7 100644 --- a/inference-router/src/lib.rs +++ b/inference-router/src/lib.rs @@ -32,6 +32,7 @@ pub mod errors; pub mod failover; pub mod forward_proxy; pub mod governance; +pub mod guardrails; pub mod handoff; pub mod inference_policy_loader; pub mod mcp; @@ -40,6 +41,7 @@ pub mod mesh; pub mod metrics; pub mod policy_envelope; pub mod policy_status; +pub mod provider; pub mod providers; pub mod proxy; pub mod rate_limiter; diff --git a/inference-router/src/metrics.rs b/inference-router/src/metrics.rs index 3ce7895c2..faeb00f53 100644 --- a/inference-router/src/metrics.rs +++ b/inference-router/src/metrics.rs @@ -38,6 +38,19 @@ pub static TOKENS_USED: LazyLock = LazyLock::new(|| { .unwrap() }); +/// Guardrail pipeline scans by backend, direction, and outcome +/// (`pass` | `flagged` | `error`). +pub static GUARDRAIL_SCANS: LazyLock = LazyLock::new(|| { + register_int_counter_vec!( + opts!( + "kars_guardrail_scans_total", + "Guardrail pipeline scans by backend, direction, and outcome" + ), + &["provider", "direction", "outcome"] + ) + .unwrap() +}); + // ── AGT Governance metrics ────────────────────────────────────────────────── /// Total AGT policy evaluations by decision (allow, deny, requires_approval, rate_limited). diff --git a/inference-router/src/provider.rs b/inference-router/src/provider.rs new file mode 100644 index 000000000..06ea713df --- /dev/null +++ b/inference-router/src/provider.rs @@ -0,0 +1,300 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Multi-provider upstream resolution. +//! +//! Maps `InferencePolicy.spec.provider` onto a concrete upstream +//! target (base URL + auth scheme). Credentials/endpoints come only +//! from the router's env / secret mounts — never from the agent. +//! +//! `spec.provider` is the sole routing selector; the pre-existing +//! `modelPreference.primary.provider` tag stays informational, so +//! adding routing doesn't reroute CRs that only set a model +//! preference. Absent/empty or unrecognised ⇒ Azure. `bedrock` +//! (unimplemented) and a provider with missing config both fail the +//! request closed rather than silently falling back to Azure. + +use crate::config::Config; + +/// Providers the router can actually forward to today. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum ProviderKind { + /// Azure OpenAI / Foundry (also GitHub Models + Copilot via + /// endpoint detection) — the Phase 1 substrate. Default. + #[default] + AzureOpenAI, + /// Anthropic Messages API — native pass-through on + /// `/v1/messages`; auth via `x-api-key` from the router-side + /// secret. + Anthropic, + /// OpenAI-compatible Ollama server — pass-through on + /// `/v1/chat/completions`; no auth. + Ollama, +} + +impl ProviderKind { + /// Kebab-case wire tag, matching the controller-side + /// `InferenceProvider::as_tag`. + #[must_use] + pub fn as_tag(&self) -> &'static str { + match self { + Self::AzureOpenAI => "azure-openai", + Self::Anthropic => "anthropic", + Self::Ollama => "ollama", + } + } +} + +/// Why a provider tag could not be turned into a forwardable upstream. +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum ProviderError { + /// Tag is recognised by the CRD schema but the router has no + /// client for it yet (`bedrock`). HTTP mapping: 501. + #[error("provider '{tag}' is not implemented by this router build")] + Unimplemented { tag: String }, + /// Provider needs an endpoint the router was not configured with. + /// HTTP mapping: 503 (operator config gap, not a caller bug). + #[error( + "provider '{provider}' selected by InferencePolicy but {env} is not configured on the router" + )] + MissingEndpoint { + provider: &'static str, + env: &'static str, + }, + /// Provider needs a credential the router was not configured + /// with. HTTP mapping: 503. + #[error( + "provider '{provider}' selected by InferencePolicy but no credential is configured ({env} or secret mount)" + )] + MissingCredential { + provider: &'static str, + env: &'static str, + }, +} + +/// Parse a policy provider tag. `Ok(None)` means "no opinion" (empty +/// or unknown tag — logged by the caller, keeps the pre-slice +/// informational-only behaviour for tags like `gemini`). +/// `Err(Unimplemented)` is reserved for tags the CRD schema accepts +/// but the router cannot serve, so declared intent fails loudly. +pub fn parse_tag(tag: &str) -> Result, ProviderError> { + match tag.trim().to_ascii_lowercase().as_str() { + "azure-openai" => Ok(Some(ProviderKind::AzureOpenAI)), + "anthropic" => Ok(Some(ProviderKind::Anthropic)), + "ollama" => Ok(Some(ProviderKind::Ollama)), + "bedrock" => Err(ProviderError::Unimplemented { + tag: "bedrock".into(), + }), + _ => Ok(None), + } +} + +/// Concrete upstream target after resolution. For non-Azure providers +/// this carries the endpoint (and credential) the proxy layer needs; +/// `AzureOpenAI` keeps the env-driven endpoint already present on +/// `UpstreamConfig`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ProviderTarget { + AzureOpenAI, + Anthropic { endpoint: String, api_key: String }, + Ollama { endpoint: String }, +} + +impl ProviderTarget { + #[must_use] + pub fn kind(&self) -> ProviderKind { + match self { + Self::AzureOpenAI => ProviderKind::AzureOpenAI, + Self::Anthropic { .. } => ProviderKind::Anthropic, + Self::Ollama { .. } => ProviderKind::Ollama, + } + } +} + +/// Resolve the upstream target from `spec.provider` (the sole routing +/// selector; see module docs). +pub fn resolve(policy_tag: Option<&str>, config: &Config) -> Result { + let kind = effective_kind(policy_tag)?; + target_for(kind, config) +} + +fn effective_kind(policy_tag: Option<&str>) -> Result { + if let Some(tag) = policy_tag.filter(|t| !t.trim().is_empty()) { + match parse_tag(tag)? { + Some(kind) => return Ok(kind), + None => { + tracing::warn!( + tag, + "InferencePolicy spec.provider tag not recognised — using azure-openai" + ); + } + } + } + Ok(ProviderKind::AzureOpenAI) +} + +fn target_for(kind: ProviderKind, config: &Config) -> Result { + match kind { + ProviderKind::AzureOpenAI => Ok(ProviderTarget::AzureOpenAI), + ProviderKind::Anthropic => { + let api_key = + config + .anthropic_api_key + .clone() + .ok_or(ProviderError::MissingCredential { + provider: "anthropic", + env: "ANTHROPIC_API_KEY", + })?; + Ok(ProviderTarget::Anthropic { + endpoint: config.anthropic_endpoint.clone(), + api_key, + }) + } + ProviderKind::Ollama => { + let endpoint = + config + .ollama_endpoint + .clone() + .ok_or(ProviderError::MissingEndpoint { + provider: "ollama", + env: "OLLAMA_ENDPOINT", + })?; + Ok(ProviderTarget::Ollama { endpoint }) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::config::{Config, RegistryMode}; + + fn cfg(anthropic_key: Option<&str>, ollama: Option<&str>) -> Config { + Config { + port: 8443, + foundry_endpoint: None, + foundry_project_endpoint: None, + azure_openai_endpoint: Some("https://contoso.openai.azure.com".into()), + default_model: "gpt-4o-mini".into(), + content_safety_enabled: false, + prompt_shields_enabled: false, + content_safety_endpoint: None, + token_budget_daily: 0, + token_budget_per_request: 0, + registry_mode: RegistryMode::Local, + registry_url: None, + provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: anthropic_key.map(String::from), + ollama_endpoint: ollama.map(String::from), + openai_moderation_endpoint: "https://api.openai.com".into(), + openai_moderation_api_key: None, + openai_moderation_model: "omni-moderation-latest".into(), + } + } + + #[test] + fn parse_recognises_supported_tags_case_insensitively() { + assert_eq!( + parse_tag("azure-openai").unwrap(), + Some(ProviderKind::AzureOpenAI) + ); + assert_eq!( + parse_tag("Anthropic").unwrap(), + Some(ProviderKind::Anthropic) + ); + assert_eq!(parse_tag(" ollama ").unwrap(), Some(ProviderKind::Ollama)); + } + + #[test] + fn parse_returns_none_for_unknown_tags() { + assert_eq!(parse_tag("gemini").unwrap(), None); + assert_eq!(parse_tag("").unwrap(), None); + assert_eq!(parse_tag("Foundry").unwrap(), None); + } + + #[test] + fn parse_rejects_bedrock_as_unimplemented() { + assert!(matches!( + parse_tag("bedrock"), + Err(ProviderError::Unimplemented { .. }) + )); + } + + #[test] + fn absent_provider_resolves_to_azure() { + assert_eq!( + resolve(None, &cfg(None, None)).unwrap(), + ProviderTarget::AzureOpenAI + ); + assert_eq!( + resolve(Some(""), &cfg(None, None)).unwrap(), + ProviderTarget::AzureOpenAI + ); + } + + #[test] + fn unknown_provider_tag_falls_back_to_azure() { + // A documented-but-unrouted tag like gemini must not error — + // it stays informational, request goes to Azure. + assert_eq!( + resolve(Some("gemini"), &cfg(Some("sk-x"), None)).unwrap(), + ProviderTarget::AzureOpenAI + ); + } + + #[test] + fn anthropic_provider_resolves_with_key() { + assert_eq!( + resolve(Some("anthropic"), &cfg(Some("sk-x"), None)).unwrap(), + ProviderTarget::Anthropic { + endpoint: "https://api.anthropic.com".into(), + api_key: "sk-x".into(), + } + ); + } + + #[test] + fn anthropic_without_key_fails_closed() { + assert!(matches!( + resolve(Some("anthropic"), &cfg(None, None)), + Err(ProviderError::MissingCredential { + provider: "anthropic", + .. + }) + )); + } + + #[test] + fn ollama_without_endpoint_fails_closed() { + assert!(matches!( + resolve(Some("ollama"), &cfg(None, None)), + Err(ProviderError::MissingEndpoint { + provider: "ollama", + .. + }) + )); + } + + #[test] + fn ollama_with_endpoint_resolves() { + assert_eq!( + resolve( + Some("ollama"), + &cfg(None, Some("http://ollama.ollama.svc:11434")) + ) + .unwrap(), + ProviderTarget::Ollama { + endpoint: "http://ollama.ollama.svc:11434".into() + } + ); + } + + #[test] + fn bedrock_is_unimplemented_not_silent() { + assert!(matches!( + resolve(Some("bedrock"), &cfg(None, None)), + Err(ProviderError::Unimplemented { .. }) + )); + } +} diff --git a/inference-router/src/proxy.rs b/inference-router/src/proxy.rs index 87da56894..5d79b145e 100644 --- a/inference-router/src/proxy.rs +++ b/inference-router/src/proxy.rs @@ -18,6 +18,7 @@ use crate::copilot_auth::{ self, COPILOT_INTEGRATION_ID, CopilotTokenCache, EDITOR_PLUGIN_VERSION, EDITOR_VERSION, }; use crate::metrics; +use crate::provider::ProviderKind; use std::sync::Arc; /// Upstream configuration for a single request. @@ -26,6 +27,39 @@ pub struct UpstreamConfig { pub endpoint: String, pub deployment: String, pub sandbox_name: String, + /// Which provider family `endpoint` belongs to — drives URL shape + /// and auth scheme. `AzureOpenAI` preserves the historic + /// behaviour (incl. GitHub Models / Copilot endpoint detection). + pub provider: ProviderKind, + /// Static API key for providers that use one (`Anthropic`). + /// Filled by `provider::resolve` from router-side config only — + /// never from the inbound request. + pub api_key: Option, +} + +impl UpstreamConfig { + /// The historic constructor shape: an Azure OpenAI / Foundry + /// upstream authenticated via Workload Identity / API-key mode. + #[must_use] + pub fn azure(endpoint: String, deployment: String, sandbox_name: String) -> Self { + Self { + endpoint, + deployment, + sandbox_name, + provider: ProviderKind::AzureOpenAI, + api_key: None, + } + } +} + +/// Credential material resolved for one upstream request. Which +/// header it lands in is provider-specific (`Authorization: Bearer` +/// vs `x-api-key` vs nothing at all for unauthenticated in-cluster +/// Ollama). +pub enum UpstreamCredential { + Bearer(String), + AnthropicApiKey(String), + None, } /// Determine the correct token audience for the upstream endpoint. @@ -51,7 +85,7 @@ fn token_audience(endpoint: &str) -> &'static str { fn build_upstream_headers( request_headers: &HeaderMap, _auth: &WorkloadIdentityAuth, - token: &str, + credential: &UpstreamCredential, endpoint: &str, ) -> Result { let mut headers = HeaderMap::new(); @@ -68,13 +102,32 @@ fn build_upstream_headers( } } - // Both API-key and Entra modes use Authorization: Bearer for the unified - // /openai/v1/ endpoint format. Azure OpenAI accepts API keys as Bearer tokens. - // Copilot also uses Bearer (with the exchanged Copilot JWT). - headers.insert( - "authorization", - HeaderValue::from_str(&format!("Bearer {token}")).context("Invalid token")?, - ); + match credential { + // Both API-key and Entra modes use Authorization: Bearer for the unified + // /openai/v1/ endpoint format. Azure OpenAI accepts API keys as Bearer tokens. + // Copilot also uses Bearer (with the exchanged Copilot JWT). + UpstreamCredential::Bearer(token) => { + headers.insert( + "authorization", + HeaderValue::from_str(&format!("Bearer {token}")).context("Invalid token")?, + ); + } + // Anthropic's Messages API authenticates with `x-api-key` and + // requires an `anthropic-version` header. The inbound SDK + // value (when present) was already copied through above — + // only the default is filled in here. + UpstreamCredential::AnthropicApiKey(key) => { + headers.insert( + "x-api-key", + HeaderValue::from_str(key).context("Invalid Anthropic API key")?, + ); + headers + .entry("anthropic-version") + .or_insert(HeaderValue::from_static("2023-06-01")); + } + // Unauthenticated upstream (in-cluster Ollama). + UpstreamCredential::None => {} + } headers .entry("content-type") .or_insert(HeaderValue::from_static("application/json")); @@ -132,6 +185,38 @@ pub async fn token_for_endpoint( } } +/// Provider-aware credential resolution for a single upstream request. +/// +/// - `AzureOpenAI` → the historic [`token_for_endpoint`] path (Azure +/// WI/IMDS, API key, or Copilot JWT depending on endpoint). +/// - `Anthropic` → the static API key `provider::resolve` copied from +/// router-side config onto `UpstreamConfig.api_key`. Its absence +/// here is a programmer error (resolution fails closed earlier), +/// surfaced as a clean 502 rather than a panic. +/// - `Ollama` → no credential. +pub async fn credential_for_upstream( + auth: &WorkloadIdentityAuth, + copilot: Option<&CopilotTokenCache>, + upstream: &UpstreamConfig, +) -> Result { + match upstream.provider { + ProviderKind::AzureOpenAI => token_for_endpoint(auth, copilot, &upstream.endpoint) + .await + .map(UpstreamCredential::Bearer), + ProviderKind::Anthropic => upstream + .api_key + .clone() + .map(UpstreamCredential::AnthropicApiKey) + .ok_or_else(|| { + anyhow::anyhow!( + "Anthropic upstream selected but no API key on UpstreamConfig — \ + provider resolution must run before forward()" + ) + }), + ProviderKind::Ollama => Ok(UpstreamCredential::None), + } +} + /// Record Prometheus metrics from a completed request. fn record_metrics( upstream: &UpstreamConfig, @@ -194,20 +279,26 @@ pub async fn forward( let (upstream_url, body) = build_upstream_url(auth, upstream, path, request_body)?; - let mode = if is_copilot_endpoint(&upstream.endpoint) { - "copilot" - } else if auth.is_api_key_mode() { - "dev" - } else { - "foundry" + let mode = match upstream.provider { + ProviderKind::Anthropic => "anthropic", + ProviderKind::Ollama => "ollama", + ProviderKind::AzureOpenAI => { + if is_copilot_endpoint(&upstream.endpoint) { + "copilot" + } else if auth.is_api_key_mode() { + "dev" + } else { + "foundry" + } + } }; tracing::info!(sandbox = %upstream.sandbox_name, model = %upstream.deployment, mode = %mode, "Forwarding inference"); - let token = token_for_endpoint(auth, copilot, &upstream.endpoint) + let credential = credential_for_upstream(auth, copilot, upstream) .await .context("Failed to acquire auth token")?; - let headers = build_upstream_headers(request_headers, auth, &token, &upstream.endpoint)?; + let headers = build_upstream_headers(request_headers, auth, &credential, &upstream.endpoint)?; tracing::info!(sandbox = %upstream.sandbox_name, url = %upstream_url, body_len = body.len(), "Sending upstream request"); @@ -430,10 +521,10 @@ pub async fn forward_stream( tracing::info!(sandbox = %upstream.sandbox_name, model = %upstream.deployment, mode = "stream", "Forwarding SSE stream"); - let token = token_for_endpoint(&auth, copilot.as_deref(), &upstream.endpoint) + let credential = credential_for_upstream(&auth, copilot.as_deref(), &upstream) .await .context("Failed to acquire auth token")?; - let headers = build_upstream_headers(&request_headers, &auth, &token, &upstream.endpoint)?; + let headers = build_upstream_headers(&request_headers, &auth, &credential, &upstream.endpoint)?; let start = Instant::now(); @@ -591,6 +682,10 @@ fn is_github_models_endpoint(endpoint: &str) -> bool { /// Uses the unified /openai/v1/ format — works with both API-key and Entra auth. /// /// Routing rules: +/// - Anthropic: no path rewrite — callers pass Messages-API paths +/// (`v1/messages`) verbatim. +/// - Ollama: OpenAI-compat lives under `/v1/` — `chat/completions` +/// becomes `{endpoint}/v1/chat/completions`. /// - GitHub Copilot (`api.githubcopilot.com`): no path rewrite; OpenClaw /// sends OpenAI-shape to `/chat/completions` and Anthropic-shape to /// `/v1/messages`. We forward those paths unchanged. @@ -602,20 +697,34 @@ fn build_upstream_url( path: &str, request_body: Bytes, ) -> Result<(String, Bytes)> { - let url = if is_github_models_endpoint(&upstream.endpoint) - || is_copilot_endpoint(&upstream.endpoint) - { - format!( + let url = match upstream.provider { + ProviderKind::Anthropic => format!( "{}/{}", upstream.endpoint.trim_end_matches('/'), path.trim_start_matches('/'), - ) - } else { - format!( - "{}/openai/v1/{}", + ), + ProviderKind::Ollama => format!( + "{}/v1/{}", upstream.endpoint.trim_end_matches('/'), - path.trim_start_matches('/'), - ) + path.trim_start_matches('/').trim_start_matches("v1/"), + ), + ProviderKind::AzureOpenAI => { + if is_github_models_endpoint(&upstream.endpoint) + || is_copilot_endpoint(&upstream.endpoint) + { + format!( + "{}/{}", + upstream.endpoint.trim_end_matches('/'), + path.trim_start_matches('/'), + ) + } else { + format!( + "{}/openai/v1/{}", + upstream.endpoint.trim_end_matches('/'), + path.trim_start_matches('/'), + ) + } + } }; let body = if let Ok(mut body_json) = serde_json::from_slice::(&request_body) { @@ -642,6 +751,7 @@ fn build_upstream_url( // requesting reasoning encryption on a stripped input is a no-op // for output and Azure rejects it on the input side anyway. if path.trim_start_matches('/').starts_with("responses") + && upstream.provider == ProviderKind::AzureOpenAI && !is_github_models_endpoint(&upstream.endpoint) && !is_copilot_endpoint(&upstream.endpoint) && let Some(obj) = body_json.as_object_mut() diff --git a/inference-router/src/routes/anthropic_messages.rs b/inference-router/src/routes/anthropic_messages.rs index 4bbc1416b..20e6c562d 100644 --- a/inference-router/src/routes/anthropic_messages.rs +++ b/inference-router/src/routes/anthropic_messages.rs @@ -25,7 +25,29 @@ 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; + +/// 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 { ( @@ -264,11 +286,93 @@ pub(super) async fn anthropic_messages( // Slice 2d.1: honour `InferencePolicy.modelPreference.primary.deployment`. crate::routes::apply_model_preference_override(&mut upstream, &policy); - // Copilot exposes a native Anthropic Messages endpoint at /v1/messages. - // Skip translation entirely and forward the body as-is, preserving the - // streaming + tool_use + multi-modal contracts of the Anthropic SDK. - if proxy::is_copilot_endpoint(&upstream.endpoint) { - return forward_anthropic_passthrough(state, sandbox_name, headers, body, upstream).await; + // Retarget at the policy-selected provider (fails closed). + if let Err(e) = crate::routes::apply_provider_resolution(&state, &mut upstream, &policy) { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "provider_resolution", + error = %e, + "InferencePolicy provider could not be resolved (anthropic route)" + ); + let status = match e { + ProviderError::Unimplemented { .. } => StatusCode::NOT_IMPLEMENTED, + _ => StatusCode::SERVICE_UNAVAILABLE, + }; + return deny_response(status, &e.to_string(), "api_error"); + } + + // Guardrail pipeline; a declared-but-unbuildable stage blocks. + let guardrail_pipeline = + match super::chat_completions::build_guardrail_pipeline(&state, &policy) { + Ok(p) => p, + Err(e) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_config", + error = %e, + "guardrail pipeline could not be built (anthropic route) — failing closed" + ); + return deny_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 Messages pass-through (provider: anthropic, or Copilot's + // native /v1/messages) — no translation. + if upstream.provider == ProviderKind::Anthropic + || proxy::is_copilot_endpoint(&upstream.endpoint) + { + return forward_anthropic_passthrough( + state, + sandbox_name, + headers, + body, + upstream, + guardrail_pipeline, + policy.digest.clone(), + ) + .await; } // Translate Anthropic -> OpenAI chat completions request shape. @@ -337,6 +441,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 +506,8 @@ async fn forward_anthropic_passthrough( headers: HeaderMap, body: Bytes, upstream: crate::proxy::UpstreamConfig, + guardrail_pipeline: Option>, + policy_digest: String, ) -> axum::response::Response { let is_stream = serde_json::from_slice::(&body) .ok() @@ -393,10 +539,27 @@ 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))); + // 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, + 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() { + if is_hop_by_hop(n.as_str()) { + continue; + } h.insert(n.clone(), v.clone()); } h.insert( @@ -454,9 +617,59 @@ async fn forward_anthropic_passthrough( } } + // Guardrail output scan (buffered, Anthropic shape). + if status.is_success() + && let Some(p) = guardrail_pipeline + .as_ref() + .filter(|p| p.covers(Direction::Output)) + { + let text = guardrails::scan_text_or_raw( + &resp_body, + guardrails::extract_anthropic_output_text, + ); + match p.scan(&text, Direction::Output).await { + Ok(None) => {} + Ok(Some(v)) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy_digest, + decision = "deny", + gate = "guardrail_output", + categories = ?v.categories, + "guardrail pipeline blocked buffered response (anthropic route)" + ); + return deny_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() { + if is_hop_by_hop(n.as_str()) { + continue; + } h.insert(n.clone(), v.clone()); } if !h.contains_key(axum::http::header::CONTENT_TYPE) { diff --git a/inference-router/src/routes/chat_completions.rs b/inference-router/src/routes/chat_completions.rs index 2b53a7e20..794f44bf2 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,227 @@ fn insert_decision_headers( } } +/// Map a provider-resolution failure onto an OpenAI-shaped error +/// response: 501 for schema-valid-but-unimplemented providers +/// (`bedrock`), 503 for router-side config gaps. Never falls back to +/// Azure — see `routes::apply_provider_resolution`. +pub(super) fn provider_error_response(e: &ProviderError) -> axum::response::Response { + let (status, code) = match e { + ProviderError::Unimplemented { .. } => { + (StatusCode::NOT_IMPLEMENTED, "provider_unimplemented") + } + ProviderError::MissingEndpoint { .. } | ProviderError::MissingCredential { .. } => { + (StatusCode::SERVICE_UNAVAILABLE, "provider_unconfigured") + } + }; + let mut resp = ( + status, + Json(serde_json::json!({ + "error": { "message": e.to_string(), "type": "provider_error", "code": code } + })), + ) + .into_response(); + insert_decision_headers(&mut resp, "blocked", "InferencePolicy", &e.to_string()); + resp +} + +/// 403 response for a confirmed guardrail violation, with the +/// canonical `x-kars-decision*` triplet attached. +pub(super) fn guardrail_violation_response(v: &GuardrailViolation) -> axum::response::Response { + let mut resp = ( + StatusCode::FORBIDDEN, + Json(serde_json::json!({ + "error": { "message": v.message(), "type": "content_policy_violation", "code": v.code() } + })), + ) + .into_response(); + insert_decision_headers(&mut resp, "blocked", "InferencePolicy", &v.message()); + resp +} + +/// Fail-closed response for a guardrail that could not run: 503 for a +/// misconfigured stage, 502 for a backend outage. +pub(super) fn guardrail_error_response(e: &GuardrailError) -> axum::response::Response { + let status = match e { + GuardrailError::Config { .. } => StatusCode::SERVICE_UNAVAILABLE, + GuardrailError::Unavailable { .. } => StatusCode::BAD_GATEWAY, + }; + let mut resp = ( + status, + Json(serde_json::json!({ + "error": { "message": e.to_string(), "type": "guardrail_error", "code": e.code() } + })), + ) + .into_response(); + insert_decision_headers(&mut resp, "blocked", "InferencePolicy", &e.to_string()); + resp +} + +/// Build the policy's guardrail pipeline, or `None` when the policy +/// declares no stages. A declared-but-unbuildable pipeline is a +/// request-blocking error (fail closed). +pub(super) fn build_guardrail_pipeline( + state: &AppState, + policy: &crate::inference_policy_loader::InferencePolicySnapshot, +) -> Result>, GuardrailError> { + if policy.guardrails.is_empty() { + return Ok(None); + } + GuardrailPipeline::from_stages(&policy.guardrails, &state.config, &state.client) + .map(|p| Some(Arc::new(p))) +} + +/// A blocked output scan, kept as data so each transport picks its +/// wire shape (buffered → HTTP 403/502/503, SSE → error frame). +pub(super) enum OutputGuardrailBlock { + Violation(GuardrailViolation), + Error(GuardrailError), +} + +impl OutputGuardrailBlock { + /// SSE frame carrying this block's own type/code, so a backend + /// outage isn't mislabelled a content violation. + pub(super) fn sse_frame(&self) -> bytes::Bytes { + match self { + Self::Violation(v) => guardrails::violation_sse_frame(v), + Self::Error(e) => guardrails::error_sse_frame(e), + } + } +} + +/// Run the output-direction guardrail stages over a buffered +/// OpenAI-shaped response body. `None` when there is nothing to do or +/// the scan passes; `Some(block)` when the response must not reach +/// the client. Emits the audit log line on every block. +pub(super) async fn scan_openai_output_guardrails( + pipeline: Option<&Arc>, + resp_body: &[u8], + sandbox_name: &str, + policy_digest: &str, +) -> Option { + let p = pipeline?; + if !p.covers(Direction::Output) { + return None; + } + let text = guardrails::scan_text_or_raw(resp_body, guardrails::extract_openai_output_text); + match p.scan(&text, Direction::Output).await { + Ok(None) => None, + Ok(Some(v)) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy_digest, + decision = "deny", + gate = "guardrail_output", + categories = ?v.categories, + "guardrail pipeline blocked buffered response" + ); + Some(OutputGuardrailBlock::Violation(v)) + } + Err(e) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy_digest, + decision = "deny", + gate = "guardrail_output", + error = %e, + "guardrail pipeline unavailable — failing closed" + ); + Some(OutputGuardrailBlock::Error(e)) + } + } +} + +/// HTTP-response wrapper over [`scan_openai_output_guardrails`] for +/// the buffered branches: `Err(response)` carries the ready-made +/// block response (403 violation / 502 unavailable / 503 config). +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)), + } +} + +/// 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, @@ -197,6 +421,84 @@ pub(super) async fn chat_completions( // Slice 2d.1: honour `InferencePolicy.modelPreference.primary.deployment`. crate::routes::apply_model_preference_override(&mut upstream, &policy); + // Retarget at the policy-selected provider; fails closed. + if let Err(e) = crate::routes::apply_provider_resolution(&state, &mut upstream, &policy) { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "provider_resolution", + error = %e, + "InferencePolicy provider could not be resolved" + ); + return provider_error_response(&e); + } + + // Anthropic serves the Messages API — an OpenAI-shaped request + // here gets an explicit 501 pointing at /anthropic/v1/messages. + 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; a declared-but-unbuildable stage blocks. + let guardrail_pipeline = match build_guardrail_pipeline(&state, &policy) { + Ok(p) => p, + Err(e) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_config", + error = %e, + "guardrail pipeline could not be built — failing closed" + ); + return guardrail_error_response(&e); + } + }; + + // Input-direction scan, before any upstream forward. + if let Some(ref p) = guardrail_pipeline + && p.covers(Direction::Input) + { + let input_text = guardrails::scan_text_or_raw(&body, guardrails::extract_openai_input_text); + match p.scan(&input_text, Direction::Input).await { + Ok(None) => {} + Ok(Some(v)) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_input", + categories = ?v.categories, + "guardrail pipeline blocked request" + ); + return guardrail_violation_response(&v); + } + Err(e) => { + tracing::warn!( + target: "inference.audit", + sandbox = %sandbox_name, + inference_policy_digest = %policy.digest, + decision = "deny", + gate = "guardrail_input", + error = %e, + "guardrail pipeline unavailable — failing closed" + ); + return guardrail_error_response(&e); + } + } + } + // Defence-in-depth tool-schema filter: even when the upstream // runtime (e.g. a raw OpenAI SDK client outside OpenClaw) sends // `tools[]` schemas the AGT plugin would normally never have @@ -254,6 +556,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 +597,20 @@ pub(super) async fn chat_completions( { budget.record_usage(&sandbox_owned, total).await; } + // Output scan on the Responses-API recovery + // path; block as an SSE frame (already committed + // to text/event-stream). + if let Some(block) = scan_openai_output_guardrails( + guardrail_for_task.as_ref(), + &chat_body, + &sandbox_owned, + &digest_for_task, + ) + .await + { + let _ = tx.send(Ok(block.sse_frame())).await; + return; + } let sse_data = format!( "data: {}\n\ndata: [DONE]\n\n", String::from_utf8_lossy(&chat_body) @@ -347,6 +665,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 +776,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 +831,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 +920,22 @@ pub(super) async fn chat_completions( } chunk }); - let body = Body::from_stream(wrapped); + // Streaming output scan; skipped when no output stages. + let guarded: futures::stream::BoxStream<'static, Result> = + match guardrail_pipeline + .as_ref() + .filter(|p| p.covers(Direction::Output)) + { + Some(p) => guardrails::guard_sse_stream( + wrapped.boxed(), + p.clone(), + guardrails::StreamDialect::OpenAiChat, + sandbox_for_guard, + policy.digest.clone(), + ), + None => wrapped.boxed(), + }; + let body = Body::from_stream(guarded); let mut response = (status, body).into_response(); if let Some(ct) = resp_headers.get("content-type") { response.headers_mut().insert("content-type", ct.clone()); @@ -665,6 +1019,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()); @@ -794,6 +1158,18 @@ 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") @@ -1099,6 +1475,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 38428bd5f..6571567a9 100644 --- a/inference-router/src/routes/mod.rs +++ b/inference-router/src/routes/mod.rs @@ -366,12 +366,45 @@ impl AppState { .and_then(|g| g.clone()) .unwrap_or_else(|| self.config.default_model.clone()); - UpstreamConfig { - endpoint, - deployment, - sandbox_name: sandbox_name.to_string(), + UpstreamConfig::azure(endpoint, deployment, sandbox_name.to_string()) + } +} + +/// Retarget `upstream` at the policy-selected provider (no-op for +/// Azure / no policy). Fails closed on `bedrock` or missing config — +/// the handler must surface the error, not fall back to Azure. +pub(crate) fn apply_provider_resolution( + state: &AppState, + upstream: &mut UpstreamConfig, + policy: &crate::inference_policy_loader::InferencePolicySnapshot, +) -> Result<(), crate::provider::ProviderError> { + let target = crate::provider::resolve(policy.provider.as_deref(), &state.config)?; + match target { + crate::provider::ProviderTarget::AzureOpenAI => {} + crate::provider::ProviderTarget::Anthropic { endpoint, api_key } => { + tracing::info!( + sandbox = %upstream.sandbox_name, + endpoint = %endpoint, + digest = %policy.digest, + "InferencePolicy provider: routing to Anthropic" + ); + upstream.endpoint = endpoint; + upstream.provider = crate::provider::ProviderKind::Anthropic; + upstream.api_key = Some(api_key); + } + crate::provider::ProviderTarget::Ollama { endpoint } => { + tracing::info!( + sandbox = %upstream.sandbox_name, + endpoint = %endpoint, + digest = %policy.digest, + "InferencePolicy provider: routing to Ollama" + ); + upstream.endpoint = endpoint; + upstream.provider = crate::provider::ProviderKind::Ollama; + upstream.api_key = None; } } + Ok(()) } /// Slice 2d.1 — apply `modelPreference.primary.deployment` from a diff --git a/inference-router/tests/agt_governance_integration.rs b/inference-router/tests/agt_governance_integration.rs index d397f43e6..190a1dde1 100644 --- a/inference-router/tests/agt_governance_integration.rs +++ b/inference-router/tests/agt_governance_integration.rs @@ -55,6 +55,12 @@ fn test_state(sandbox: &str, admin_token: Option<&str>) -> AppState { registry_mode: RegistryMode::Local, registry_url: None, provider_override: None, + anthropic_endpoint: "https://api.anthropic.com".into(), + anthropic_api_key: None, + ollama_endpoint: None, + openai_moderation_endpoint: "https://api.openai.com".into(), + openai_moderation_api_key: None, + openai_moderation_model: "omni-moderation-latest".into(), }), budget: TokenBudgetTracker::new(1_000_000, 100_000), policy_provider: Arc::clone(&governance) as Arc, diff --git a/inference-router/tests/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();