From ab3a9a7ce1c31de677e958278d7e7d1ed5ea6326 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Tue, 8 Sep 2026 04:08:26 +0200 Subject: [PATCH 1/4] feat(router): add scoped governed services and bounded telemetry Add separately authenticated operator controls, UID-qualified request scopes, bounded cancellation and policy-gated waits, metadata-only observations, and narrow router-only credential projection. Preserve existing policy enforcement and provider behavior. Add real Kind authentication/mount/reset coverage; independent review, hosted qualification and genuine audit sign-offs remain pending. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- controller/src/providers/signing.rs | 7 + .../src/reconciler/governed_services.rs | 242 ++++++++ controller/src/reconciler/mod.rs | 13 + docs/governed-services.md | 127 +++++ .../2026-09-08-governed-router-services.md | 81 +++ inference-router/src/access_request.rs | 347 ++++++++++++ inference-router/src/access_request/scope.rs | 87 +++ inference-router/src/egress_blocked.rs | 43 +- inference-router/src/failover.rs | 1 + inference-router/src/forward_proxy.rs | 10 + .../src/forward_proxy/governed_tests.rs | 42 ++ inference-router/src/governed_services.rs | 220 ++++++++ inference-router/src/lib.rs | 3 + inference-router/src/main.rs | 19 +- inference-router/src/proxy.rs | 133 ++++- inference-router/src/routes/access_request.rs | 269 +++++++++ inference-router/src/routes/egress.rs | 181 +++++- inference-router/src/routes/governance.rs | 40 ++ inference-router/src/routes/mcp.rs | 19 + inference-router/src/routes/mod.rs | 18 +- inference-router/src/routes/model_routing.rs | 1 + inference-router/src/routes/task_telemetry.rs | 91 +++ inference-router/src/task_telemetry.rs | 259 +++++++++ inference-router/src/task_telemetry/mcp.rs | 79 +++ .../src/task_telemetry/observe.rs | 182 ++++++ .../src/task_telemetry/observe/stream.rs | 197 +++++++ inference-router/src/task_telemetry/parse.rs | 187 ++++++ .../tests/agt_governance_integration.rs | 1 + .../tests/anthropic_buffered_guardrail.rs | 1 + .../tests/chat_output_guardrail_nonjson.rs | 1 + .../tests/common/governed_services.rs | 196 +++++++ .../tests/egress_blocked_endpoint.rs | 1 + inference-router/tests/failover_walk.rs | 3 + inference-router/tests/foundry_route_guard.rs | 1 + .../tests/governed_access_services.rs | 533 ++++++++++++++++++ .../tests/governed_egress_wait.rs | 254 +++++++++ inference-router/tests/governed_telemetry.rs | 484 ++++++++++++++++ .../tests/multi_provider_guardrails.rs | 2 + .../tests/policy_status_endpoint.rs | 1 + inference-router/tests/proxy_fake_upstream.rs | 3 + tests/e2e/governed-services.sh | 134 +++++ tests/e2e/run.sh | 6 + 42 files changed, 4469 insertions(+), 50 deletions(-) create mode 100644 controller/src/reconciler/governed_services.rs create mode 100644 docs/governed-services.md create mode 100644 docs/security-audits/2026-09-08-governed-router-services.md create mode 100644 inference-router/src/access_request.rs create mode 100644 inference-router/src/access_request/scope.rs create mode 100644 inference-router/src/forward_proxy/governed_tests.rs create mode 100644 inference-router/src/governed_services.rs create mode 100644 inference-router/src/routes/access_request.rs create mode 100644 inference-router/src/routes/task_telemetry.rs create mode 100644 inference-router/src/task_telemetry.rs create mode 100644 inference-router/src/task_telemetry/mcp.rs create mode 100644 inference-router/src/task_telemetry/observe.rs create mode 100644 inference-router/src/task_telemetry/observe/stream.rs create mode 100644 inference-router/src/task_telemetry/parse.rs create mode 100644 inference-router/tests/common/governed_services.rs create mode 100644 inference-router/tests/governed_access_services.rs create mode 100644 inference-router/tests/governed_egress_wait.rs create mode 100644 inference-router/tests/governed_telemetry.rs create mode 100644 tests/e2e/governed-services.sh diff --git a/controller/src/providers/signing.rs b/controller/src/providers/signing.rs index bcc0e5f05..5225bc718 100644 --- a/controller/src/providers/signing.rs +++ b/controller/src/providers/signing.rs @@ -168,9 +168,16 @@ pub fn sha256_hex(bytes: &[u8]) -> String { use std::fmt::Write; let _ = write!(out, "{b:02x}"); } + out } +/// Operator-only service credential, generated by the standard CSPRNG. +pub fn generate_service_token() -> String { + use rand::distr::{Alphanumeric, SampleString}; + Alphanumeric.sample_string(&mut rand::rng(), 64) +} + /// Preserve the existing 128-bit content identifier used by skills, profiles /// and commons. Signed receipt payloads continue using the full SHA-256 digest. pub fn content_digest(bytes: &[u8]) -> String { diff --git a/controller/src/reconciler/governed_services.rs b/controller/src/reconciler/governed_services.rs new file mode 100644 index 000000000..9543653c5 --- /dev/null +++ b/controller/src/reconciler/governed_services.rs @@ -0,0 +1,242 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Narrow router service identity/credential projection; no approvals or grants. + +use crate::{crd::KarsSandbox, kars_task::KarsTask}; +use k8s_openapi::api::core::v1::{Namespace, Secret}; +use kube::{Api, Client, ResourceExt, api::PostParams}; +use serde_json::{Value, json}; + +const SECRET: &str = "router-services-admin"; +const SOURCE_UID: &str = "kars.azure.com/sandbox-uid"; +const NAMESPACE_UID: &str = "kars.azure.com/namespace-uid"; + +fn api_error(error: kube::Error) -> String { + match error { + kube::Error::Api(status) => { + format!("Governed service credential API status {}", status.code) + } + _ => "Governed service credential API/transport failure".into(), + } +} + +fn authorized_task( + task: &KarsTask, + workspace: &str, + sandbox: &str, + name: &str, + uid: &str, +) -> Option { + let status = task.status.as_ref()?; + let authorization = task.spec.authorization_digest(); + (task.metadata.namespace.as_deref() == Some(workspace) + && task.metadata.name.as_deref() == Some(name) + && task.metadata.uid.as_deref() == Some(uid) + && task + .metadata + .generation + .is_some_and(|generation| generation > 0) + && task.metadata.deletion_timestamp.is_none() + && status.observed_generation == task.metadata.generation + && status.phase.as_deref() == Some("Ready") + && status + .conditions + .iter() + .flatten() + .any(|condition| condition.type_ == "Ready" && condition.status == "True") + && status.envelope_digest.as_deref() == Some(authorization.as_str()) + && status + .sandbox_ref + .as_ref() + .map(|reference| reference.name.as_str()) + == Some(sandbox) + && task + .spec + .execution + .as_ref() + .is_some_and(|execution| execution.launch) + && crate::kars_task::validate_execution_contract(&task.spec).is_ok()) + .then_some(authorization) +} + +pub async fn ensure( + client: &Client, + sandbox: &KarsSandbox, + namespace: &Namespace, +) -> Result { + // Reuse the authoritative namespace claim path, not labels or a caller's + // requested namespace. Recreated CRs/namespaces cannot inherit this token. + let (live, owned) = super::namespace_ownership::ensure(client, sandbox) + .await + .map_err(|_| "Governed service namespace authority could not be verified")?; + let owned = owned.ok_or("Governed service namespace is absent")?; + if live.metadata.uid != sandbox.metadata.uid || owned.metadata.uid != namespace.metadata.uid { + return Err("Governed service namespace or Sandbox incarnation changed".into()); + } + let sandbox_uid = live + .metadata + .uid + .as_deref() + .filter(|uid| !uid.is_empty()) + .ok_or("Sandbox UID missing")?; + let namespace_uid = owned + .metadata + .uid + .as_deref() + .filter(|uid| !uid.is_empty()) + .ok_or("Namespace UID missing")?; + let workspace = live.namespace().ok_or("Sandbox workspace missing")?; + let name = live.name_any(); + let mut task_identity = Value::Null; + let mut task_authorization = None; + let mut task_generation = None; + if let Some(owner) = live.metadata.owner_references.as_ref().and_then(|owners| { + owners.iter().find(|owner| { + owner.kind == "KarsTask" + && owner.api_version == "kars.azure.com/v1alpha1" + && owner.controller == Some(true) + }) + }) { + let tasks: Api = Api::namespaced(client.clone(), &workspace); + let task = tasks.get(&owner.name).await.map_err(api_error)?; + let authorization = authorized_task(&task, &workspace, &name, &owner.name, &owner.uid) + .ok_or("Task UID/effective authorization does not bind this Sandbox")?; + task_identity = json!({"namespace":workspace,"name":owner.name,"uid":owner.uid}); + task_authorization = Some(authorization); + task_generation = task.metadata.generation; + } + let secrets: Api = Api::namespaced(client.clone(), &owned.name_any()); + if let Some(secret) = secrets.get_opt(SECRET).await.map_err(api_error)? { + let annotations = secret.metadata.annotations.as_ref(); + if secret.metadata.deletion_timestamp.is_some() + || secret.metadata.uid.as_deref().is_none_or(str::is_empty) + || secret + .metadata + .resource_version + .as_deref() + .is_none_or(str::is_empty) + || secret.metadata.name.as_deref() != Some(SECRET) + || secret.metadata.namespace.as_deref() != Some(owned.name_any().as_str()) + || secret + .metadata + .owner_references + .as_ref() + .is_some_and(|owners| !owners.is_empty()) + || annotations + .and_then(|annotations| annotations.get(SOURCE_UID)) + .map(String::as_str) + != Some(sandbox_uid) + || annotations + .and_then(|annotations| annotations.get(NAMESPACE_UID)) + .map(String::as_str) + != Some(namespace_uid) + || secret + .data + .as_ref() + .and_then(|data| data.get("control-token")) + .is_none_or(|value| { + value.0.len() != 64 || value.0.iter().any(|byte| !byte.is_ascii_graphic()) + }) + { + return Err( + "Existing governed service credential has conflicting ownership or invalid data" + .into(), + ); + } + } else { + let secret:Secret=serde_json::from_value(json!({ + "apiVersion":"v1","kind":"Secret","metadata":{"name":SECRET,"namespace":owned.name_any(), + "labels":{"app.kubernetes.io/managed-by":"kars-controller"}, + "annotations":{SOURCE_UID:sandbox_uid,NAMESPACE_UID:namespace_uid}}, + "stringData":{"control-token":crate::providers::signing::generate_service_token()}, + })).map_err(|_|"Governed service credential serialization failed")?; + secrets + .create(&PostParams::default(), &secret) + .await + .map_err(api_error)?; + } + Ok( + json!({"sandbox":{"namespace":workspace,"name":name,"uid":sandbox_uid}, + "namespace_uid":namespace_uid,"task":task_identity,"task_authorization":task_authorization, + "task_generation":task_generation,"managed":true}), + ) +} + +pub fn mount(pod: &mut Value) { + pod["volumes"] + .as_array_mut() + .expect("pod volumes array") + .push(json!({ + "name":"governed-services-control","secret":{"secretName":SECRET, + "items":[{"key":"control-token","path":"control-token"}]}})); + for container in pod["containers"] + .as_array_mut() + .expect("pod containers array") + { + if container["name"] == "inference-router" { + container["volumeMounts"].as_array_mut().expect("router mounts array").push(json!({ + "name":"governed-services-control","mountPath":"/etc/kars/services","readOnly":true})); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + fn task() -> KarsTask { + let mut task:KarsTask=serde_json::from_value(json!({ + "apiVersion":"kars.azure.com/v1alpha1","kind":"KarsTask", + "metadata":{"name":"task","namespace":"workspace","uid":"task-uid","generation":1}, + "spec":{"objective":"inspect","envelope":{"tier":1,"authorityCeiling":1,"delegationDepth":0},"execution":{"launch":true}}, + "status":{"phase":"Ready","observedGeneration":1,"sandboxRef":{"name":"sandbox"}, + "conditions":[{"type":"Ready","status":"True","reason":"Valid","message":"validated","lastTransitionTime":"2026-09-08T00:00:00Z"}]}, + })).unwrap(); + task.status.as_mut().unwrap().envelope_digest = Some(task.spec.authorization_digest()); + task + } + #[test] + fn service_task_attribution_requires_full_authorization_and_current_uid() { + let mut task = task(); + assert!(authorized_task(&task, "workspace", "sandbox", "task", "task-uid").is_some()); + assert!(authorized_task(&task, "other", "sandbox", "task", "task-uid").is_none()); + assert!(authorized_task(&task, "workspace", "sandbox", "task", "old-task-uid").is_none()); + task.spec.blueprint = Some(crate::kars_task::TaskBlueprint { + instructions: Some("changed after approval".into()), + ..Default::default() + }); + assert!(authorized_task(&task, "workspace", "sandbox", "task", "task-uid").is_none()); + } + #[test] + fn envelope_only_hash_and_unsupported_launch_budgets_never_authorize_projection() { + let mut task = task(); + task.status.as_mut().unwrap().envelope_digest = Some(task.spec.envelope.digest()); + assert!(authorized_task(&task, "workspace", "sandbox", "task", "task-uid").is_none()); + task.spec.envelope.budget = Some(crate::kars_task::TaskBudget { + tokens: Some(10), + ..Default::default() + }); + task.status.as_mut().unwrap().envelope_digest = Some(task.spec.authorization_digest()); + assert!(authorized_task(&task, "workspace", "sandbox", "task", "task-uid").is_none()); + } + #[test] + fn control_credential_is_never_mounted_in_agent_containers() { + let mut pod = json!({"volumes":[],"containers":[ + {"name":"openclaw","volumeMounts":[]},{"name":"agent","volumeMounts":[]}, + {"name":"inference-router","volumeMounts":[]}]}); + mount(&mut pod); + assert_eq!(pod["containers"][0]["volumeMounts"], json!([])); + assert_eq!(pod["containers"][1]["volumeMounts"], json!([])); + assert_eq!( + pod["containers"][2]["volumeMounts"][0]["mountPath"], + "/etc/kars/services" + ); + } + #[test] + fn standard_control_tokens_are_bounded_and_distinct() { + let first = crate::providers::signing::generate_service_token(); + assert_eq!(first.len(), 64); + assert!(first.bytes().all(|byte| byte.is_ascii_alphanumeric())); + assert_ne!(first, crate::providers::signing::generate_service_token()); + } +} diff --git a/controller/src/reconciler/mod.rs b/controller/src/reconciler/mod.rs index 8e86917fe..d7dabe3cc 100644 --- a/controller/src/reconciler/mod.rs +++ b/controller/src/reconciler/mod.rs @@ -38,6 +38,7 @@ pub(crate) mod byo_contract; mod credential_sources; mod dev_env; pub(crate) mod governance_mounts; +mod governed_services; mod inference; mod mcp_egress; pub(crate) mod namespace_ownership; @@ -1614,6 +1615,15 @@ async fn reconcile(sandbox: Arc, ctx: Arc) -> Result, ctx: Arc) -> Result, ctx: Arc) -> Result, + pub port: Option, +} + +#[derive(Clone, Serialize)] +pub struct Entry { + pub request_id: String, + pub scope_id: String, + pub kind: String, + pub target: String, + pub reason: String, + pub tier: Option, + pub port: Option, + pub status: Status, + pub decision: Option, + pub count: u32, + pub first_seen_unix: u64, + pub expires_at_unix: u64, + #[serde(skip)] + expires: Instant, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Error { + Invalid, + StaleScope, + Full, + RateLimited, + Missing, + Expired, + WaitTimeout, + Terminal, + Unavailable, +} + +struct Inner { + scope: Scope, + generation: u64, + sequence: u64, + requests: VecDeque, + window: Instant, + rate: u32, + cancellation: CancellationToken, +} + +pub struct AccessRequestBuffer { + inner: Mutex, + changed: watch::Sender, + capacity: usize, + rate_limit: u32, + ttl: Duration, +} + +impl AccessRequestBuffer { + pub fn new(identity: Identity) -> Self { + Self::with_limits(identity, CAPACITY, REQUESTS_PER_MINUTE, TTL) + } + + pub fn with_limits( + identity: Identity, + capacity: usize, + rate_limit: u32, + ttl: Duration, + ) -> Self { + let scope = Scope { + id: format!("{:032x}:0", rand::random::()), + identity, + assignment_id: None, + }; + Self { + inner: Mutex::new(Inner { + scope, + generation: 0, + sequence: 0, + requests: VecDeque::new(), + window: Instant::now(), + rate: 0, + cancellation: CancellationToken::new(), + }), + changed: watch::channel(0).0, + capacity: capacity.clamp(1, 256), + rate_limit: rate_limit.clamp(1, 128), + ttl: ttl.min(Duration::from_secs(3600)), + } + } + + fn notify(&self) { + self.changed + .send_modify(|value| *value = value.saturating_add(1)); + } + pub fn subscribe(&self) -> watch::Receiver { + self.changed.subscribe() + } + pub fn scope(&self) -> Result { + Ok(self + .inner + .lock() + .map_err(|_| Error::Unavailable)? + .scope + .clone()) + } + + fn check(inner: &Inner, scope: &str) -> Result<(), Error> { + if inner.scope.id == scope { + Ok(()) + } else { + Err(Error::StaleScope) + } + } + fn expire(inner: &mut Inner) { + for entry in &mut inner.requests { + if entry.expires <= Instant::now() + && matches!(entry.status, Status::Pending | Status::Approved) + { + entry.status = Status::Expired; + } + } + } + + pub fn record(&self, mut request: Request) -> Result<(Entry, bool), Error> { + validate(&mut request)?; + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::check(&inner, &request.scope_id)?; + if inner.window.elapsed() >= Duration::from_secs(60) { + inner.window = Instant::now(); + inner.rate = 0; + } + if inner.rate >= self.rate_limit { + return Err(Error::RateLimited); + } + inner.rate += 1; + Self::expire(&mut inner); + if let Some(entry) = inner.requests.iter_mut().find(|entry| { + entry.kind == request.kind + && entry.target == request.target + && entry.port == request.port + && entry.tier == request.tier + && entry.expires > Instant::now() + }) { + // An agent cannot rewrite the reviewed payload or extend its lifetime. + entry.count = entry.count.saturating_add(1); + return Ok((entry.clone(), false)); + } + if inner.requests.len() >= self.capacity { + if let Some(index) = inner + .requests + .iter() + .position(|entry| entry.status != Status::Pending) + { + inner.requests.remove(index); + } else { + return Err(Error::Full); + } + } + inner.sequence = inner.sequence.checked_add(1).ok_or(Error::Unavailable)?; + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + let entry = Entry { + request_id: format!("{}:{}", inner.scope.id, inner.sequence), + scope_id: inner.scope.id.clone(), + kind: request.kind, + target: request.target, + reason: request.reason, + tier: request.tier, + port: request.port, + status: Status::Pending, + decision: None, + count: 1, + first_seen_unix: now, + expires_at_unix: now.saturating_add(self.ttl.as_secs()), + expires: Instant::now() + self.ttl, + }; + inner.requests.push_back(entry.clone()); + drop(inner); + self.notify(); + Ok((entry, true)) + } + + pub fn snapshot(&self, scope: &str) -> Result, Error> { + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::check(&inner, scope)?; + Self::expire(&mut inner); + Ok(inner.requests.iter().cloned().collect()) + } + + pub fn transition(&self, scope: &str, id: &str, status: Status) -> Result { + if !matches!( + status, + Status::Approved | Status::Denied | Status::Cancelled + ) { + return Err(Error::Invalid); + } + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::check(&inner, scope)?; + Self::expire(&mut inner); + let entry = inner + .requests + .iter_mut() + .find(|entry| entry.request_id == id) + .ok_or(Error::Missing)?; + if entry.expires <= Instant::now() { + return Err(Error::Expired); + } + if entry.status != Status::Pending + && !(status == Status::Cancelled && entry.status == Status::Approved) + { + return Err(Error::Terminal); + } + if matches!(status, Status::Approved | Status::Denied) { + entry.decision = Some(status); + } + entry.status = status; + let entry = entry.clone(); + drop(inner); + self.notify(); + Ok(entry) + } + + pub fn reset( + &self, + expected_scope: &str, + assignment_id: Option, + ) -> Result<(Scope, usize), Error> { + if assignment_id + .as_deref() + .is_some_and(|id| !scope::identifier(id, 253)) + { + return Err(Error::Invalid); + } + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::check(&inner, expected_scope)?; + inner.generation = inner.generation.checked_add(1).ok_or(Error::Unavailable)?; + let instance = inner.scope.id.split(':').next().ok_or(Error::Unavailable)?; + inner.scope.id = format!("{instance}:{}", inner.generation); + inner.scope.assignment_id = assignment_id; + let cleared = inner.requests.len(); + inner.requests.clear(); + inner.cancellation.cancel(); + inner.cancellation = CancellationToken::new(); + let scope = inner.scope.clone(); + drop(inner); + self.notify(); + Ok((scope, cleared)) + } + + pub fn cancellation(&self, scope: &str) -> Result { + let inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::check(&inner, scope)?; + Ok(inner.cancellation.clone()) + } + + pub fn entry(&self, scope: &str, id: &str) -> Result { + self.snapshot(scope)? + .into_iter() + .find(|entry| entry.request_id == id) + .ok_or(Error::Missing) + } +} + +fn validate(request: &mut Request) -> Result<(), Error> { + request.kind = request.kind.trim().to_ascii_lowercase(); + request.target = request.target.trim().to_string(); + if ![ + "egress", + "tool", + "skill", + "mcp", + "command", + "permission", + "clarification", + "tier", + ] + .contains(&request.kind.as_str()) + || request.reason.len() > 512 + || request.reason.chars().any(|c| c.is_control() && c != '\n') + { + return Err(Error::Invalid); + } + if request.kind == "tier" { + if !request.tier.is_some_and(|tier| (1..=5).contains(&tier)) + || !request.target.is_empty() + || request.port.is_some() + { + return Err(Error::Invalid); + } + } else { + if !scope::identifier(&request.target, 253) || request.tier.is_some() { + return Err(Error::Invalid); + } + if request.kind == "egress" { + if request.target.contains(['/', ':']) { + return Err(Error::Invalid); + } + request.target = + crate::egress_blocked::normalize_host(&request.target).ok_or(Error::Invalid)?; + request.port = Some(request.port.unwrap_or(443)); + if request.port == Some(0) { + return Err(Error::Invalid); + } + } else if request.port.is_some() { + return Err(Error::Invalid); + } + } + Ok(()) +} diff --git a/inference-router/src/access_request/scope.rs b/inference-router/src/access_request/scope.rs new file mode 100644 index 000000000..dae8b1a32 --- /dev/null +++ b/inference-router/src/access_request/scope.rs @@ -0,0 +1,87 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct ResourceIdentity { + pub namespace: String, + pub name: String, + pub uid: String, +} + +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(deny_unknown_fields)] +pub struct Identity { + pub sandbox: ResourceIdentity, + pub namespace_uid: String, + #[serde(default)] + pub task: Option, + #[serde(default)] + pub task_authorization: Option, + #[serde(default)] + pub task_generation: Option, + pub managed: bool, +} + +#[derive(Clone, Debug, Serialize)] +pub struct Scope { + pub id: String, + pub identity: Identity, + pub assignment_id: Option, +} + +pub fn identifier(value: &str, max: usize) -> bool { + !value.is_empty() + && value.len() <= max + && value + .bytes() + .all(|c| c.is_ascii_alphanumeric() || b"-._:/".contains(&c)) +} + +impl Identity { + pub fn standalone(name: &str) -> Self { + Self { + sandbox: ResourceIdentity { + namespace: "standalone".into(), + name: name.into(), + uid: "process-local".into(), + }, + namespace_uid: "process-local".into(), + task: None, + task_authorization: None, + task_generation: None, + managed: false, + } + } + pub fn valid(&self, sandbox_name: &str) -> bool { + let valid = |identity: &ResourceIdentity| { + identifier(&identity.namespace, 253) + && identifier(&identity.name, 253) + && identifier(&identity.uid, 128) + }; + self.sandbox.name == sandbox_name + && valid(&self.sandbox) + && identifier(&self.namespace_uid, 128) + && self + .task + .as_ref() + .is_none_or(|task| valid(task) && task.namespace == self.sandbox.namespace) + && match &self.task { + Some(_) => { + self.task_generation + .is_some_and(|generation| generation > 0) + && self + .task_authorization + .as_deref() + .and_then(|digest| digest.strip_prefix("sha256:")) + .is_some_and(|digest| { + digest.len() == 64 + && digest.bytes().all(|byte| byte.is_ascii_hexdigit()) + }) + } + None => self.task_authorization.is_none() && self.task_generation.is_none(), + } + } +} diff --git a/inference-router/src/egress_blocked.rs b/inference-router/src/egress_blocked.rs index 72646c1f1..7ca40d4f6 100644 --- a/inference-router/src/egress_blocked.rs +++ b/inference-router/src/egress_blocked.rs @@ -9,7 +9,7 @@ //! paths, headers, query strings, or payload data are ever stored. use std::collections::{HashMap, VecDeque}; -use std::sync::Mutex; +use std::sync::{Arc, Mutex, Weak}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; /// Default capacity when the operator does not override. @@ -51,6 +51,7 @@ pub enum RecordOutcome { /// Bounded ring buffer of blocked-attempt records. pub struct BlockedBuffer { + services: Mutex>>, inner: Mutex, capacity: usize, rate_limit_window: Duration, @@ -71,6 +72,7 @@ struct RateState { impl BlockedBuffer { pub fn new(capacity: usize, rate_limit_window: Duration, rate_limit_per_source: u32) -> Self { Self { + services: Mutex::new(None), inner: Mutex::new(Inner { by_key: HashMap::new(), order: VecDeque::new(), @@ -96,7 +98,42 @@ impl BlockedBuffer { /// possible results. Hostname-only — never store paths, headers, or /// payload data. pub fn record(&self, source_sandbox: &str, host: &str, port: u16) -> RecordOutcome { - self.record_at(Instant::now(), unix_now(), source_sandbox, host, port) + let result = self.record_at(Instant::now(), unix_now(), source_sandbox, host, port); + if matches!( + result, + RecordOutcome::Recorded | RecordOutcome::Deduplicated + ) && let Some(services) = self + .services + .lock() + .ok() + .and_then(|services| services.as_ref().and_then(Weak::upgrade)) + && let Ok(scope) = services.requests.scope() + && scope.identity.sandbox.name == source_sandbox + { + let suffix = format!(":{port}"); + let host = host.strip_suffix(&suffix).unwrap_or(host); + let _ = services.record_egress_in_scope(&scope.id, host, port); + services.telemetry.record_policy(&scope.id, "egress", false); + } + result + } + + pub fn bind_services(&self, services: &Arc) { + if let Ok(mut target) = self.services.lock() { + *target = Some(Arc::downgrade(services)); + } + } + pub fn observe_allowed(&self, source_sandbox: &str) { + if let Some(services) = self + .services + .lock() + .ok() + .and_then(|services| services.as_ref().and_then(Weak::upgrade)) + && let Ok(scope) = services.requests.scope() + && scope.identity.sandbox.name == source_sandbox + { + services.telemetry.record_policy(&scope.id, "egress", true); + } } /// Testable record entry point — accepts an injected clock pair. @@ -269,7 +306,7 @@ fn unix_now() -> u64 { } /// Lowercase ASCII, strip a single trailing dot, reject empty / IP literal. -fn normalize_host(host: &str) -> Option { +pub(crate) fn normalize_host(host: &str) -> Option { let trimmed = host.trim(); if trimmed.is_empty() { return None; diff --git a/inference-router/src/failover.rs b/inference-router/src/failover.rs index 1bf916150..ecb522045 100644 --- a/inference-router/src/failover.rs +++ b/inference-router/src/failover.rs @@ -466,6 +466,7 @@ mod tests { fn upstream(dep: &str) -> UpstreamConfig { UpstreamConfig { + telemetry: None, endpoint: "https://example.openai.azure.com".into(), deployment: dep.to_string(), sandbox_name: "sbx".into(), diff --git a/inference-router/src/forward_proxy.rs b/inference-router/src/forward_proxy.rs index f5acae588..7e0541dd1 100644 --- a/inference-router/src/forward_proxy.rs +++ b/inference-router/src/forward_proxy.rs @@ -21,6 +21,10 @@ use tokio_util::sync::CancellationToken; use crate::blocklist::Blocklist; use crate::egress_blocked::BlockedBuffer; +#[cfg(test)] +#[path = "forward_proxy/governed_tests.rs"] +mod governed_tests; + /// Maximum concurrent tunnel connections (prevents resource exhaustion). const MAX_CONCURRENT_TUNNELS: usize = 256; @@ -324,6 +328,8 @@ async fn handle_connect( return Ok(()); } + blocked_egress.observe_allowed(sandbox); + // Resolve DNS immediately after policy check and validate against private IPs let resolved = match resolve_and_validate(&domain, port, sandbox, blocked_egress).await { Ok(addr) => addr, @@ -416,6 +422,8 @@ async fn handle_http( return Ok(()); } + blocked_egress.observe_allowed(sandbox); + // Resolve + validate (prevents DNS rebinding to private IPs) let (host, port) = parse_host_port(&domain, 80); let resolved = match resolve_and_validate(&host, port, sandbox, blocked_egress).await { @@ -495,6 +503,8 @@ async fn handle_tls_redirect( return Ok(()); } + blocked_egress.observe_allowed(sandbox); + // Resolve + validate (prevents DNS rebinding to private IPs) let resolved = match resolve_and_validate(&domain, 443, sandbox, blocked_egress).await { Ok(addr) => addr, diff --git a/inference-router/src/forward_proxy/governed_tests.rs b/inference-router/src/forward_proxy/governed_tests.rs new file mode 100644 index 000000000..3e6fff692 --- /dev/null +++ b/inference-router/src/forward_proxy/governed_tests.rs @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::{access_request::Identity, governed_services::GovernedServices}; + +#[tokio::test] +async fn denied_connect_records_a_scoped_request_without_an_implicit_wait() { + let services = Arc::new(GovernedServices::new(Identity::standalone("test"), None)); + let scope = services.requests.scope().unwrap().id; + let blocked = Arc::new(BlockedBuffer::with_defaults()); + blocked.bind_services(&services); + let blocklist = Blocklist::new(None).await; + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + let (socket, _) = listener.accept().await.unwrap(); + handle_connection(socket, &blocklist, "test", &blocked) + .await + .unwrap(); + }); + let mut client = TcpStream::connect(address).await.unwrap(); + client + .write_all(b"CONNECT blocked.example:443 HTTP/1.1\r\nHost: blocked.example\r\n\r\n") + .await + .unwrap(); + let mut response = Vec::new(); + tokio::time::timeout( + std::time::Duration::from_secs(1), + client.read_to_end(&mut response), + ) + .await + .unwrap() + .unwrap(); + server.await.unwrap(); + assert!(String::from_utf8_lossy(&response).starts_with("HTTP/1.1 403")); + let entries = services.requests.snapshot(&scope).unwrap(); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0].target, "blocked.example"); + assert_eq!(entries[0].port, Some(443)); + assert_eq!(services.wait_slots(), 16); +} diff --git a/inference-router/src/governed_services.rs b/inference-router/src/governed_services.rs new file mode 100644 index 000000000..298923f5a --- /dev/null +++ b/inference-router/src/governed_services.rs @@ -0,0 +1,220 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use crate::{ + access_request::{AccessRequestBuffer, Error, Identity, Request, Scope, Status}, + blocklist::Blocklist, + task_telemetry::TaskTelemetry, +}; +use std::{net::IpAddr, sync::Arc, time::Duration}; +use tokio::sync::Semaphore; +use tokio_util::sync::CancellationToken; + +pub const CONTROL_TOKEN_PATH: &str = "/etc/kars/services/control-token"; +pub const MAX_WAIT: Duration = Duration::from_secs(300); +const WAIT_CAPACITY: usize = 16; + +pub struct GovernedServices { + pub requests: AccessRequestBuffer, + pub telemetry: Arc, + control_token: Option, + pub allow_ips: Option>, + pub identity_valid: bool, + pub shutdown: CancellationToken, + waits: Arc, + reset_gate: std::sync::Mutex<()>, +} + +impl Default for GovernedServices { + fn default() -> Self { + Self::new(Identity::standalone("unknown"), None) + } +} + +impl GovernedServices { + pub fn new(identity: Identity, control_token: Option) -> Self { + let control_token = control_token.filter(|token| { + (32..=256).contains(&token.len()) + && token.is_ascii() + && !token.chars().any(|c| c.is_whitespace() || c.is_control()) + }); + let requests = AccessRequestBuffer::new(identity); + let scope = requests + .scope() + .expect("fresh service state is not poisoned"); + Self { + requests, + telemetry: Arc::new(TaskTelemetry::new(scope.id)), + control_token, + allow_ips: None, + identity_valid: true, + shutdown: CancellationToken::new(), + waits: Arc::new(Semaphore::new(WAIT_CAPACITY)), + reset_gate: std::sync::Mutex::new(()), + } + } + pub fn from_env(sandbox: &str) -> Self { + let identity = std::env::var("KARS_SERVICE_IDENTITY_JSON") + .ok() + .map(|raw| serde_json::from_str::(&raw).ok()); + let valid = identity.as_ref().is_none_or(|identity| { + identity + .as_ref() + .is_some_and(|identity| identity.valid(sandbox)) + }); + let identity = identity + .flatten() + .unwrap_or_else(|| Identity::standalone(sandbox)); + let token = std::env::var("KARS_SERVICES_ADMIN_TOKEN") + .ok() + .or_else(|| std::fs::read_to_string(CONTROL_TOKEN_PATH).ok()) + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()); + let mut services = Self::new(identity, token); + services.identity_valid = valid; + services.allow_ips = std::env::var("ROUTER_ADMIN_ALLOW_IPS") + .ok() + .filter(|value| !value.trim().is_empty()) + .map(|value| { + value + .split(',') + .map(str::trim) + .map(str::parse::) + .collect::, _>>() + .unwrap_or_default() + }); + services + } + pub fn control_authorized(&self, provided: Option<&str>) -> bool { + self.identity_valid + && self + .control_token + .as_deref() + .zip(provided) + .is_some_and(|(expected, provided)| { + crate::handoff::constant_time_eq(expected.as_bytes(), provided.as_bytes()) + }) + } + pub fn control_configured(&self) -> bool { + self.identity_valid && self.control_token.is_some() + } + pub fn reset( + &self, + expected: &str, + assignment: Option, + ) -> Result<(Scope, usize), Error> { + let _guard = self.reset_gate.lock().map_err(|_| Error::Unavailable)?; + let result = self.requests.reset(expected, assignment)?; + self.telemetry.reset(result.0.id.clone()); + Ok(result) + } + pub fn record_egress_in_scope( + &self, + scope: &str, + host: &str, + port: u16, + ) -> Result { + self.requests + .record(Request { + scope_id: scope.into(), + kind: "egress".into(), + target: host.into(), + reason: "Egress policy blocked this destination".into(), + tier: None, + port: Some(port), + }) + .map(|(entry, _)| entry) + } + pub fn wait_slots(&self) -> usize { + self.waits.available_permits() + } + + pub async fn wait_for_decision( + &self, + scope: &str, + id: &str, + timeout: Duration, + ) -> Result { + let _permit = self + .waits + .clone() + .try_acquire_owned() + .map_err(|_| Error::Full)?; + let cancellation = self.requests.cancellation(scope)?; + let mut changed = self.requests.subscribe(); + let future = async { + loop { + let entry = self.requests.entry(scope, id)?; + if entry.status != Status::Pending { + return Ok(entry); + } + tokio::select! { + _ = cancellation.cancelled() => return Err(Error::StaleScope), + _ = self.shutdown.cancelled() => return Err(Error::Unavailable), + _ = changed.changed() => {}, + _ = tokio::time::sleep(Duration::from_millis(100)) => {}, + } + } + }; + tokio::time::timeout(timeout.min(MAX_WAIT), future) + .await + .map_err(|_| Error::WaitTimeout)? + } + + pub async fn wait_for_egress( + &self, + blocklist: &Blocklist, + scope: &str, + id: &str, + target: &str, + sandbox: &str, + timeout: Duration, + ) -> Result<(), Error> { + let url = reqwest::Url::parse(target).map_err(|_| Error::Invalid)?; + let host = url.host_str().ok_or(Error::Invalid)?; + let port = url.port_or_known_default().ok_or(Error::Invalid)?; + let owner = self.requests.scope()?; + if owner.id != scope || owner.identity.sandbox.name != sandbox { + return Err(Error::StaleScope); + } + let _permit = self + .waits + .clone() + .try_acquire_owned() + .map_err(|_| Error::Full)?; + let cancellation = self.requests.cancellation(scope)?; + let mut changed = self.requests.subscribe(); + let future = async { + loop { + let entry = self.requests.entry(scope, id)?; + if entry.kind != "egress" + || !entry.target.eq_ignore_ascii_case(host) + || entry.port != Some(port) + { + return Err(Error::Invalid); + } + match entry.status { + Status::Denied | Status::Cancelled => return Err(Error::Terminal), + Status::Expired => return Err(Error::Expired), + Status::Approved => { + // The decision API never changes enforcement. Only the + // normal signed-policy path can make this check pass. + if blocklist.check_egress(target, sandbox).await.is_ok() { + return Ok(()); + } + } + Status::Pending => {} + } + tokio::select! { + _ = cancellation.cancelled() => return Err(Error::StaleScope), + _ = self.shutdown.cancelled() => return Err(Error::Unavailable), + _ = changed.changed() => {}, + _ = tokio::time::sleep(Duration::from_millis(100)) => {}, + } + } + }; + tokio::time::timeout(timeout.min(MAX_WAIT), future) + .await + .map_err(|_| Error::WaitTimeout)? + } +} diff --git a/inference-router/src/lib.rs b/inference-router/src/lib.rs index a8cc38a83..9257c75db 100644 --- a/inference-router/src/lib.rs +++ b/inference-router/src/lib.rs @@ -15,6 +15,7 @@ pub mod a2a; pub mod a2a_mtls; +pub mod access_request; pub mod audit; pub mod audit_jsonl; pub mod audit_sink; @@ -32,6 +33,7 @@ pub mod errors; pub mod failover; pub mod forward_proxy; pub mod governance; +pub mod governed_services; pub mod guardrails; pub mod handoff; pub mod inference_policy_loader; @@ -49,6 +51,7 @@ pub mod routes; pub mod safety; pub mod sidecar_client; pub mod spawn; +pub mod task_telemetry; pub mod telemetry; /// Select RustCrypto for JWT signing and verification. Workspace builds also diff --git a/inference-router/src/main.rs b/inference-router/src/main.rs index 27d6809dc..282e2a855 100644 --- a/inference-router/src/main.rs +++ b/inference-router/src/main.rs @@ -323,6 +323,7 @@ async fn main() -> Result<()> { .unwrap_or_default(), ); + let services_shutdown = state.services.shutdown.clone(); let app = { // Public routes — no admin token required (health, metrics, inference, Foundry proxies, mesh) let public = Router::new() @@ -449,16 +450,19 @@ async fn main() -> Result<()> { let memory_binding_for_platform = state.memory_binding.clone(); let policy_status_for_platform = state.policy_status.clone(); + let telemetry = state.services.telemetry.clone(); + let services = routes::governed_service_routes(state.clone()).with_state(state.clone()); let merged = public .merge(protected) .merge(handoff_init) .merge(handoff_mutations) .merge(handoff_status) .with_state(state) - .merge(build_mcp_router().await) + .merge(build_mcp_router(Some(telemetry.clone())).await) .merge(build_platform_mcp_router( Some(memory_binding_for_platform), Some(policy_status_for_platform), + Some(telemetry), )); let merged = if let Some(a2a) = a2a_router_opt { @@ -475,6 +479,9 @@ async fn main() -> Result<()> { .and_then(|s| s.parse::().ok()) .unwrap_or(256), )) + // Operator controls must remain reachable while inference requests + // or bounded approval waits occupy their own concurrency limits. + .merge(services) // r6 — trace-id middleware is outermost so every request gets a // trace span before any other layer runs (concurrency limit, // connection_close, auth gates all log inside the span). @@ -561,6 +568,7 @@ async fn main() -> Result<()> { let (signal_fired_tx, signal_fired_rx) = tokio::sync::oneshot::channel::<()>(); let shutdown_fut = async move { shutdown_signal().await; + services_shutdown.cancel(); let _ = signal_fired_tx.send(()); }; @@ -623,7 +631,9 @@ async fn main() -> Result<()> { /// falling back to the unauthenticated dev route. Operators see a clear /// startup-time error instead of a route that quietly serves /// unauthenticated MCP traffic. -async fn build_mcp_router() -> Router { +async fn build_mcp_router( + telemetry: Option>, +) -> Router { use kars_inference_router::mcp::forwarder::RouterToolDispatcher; use kars_inference_router::mcp::oauth::OAuthVerifierConfig; use kars_inference_router::mcp::registry; @@ -675,6 +685,7 @@ async fn build_mcp_router() -> Router { }; let mut state = routes::McpRouteState::standard(); + state.task_telemetry = telemetry; if let Some(d) = dispatcher_arc { state = state.with_tools(d); } @@ -778,8 +789,10 @@ fn build_platform_mcp_router( policy_status: Option< std::sync::Arc, >, + telemetry: Option>, ) -> Router { - let state = routes::McpRouteState::platform(memory_binding, policy_status); + let mut state = routes::McpRouteState::platform(memory_binding, policy_status); + state.task_telemetry = telemetry; tracing::info!( "Mounting /platform/mcp (Foundry-shim discovery surface, loopback-only, no OAuth)" ); diff --git a/inference-router/src/proxy.rs b/inference-router/src/proxy.rs index 0d40e7ed7..122ff9605 100644 --- a/inference-router/src/proxy.rs +++ b/inference-router/src/proxy.rs @@ -34,6 +34,7 @@ mod authentication_tests; /// Upstream configuration for a single request. #[derive(Clone)] pub struct UpstreamConfig { + pub telemetry: Option>, pub endpoint: String, pub deployment: String, pub sandbox_name: String, @@ -51,11 +52,18 @@ pub struct UpstreamConfig { } impl UpstreamConfig { + fn telemetry_provider(&self) -> &str { + match &self.authentication { + AuthenticationProvenance::Named { provider_id } => provider_id, + AuthenticationProvenance::LegacyDefault => self.provider.as_tag(), + } + } /// 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 { + telemetry: None, endpoint, deployment, sandbox_name, @@ -255,9 +263,22 @@ pub async fn forward( request_body: Bytes, ) -> Result<(StatusCode, HeaderMap, Bytes)> { let start = Instant::now(); + let mut observation = upstream.telemetry.as_ref().and_then(|telemetry| { + telemetry.begin( + path, + upstream.telemetry_provider(), + &upstream.deployment, + &request_body, + ) + }); - let (upstream_url, body) = build_upstream_url(auth, upstream, path, request_body) - .map_err(ForwardFailure::configuration)?; + let (upstream_url, body) = + build_upstream_url(auth, upstream, path, request_body).map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("configuration_error"); + } + ForwardFailure::configuration(error) + })?; let mode = match upstream.provider { ProviderKind::Anthropic => "anthropic", @@ -276,12 +297,22 @@ pub async fn forward( let credential = credential_for_upstream(auth, copilot, upstream) .await - .map_err(ForwardFailure::authentication)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("authentication_error"); + } + ForwardFailure::authentication(error) + })?; let headers = build_upstream_headers(request_headers, auth, &credential, &upstream.endpoint) - .map_err(ForwardFailure::configuration)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("configuration_error"); + } + ForwardFailure::configuration(error) + })?; - tracing::info!(sandbox = %upstream.sandbox_name, url = %upstream_url, body_len = body.len(), "Sending upstream request"); + tracing::info!(sandbox = %upstream.sandbox_name, body_len = body.len(), "Sending upstream request"); let retryable = is_idempotent(&method, path); let response = send_with_retry( @@ -293,11 +324,19 @@ pub async fn forward( retryable, &upstream.sandbox_name, ) - .await?; + .await + .inspect_err(|_| { + if let Some(observation) = observation.as_mut() { + observation.fail("transport_error"); + } + })?; let status = StatusCode::from_u16(response.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY); let response_headers = response.headers().clone(); + if let Some(observation) = observation.as_mut() { + observation.headers(status.as_u16()); + } // r6 — surface Azure-side request ids so one log line carries both our // trace_id (from the outer tracing span) and Azure's correlation ids. @@ -312,13 +351,18 @@ pub async fn forward( .and_then(|v| v.to_str().ok()) .unwrap_or(""); - let response_body = response - .bytes() - .await - .map_err(|error| ForwardFailure::response_body(status, error))?; + let response_body = response.bytes().await.map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("response_body_error"); + } + ForwardFailure::response_body(status, error) + })?; let latency = start.elapsed(); record_metrics(upstream, status, latency, &response_body); + if let Some(observation) = observation.as_mut() { + observation.buffered(status.as_u16(), &response_body); + } tracing::info!( sandbox = %upstream.sandbox_name, @@ -485,6 +529,14 @@ pub async fn forward_stream( HeaderMap, futures::stream::BoxStream<'static, Result>, )> { + let mut observation = upstream.telemetry.as_ref().and_then(|telemetry| { + telemetry.begin( + path, + upstream.telemetry_provider(), + &upstream.deployment, + &request_body, + ) + }); // Inject stream_options.include_usage into request body so the final // SSE chunk contains a `usage` object with token counts. ONLY for // chat/completions — Anthropic Messages API (/v1/messages) rejects @@ -501,15 +553,30 @@ pub async fn forward_stream( inject_stream_usage(request_body) }; let (upstream_url, body) = build_upstream_url(&auth, &upstream, path, body_with_usage) - .map_err(ForwardFailure::configuration)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("configuration_error"); + } + ForwardFailure::configuration(error) + })?; tracing::info!(sandbox = %upstream.sandbox_name, model = %upstream.deployment, mode = "stream", "Forwarding SSE stream"); let credential = credential_for_upstream(&auth, copilot.as_deref(), &upstream) .await - .map_err(ForwardFailure::authentication)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("authentication_error"); + } + ForwardFailure::authentication(error) + })?; let headers = build_upstream_headers(&request_headers, &auth, &credential, &upstream.endpoint) - .map_err(ForwardFailure::configuration)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("configuration_error"); + } + ForwardFailure::configuration(error) + })?; let start = Instant::now(); @@ -520,11 +587,19 @@ pub async fn forward_stream( .timeout(INFERENCE_REQUEST_TIMEOUT) .send() .await - .map_err(ForwardFailure::transport)?; + .map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("transport_error"); + } + ForwardFailure::transport(error) + })?; let status = StatusCode::from_u16(response.status().as_u16()).unwrap_or(StatusCode::BAD_GATEWAY); let response_headers = response.headers().clone(); + if let Some(observation) = observation.as_mut() { + observation.headers(status.as_u16()); + } // r6 — log Azure correlation ids for the stream path too. Emitted at // stream-start because headers arrive before any bytes. @@ -555,22 +630,24 @@ pub async fn forward_stream( .inc(); // On non-success, the upstream body is a short JSON error (not an SSE - // stream). Eagerly drain it, log the contents (capped), and forward as a + // stream). Eagerly drain it and forward as a // single chunk so callers see the actual reason. Without this we only // see "status=413" and have to guess at causes (token cap? bytes cap? - // schema break?). Cap at 4 KiB so a misbehaving upstream can't blow logs. + // schema break?). Do not put provider API bodies into logs. if !status.is_success() { - let body_bytes = response - .bytes() - .await - .map_err(|error| ForwardFailure::response_body(status, error))?; - let preview = String::from_utf8_lossy(&body_bytes); - let preview_trimmed: String = preview.chars().take(2048).collect(); + let body_bytes = response.bytes().await.map_err(|error| { + if let Some(observation) = observation.as_mut() { + observation.fail("response_body_error"); + } + ForwardFailure::response_body(status, error) + })?; + if let Some(observation) = observation.as_mut() { + observation.buffered(status.as_u16(), &body_bytes); + } tracing::warn!( sandbox = %upstream.sandbox_name, status = %status.as_u16(), body_len = body_bytes.len(), - body = %preview_trimmed, "Upstream returned non-success status" ); let stream = futures::stream::once(async move { Ok::<_, reqwest::Error>(body_bytes) }); @@ -632,7 +709,15 @@ pub async fn forward_stream( chunk }); - Ok((status, response_headers, metered.boxed())) + let is_sse = response_headers + .get("content-type") + .and_then(|header| header.to_str().ok()) + .is_some_and(|value| value.starts_with("text/event-stream")); + Ok(( + status, + response_headers, + crate::task_telemetry::observe::wrap_stream(metered.boxed(), observation, is_sse), + )) } /// Inject `stream_options: { include_usage: true }` into the request body diff --git a/inference-router/src/routes/access_request.rs b/inference-router/src/routes/access_request.rs new file mode 100644 index 000000000..e15de9197 --- /dev/null +++ b/inference-router/src/routes/access_request.rs @@ -0,0 +1,269 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::AppState; +use crate::access_request::{Error, Request, Status}; +use axum::{ + Json, Router, + extract::{ConnectInfo, DefaultBodyLimit, Path, Query, State}, + http::{HeaderMap, StatusCode}, + middleware::{self, Next}, + response::{IntoResponse, Response}, + routing::{get, post}, +}; +use serde::Deserialize; +use serde_json::json; +use std::{net::SocketAddr, time::Duration}; + +pub fn routes(state: AppState) -> Router { + let agent = Router::new() + .route("/v1/access-request", post(create)) + .route("/v1/access-requests", get(list_agent)) + .route("/v1/access-requests/{id}/cancel", post(cancel)) + .route("/v1/access-requests/{id}/wait", get(wait)) + .merge(super::task_telemetry::routes()) + .route_layer(middleware::from_fn_with_state(state.clone(), agent_auth)) + .layer(tower::limit::ConcurrencyLimitLayer::new(24)); + let admin = Router::new() + .route("/internal/access-requests", get(inspect)) + .route("/internal/access-requests/reset", post(reset)) + .route("/internal/access-requests/decision", post(decide)) + .route_layer(middleware::from_fn_with_state(state, control_auth)) + .layer(tower::limit::ConcurrencyLimitLayer::new(8)); + agent.merge(admin).layer(DefaultBodyLimit::max(8192)) +} + +fn bearer(headers: &HeaderMap) -> Option<&str> { + headers + .get("authorization") + .and_then(|header| header.to_str().ok()) + .and_then(|header| header.strip_prefix("Bearer ")) +} + +async fn control_auth( + State(state): State, + request: axum::extract::Request, + next: Next, +) -> Response { + if !state.services.identity_valid { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "Service identity is unavailable", + ) + .into_response(); + } + if !state.services.control_configured() { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "Service control credential is not configured", + ) + .into_response(); + } + if !state.services.control_authorized(bearer(request.headers())) { + return ( + StatusCode::UNAUTHORIZED, + "Service control credential required; agent admin credentials are not accepted", + ) + .into_response(); + } + if let Some(allowed) = &state.services.allow_ips { + let remote = request + .extensions() + .get::>() + .map(|info| info.0.ip()); + if remote.is_none_or(|remote| !allowed.contains(&remote)) { + return (StatusCode::FORBIDDEN, "Control origin is not allowed").into_response(); + } + } + next.run(request).await +} + +async fn agent_auth( + State(state): State, + request: axum::extract::Request, + next: Next, +) -> Response { + if !state.services.identity_valid { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "Service identity is unavailable", + ) + .into_response(); + } + let local = request + .extensions() + .get::>() + .is_some_and(|info| info.0.ip().is_loopback()); + if !local && !state.services.control_authorized(bearer(request.headers())) { + return ( + StatusCode::UNAUTHORIZED, + "Same-router agent or service control authentication required", + ) + .into_response(); + } + if !local && let Some(allowed) = &state.services.allow_ips { + let remote = request + .extensions() + .get::>() + .map(|info| info.0.ip()); + if remote.is_none_or(|ip| !allowed.contains(&ip)) { + return (StatusCode::FORBIDDEN, "Control origin is not allowed").into_response(); + } + } + next.run(request).await +} + +pub(super) fn error(error: Error) -> Response { + let (status, code) = match error { + Error::Invalid => (StatusCode::BAD_REQUEST, "invalid_request"), + Error::StaleScope => (StatusCode::CONFLICT, "stale_scope"), + Error::Full => (StatusCode::TOO_MANY_REQUESTS, "capacity_exhausted"), + Error::RateLimited => (StatusCode::TOO_MANY_REQUESTS, "rate_limited"), + Error::Missing => (StatusCode::NOT_FOUND, "request_not_found"), + Error::Expired => (StatusCode::GONE, "request_expired"), + Error::WaitTimeout => (StatusCode::REQUEST_TIMEOUT, "wait_timed_out"), + Error::Terminal => (StatusCode::CONFLICT, "request_terminal"), + Error::Unavailable => (StatusCode::SERVICE_UNAVAILABLE, "service_unavailable"), + }; + (status, Json(json!({"error":code}))).into_response() +} +pub(super) fn scope_header(headers: &HeaderMap) -> Result<&str, Error> { + headers + .get("x-kars-service-scope") + .and_then(|header| header.to_str().ok()) + .filter(|id| id.len() <= 128) + .ok_or(Error::StaleScope) +} + +async fn create(State(state): State, Json(request): Json) -> Response { + match state.services.requests.record(request) { + Ok((entry, new)) => ( + StatusCode::ACCEPTED, + Json(json!({"status":"queued","new":new, + "kind":entry.kind,"target":entry.target,"request":entry,"enforcement_changed":false})), + ) + .into_response(), + Err(err) => error(err), + } +} +async fn list_agent(State(state): State, headers: HeaderMap) -> Response { + let scope = match state.services.requests.scope() { + Ok(scope) => scope, + Err(err) => return error(err), + }; + if let Some(requested) = headers.get("x-kars-service-scope") + && requested.to_str().ok() != Some(scope.id.as_str()) + { + return error(Error::StaleScope); + } + match state.services.requests.snapshot(&scope.id) { + Ok(entries) => Json(json!({"scope":scope,"requests":entries,"enforcement_changed":false})) + .into_response(), + Err(err) => error(err), + } +} +async fn inspect(State(state): State, headers: HeaderMap) -> Response { + let scope = match state.services.requests.scope() { + Ok(scope) => scope, + Err(err) => return error(err), + }; + if let Some(requested) = headers.get("x-kars-service-scope") + && requested.to_str().ok() != Some(scope.id.as_str()) + { + return error(Error::StaleScope); + } + match state.services.requests.snapshot(&scope.id) { + Ok(entries) => Json( + json!({"schema_version":1,"scope":scope,"sandbox":state.sandbox_name.as_str(), + "count":entries.len(),"entries":entries,"enforcement_changed":false}), + ) + .into_response(), + Err(err) => error(err), + } +} + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct ScopeBody { + scope_id: String, +} +async fn cancel( + State(state): State, + Path(id): Path, + Json(body): Json, +) -> Response { + match state + .services + .requests + .transition(&body.scope_id, &id, Status::Cancelled) + { + Ok(entry) => Json(entry).into_response(), + Err(err) => error(err), + } +} +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct Reset { + scope_id: String, + assignment_id: Option, +} +async fn reset(State(state): State, Json(body): Json) -> Response { + match state.services.reset(&body.scope_id, body.assignment_id) { + Ok((scope, cleared)) => { + Json(json!({"scope":scope,"cleared":cleared,"enforcement_changed":false})) + .into_response() + } + Err(err) => error(err), + } +} +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct Decision { + scope_id: String, + request_id: String, + verdict: Status, +} +async fn decide(State(state): State, Json(body): Json) -> Response { + if !matches!(body.verdict, Status::Approved | Status::Denied) { + return error(Error::Invalid); + } + match state + .services + .requests + .transition(&body.scope_id, &body.request_id, body.verdict) + { + Ok(entry) => Json(json!({"request":entry,"enforcement_changed":false})).into_response(), + Err(err) => error(err), + } +} +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct Wait { + scope_id: String, + #[serde(default = "default_wait")] + timeout_ms: u64, +} +fn default_wait() -> u64 { + 30_000 +} +async fn wait( + State(state): State, + Path(id): Path, + Query(query): Query, +) -> Response { + if query.timeout_ms == 0 || query.timeout_ms > 300_000 { + return error(Error::Invalid); + } + match state + .services + .wait_for_decision( + &query.scope_id, + &id, + Duration::from_millis(query.timeout_ms), + ) + .await + { + Ok(entry) => Json(entry).into_response(), + Err(err) => error(err), + } +} diff --git a/inference-router/src/routes/egress.rs b/inference-router/src/routes/egress.rs index cf5964ca5..3178d6835 100644 --- a/inference-router/src/routes/egress.rs +++ b/inference-router/src/routes/egress.rs @@ -113,10 +113,30 @@ async fn egress_fetch( State(state): State, Json(req): Json, ) -> impl IntoResponse { + let started = std::time::Instant::now(); + let telemetry_scope = state.services.telemetry.cursor().0; let url = req.get("url").and_then(|v| v.as_str()).unwrap_or(""); let method = req.get("method").and_then(|v| v.as_str()).unwrap_or("GET"); let req_body = req.get("body").and_then(|v| v.as_str()).unwrap_or(""); let req_headers = req.get("headers").and_then(|v| v.as_object()); + let wait_ms = req + .get("wait_for_approval_ms") + .and_then(serde_json::Value::as_u64) + .unwrap_or(0); + if req + .get("wait_for_approval_ms") + .is_some_and(|value| value.as_u64().is_none()) + { + return super::access_request::error(crate::access_request::Error::Invalid); + } + if wait_ms > 300_000 { + return super::access_request::error(crate::access_request::Error::Invalid); + } + if wait_ms > 0 + && req.get("scope_id").and_then(serde_json::Value::as_str) != Some(telemetry_scope.as_str()) + { + return super::access_request::error(crate::access_request::Error::StaleScope); + } if url.is_empty() { return errors::flat(StatusCode::BAD_REQUEST, "Missing 'url' field").into_response(); @@ -134,7 +154,15 @@ async fn egress_fetch( } }; if is_private { - tracing::warn!(url = %url, "Egress fetch blocked: private/internal target"); + tracing::warn!("Egress fetch blocked: private/internal target"); + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + Some(false), + Some(403), + 0, + true, + ); return ( StatusCode::FORBIDDEN, Json(serde_json::json!({ @@ -151,18 +179,88 @@ async fn egress_fetch( // Check egress access: blocklist → allowlist (Strict denies the rest). if let Err(reason) = state.blocklist.check_egress(url, sandbox).await { - tracing::warn!(url = %url, reason = %reason, "Egress fetch denied"); - return (StatusCode::FORBIDDEN, Json(serde_json::json!({ + tracing::warn!("Egress fetch denied by policy"); + let target = reqwest::Url::parse(url).ok().and_then(|url| { + url.host_str() + .zip(url.port_or_known_default()) + .map(|(host, port)| (host.to_string(), port)) + }); + let entry = target.as_ref().and_then(|(host, port)| { + let now = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_secs(); + state + .blocked_egress + .record_at(std::time::Instant::now(), now, sandbox, host, *port); + state + .services + .record_egress_in_scope(&telemetry_scope, host, *port) + .ok() + }); + state + .services + .telemetry + .record_policy(&telemetry_scope, "egress", false); + if wait_ms > 0 { + let Some(entry) = entry else { + return super::access_request::error(crate::access_request::Error::Full); + }; + if let Err(error) = state + .services + .wait_for_egress( + &state.blocklist, + &telemetry_scope, + &entry.request_id, + url, + sandbox, + std::time::Duration::from_millis(wait_ms), + ) + .await + { + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + Some(false), + Some(403), + started.elapsed().as_millis() as u64, + false, + ); + return super::access_request::error(error); + } + } else { + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + Some(false), + Some(403), + 0, + true, + ); + return (StatusCode::FORBIDDEN, Json(serde_json::json!({ "error": reason, "url": url, "action": "If legitimate, an operator can add the host to the baseline allowlist ('kars egress --approve ', which re-signs) or grant it temporarily ('kars egress allow-extra --host --ttl --reason ').", - }))).into_response(); + }))).into_response(); + } + } + if wait_ms > 0 + && state + .services + .requests + .scope() + .map(|scope| scope.id) + .ok() + .as_deref() + != Some(telemetry_scope.as_str()) + { + return super::access_request::error(crate::access_request::Error::StaleScope); } // Record in learn mode state.blocklist.record_learned(url).await; - tracing::info!(url = %url, method = %method, "Egress fetch proxied"); + tracing::info!("Egress fetch proxied"); // Build and send the request let http_method = match method.to_uppercase().as_str() { @@ -236,32 +334,79 @@ async fn egress_fetch( }) .collect(); // Cap response body to prevent OOM - let body_bytes = resp.bytes().await.unwrap_or_default(); - let body = if body_bytes.len() > MAX_RESPONSE_BYTES { - let truncated = String::from_utf8_lossy(&body_bytes[..MAX_RESPONSE_BYTES]); - format!( - "{}... [truncated at {} bytes]", - truncated, MAX_RESPONSE_BYTES - ) - } else { - String::from_utf8_lossy(&body_bytes).into_owned() - }; + use futures::StreamExt; + let mut stream = resp.bytes_stream(); + let mut body_bytes = Vec::new(); + let mut truncated = false; + while let Some(chunk) = stream.next().await { + let chunk = match chunk { + Ok(chunk) => chunk, + Err(_) => { + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + Some(false), + Some(status), + started.elapsed().as_millis() as u64, + false, + ); + return errors::flat( + StatusCode::BAD_GATEWAY, + "Upstream response body failed", + ) + .into_response(); + } + }; + let available = MAX_RESPONSE_BYTES - body_bytes.len(); + body_bytes.extend_from_slice(&chunk[..chunk.len().min(available)]); + if chunk.len() > available { + truncated = true; + break; + } + } + let mut body = String::from_utf8_lossy(&body_bytes).into_owned(); + if truncated { + body.push_str(&format!("... [truncated at {MAX_RESPONSE_BYTES} bytes]")); + } + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + if status >= 400 { + Some(false) + } else if truncated { + None + } else { + Some(true) + }, + Some(status), + started.elapsed().as_millis() as u64, + !truncated, + ); ( StatusCode::OK, Json(serde_json::json!({ "status": status, "headers": resp_headers, "body": body, + "truncated": truncated, })), ) .into_response() } - Err(e) => { - tracing::warn!(url = %url, error = %e, "Egress fetch failed"); + Err(_) => { + tracing::warn!("Egress fetch failed"); + state.services.telemetry.record_router_tool( + &telemetry_scope, + "http_fetch", + Some(false), + None, + started.elapsed().as_millis() as u64, + false, + ); ( StatusCode::BAD_GATEWAY, Json(serde_json::json!({ - "error": format!("Request failed: {}", e), + "error": "Upstream request failed", "url": url, })), ) diff --git a/inference-router/src/routes/governance.rs b/inference-router/src/routes/governance.rs index 258fea91e..20e6db04d 100644 --- a/inference-router/src/routes/governance.rs +++ b/inference-router/src/routes/governance.rs @@ -60,12 +60,17 @@ async fn agt_evaluate( .and_then(|v| v.as_str()) .unwrap_or("unknown"); let extra_context = body.get("context").cloned(); + let telemetry_scope = state.services.telemetry.cursor().0; // Per-tool sliding window rate limit (AGT McpSlidingRateLimiter) // Extract tool name from action format "tool:exec_command" or "tool:http_fetch" if let Some(tool_name) = action.strip_prefix("tool:") { let (allowed, retry_after) = state.governance.check_tool_rate(tool_name); if !allowed { + state + .services + .telemetry + .record_policy(&telemetry_scope, action, false); return Json(serde_json::json!({ "allowed": false, "reason": format!("per-tool rate limit exceeded for '{}'", tool_name), @@ -79,6 +84,41 @@ async fn agt_evaluate( let result = state .governance .evaluate(agent_id, action, extra_context.as_ref()); + if agent_id == state.sandbox_name.as_str() + && let Some(tool) = action.strip_prefix("tool:") + { + let allowed = result.get("allowed").and_then(serde_json::Value::as_bool) == Some(true); + state + .services + .telemetry + .record_policy(&telemetry_scope, action, allowed); + if allowed { + if let Some(context) = extra_context.as_ref() + && context.get("scope_id").and_then(serde_json::Value::as_str) + == Some(telemetry_scope.as_str()) + && let Some(id) = context + .get("tool_call_id") + .and_then(serde_json::Value::as_str) + { + state + .services + .telemetry + .authorize_harness_tool(&telemetry_scope, id, tool); + } + } else { + let _ = state + .services + .requests + .record(crate::access_request::Request { + scope_id: telemetry_scope, + kind: "tool".into(), + target: tool.into(), + reason: "Governance denied the requested tool".into(), + tier: None, + port: None, + }); + } + } Json(result).into_response() } diff --git a/inference-router/src/routes/mcp.rs b/inference-router/src/routes/mcp.rs index 7be0798ce..7066d1230 100644 --- a/inference-router/src/routes/mcp.rs +++ b/inference-router/src/routes/mcp.rs @@ -62,6 +62,7 @@ pub const MCP_SESSION_HEADER: &str = "Mcp-Session-Id"; /// Per-router MCP state. Cheap to clone (everything inside is `Arc`). #[derive(Clone)] pub struct McpRouteState { + pub task_telemetry: Option>, pub config: Arc, pub minter: Arc, pub tools: Arc, @@ -76,6 +77,7 @@ impl McpRouteState { /// `McpServer.spec.url`; see [`with_tools`]. pub fn standard() -> Self { Self { + task_telemetry: None, config: Arc::new(InitializeConfig::default()), minter: Arc::new(OsRngSessionMinter), tools: Arc::new(SyncToAsync::new(EchoDispatcher::standard())), @@ -123,6 +125,7 @@ impl McpRouteState { config: Arc::new(InitializeConfig::default()), minter: Arc::new(OsRngSessionMinter), tools: Arc::new(dispatcher), + task_telemetry: None, } } } @@ -188,6 +191,10 @@ async fn method_not_allowed() -> impl IntoResponse { } async fn post_mcp(State(state): State, headers: HeaderMap, body: Bytes) -> Response { + let telemetry_scope = state + .task_telemetry + .as_ref() + .map(|telemetry| telemetry.cursor().0); let accept = headers .get(header::ACCEPT) .and_then(|v| v.to_str().ok()) @@ -204,6 +211,15 @@ async fn post_mcp(State(state): State, headers: HeaderMap, body: Some(state.tools.as_ref()), ) .await; + if let (Some(telemetry), Some(scope)) = (&state.task_telemetry, telemetry_scope) { + crate::task_telemetry::mcp::record( + telemetry, + &scope, + &body, + &outcome, + started.elapsed().as_millis() as u64, + ); + } let status_label: &str = match &outcome { ProcessOutcome::JsonRpcResponse { .. } => "200", @@ -304,6 +320,7 @@ mod tests { fn test_state() -> McpRouteState { McpRouteState { + task_telemetry: None, config: Arc::new(InitializeConfig::default()), minter: Arc::new(FixedMinter("test-session-001")), tools: Arc::new(SyncToAsync::new(EchoDispatcher::standard())), @@ -493,6 +510,7 @@ mod tests { McpRouteState { config: Arc::new(InitializeConfig::default()), minter: Arc::new(FixedMinter("platform-session-001")), + task_telemetry: None, tools: Arc::new(crate::mcp::PlatformDispatcher::with_base_url( "http://127.0.0.1:1", )), @@ -679,6 +697,7 @@ mod tests { // Real platform dispatcher, pointed at the mock upstream. let state = McpRouteState { + task_telemetry: None, config: Arc::new(InitializeConfig::default()), minter: Arc::new(FixedMinter("platform-session-001")), tools: Arc::new(crate::mcp::PlatformDispatcher::with_base_url( diff --git a/inference-router/src/routes/mod.rs b/inference-router/src/routes/mod.rs index ff8d134a6..a2eb5a6ea 100644 --- a/inference-router/src/routes/mod.rs +++ b/inference-router/src/routes/mod.rs @@ -44,7 +44,10 @@ pub use governance::sensitive_agt_routes; mod mesh; pub use mesh::mesh_routes; +mod access_request; mod mesh_token; +mod task_telemetry; +pub use access_request::routes as governed_service_routes; mod model_routing; pub use mesh_token::mesh_token_routes; @@ -71,6 +74,7 @@ pub use a2a::{A2aRouteState, a2a_routes}; /// Shared application state. #[derive(Clone)] pub struct AppState { + pub services: Arc, pub auth: Arc, /// GitHub Copilot token cache. Constructed unconditionally (cheap, lazy) /// — `proxy::forward()` only consults it when the upstream endpoint is @@ -310,7 +314,13 @@ impl AppState { ) .await; + let services = Arc::new(crate::governed_services::GovernedServices::from_env( + &sandbox_name, + )); + let blocked_egress = Arc::new(BlockedBuffer::with_defaults()); + blocked_egress.bind_services(&services); Ok(Self { + services, auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(CopilotTokenCache::from_env()), client: client.clone(), @@ -321,7 +331,7 @@ impl AppState { signing_provider: Arc::clone(&governance) as Arc, governance, blocklist, - blocked_egress: Arc::new(BlockedBuffer::with_defaults()), + blocked_egress, sandbox_name: Arc::new(sandbox_name), inbox: Arc::new(MeshInbox::new()), mesh_metrics: Arc::new(MeshMetrics::new()), @@ -367,7 +377,11 @@ impl AppState { .and_then(|g| g.clone()) .unwrap_or_else(|| self.config.default_model.clone()); - UpstreamConfig::azure(endpoint, deployment, sandbox_name.to_string()) + let mut upstream = UpstreamConfig::azure(endpoint, deployment, sandbox_name.to_string()); + if self.services.identity_valid { + upstream.telemetry = Some(self.services.telemetry.clone()); + } + upstream } } diff --git a/inference-router/src/routes/model_routing.rs b/inference-router/src/routes/model_routing.rs index 61aa5d1ca..9ecd35dc0 100644 --- a/inference-router/src/routes/model_routing.rs +++ b/inference-router/src/routes/model_routing.rs @@ -376,6 +376,7 @@ mod tests { policy_status.clone(), )); AppState { + services: Default::default(), auth: Arc::new(crate::auth::WorkloadIdentityAuth::new()), copilot: Arc::new(crate::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/src/routes/task_telemetry.rs b/inference-router/src/routes/task_telemetry.rs new file mode 100644 index 000000000..999f81646 --- /dev/null +++ b/inference-router/src/routes/task_telemetry.rs @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + AppState, + access_request::{error, scope_header}, +}; +use crate::access_request::Error; +use axum::{ + Json, Router, + extract::{Query, State}, + http::HeaderMap, + response::{IntoResponse, Response}, + routing::{get, post}, +}; +use serde::Deserialize; +use serde_json::json; + +pub(super) fn routes() -> Router { + Router::new() + .route("/telemetry/cursor", get(cursor)) + .route("/telemetry/trace", get(trace)) + .route("/telemetry/tool", post(tool)) + .route("/telemetry/budget", get(budget)) +} +async fn cursor(State(state): State) -> Response { + let (scope_id, cursor) = state.services.telemetry.cursor(); + Json(json!({"scope_id":scope_id,"cursor":cursor,"durable":false})).into_response() +} +#[derive(Deserialize)] +struct Since { + #[serde(default)] + since: u64, +} +async fn trace( + State(state): State, + headers: HeaderMap, + Query(query): Query, +) -> Response { + let scope = match scope_header(&headers) { + Ok(scope) => scope, + Err(err) => return error(err), + }; + match state.services.telemetry.snapshot(scope, query.since) { + Some(trace) => Json(trace).into_response(), + None => error(Error::StaleScope), + } +} +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct Tool { + scope_id: String, + call_id: String, + ok: bool, + latency_ms: u64, +} +async fn tool(State(state): State, Json(body): Json) -> Response { + if state.services.telemetry.complete_harness_tool( + &body.scope_id, + &body.call_id, + body.ok, + body.latency_ms, + ) { + Json(json!({"recorded":true,"source":"harness-reported","enforcement_changed":false})) + .into_response() + } else { + error(Error::Terminal) + } +} +async fn budget(State(state): State, headers: HeaderMap) -> Response { + let scope = match scope_header(&headers) { + Ok(scope) => scope, + Err(err) => return error(err), + }; + if state.services.telemetry.cursor().0 != scope { + return error(Error::StaleScope); + } + let (daily, monthly) = state.budget.get_usage(&state.sandbox_name).await; + let policy = crate::inference_policy_loader::current_snapshot(&state.inference_policy).await; + let daily_limit = policy + .daily_tokens + .unwrap_or(state.config.token_budget_daily); + Json( + json!({"scope_id":scope,"daily_observed_tokens":daily,"monthly_observed_tokens":monthly, + "daily_limit":(daily_limit>0).then_some(daily_limit), + "monthly_limit":policy.monthly_tokens.filter(|limit|*limit>0), + "shared_task_budget_enforced":false,"coverage":"existing-per-sandbox-counter-only", + "assignment_reset_resets_budget":false}), + ) + .into_response() +} diff --git a/inference-router/src/task_telemetry.rs b/inference-router/src/task_telemetry.rs new file mode 100644 index 000000000..ff5779529 --- /dev/null +++ b/inference-router/src/task_telemetry.rs @@ -0,0 +1,259 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +//! Bounded router observations, not a durable execution or budget ledger. +//! No prompts, tool arguments/results, URLs, headers or credentials are retained. + +use serde_json::{Value, json}; +use std::{ + collections::{HashMap, HashSet, VecDeque}, + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +pub mod mcp; +pub mod observe; +pub mod parse; + +const MAX_EVENTS: usize = 1024; +const MAX_PENDING: usize = 256; +const TOOL_TTL: Duration = Duration::from_secs(20 * 60); + +struct Pending { + name: String, + at: Instant, + round: u64, +} +struct Inner { + scope: String, + seq: u64, + round: u64, + dropped: u64, + events: VecDeque, + model_tools: HashMap, + seen_model: HashSet, + harness_tools: HashMap, + seen_harness: HashSet, +} + +pub struct TaskTelemetry { + inner: Mutex, +} + +impl TaskTelemetry { + pub fn new(scope: String) -> Self { + Self { + inner: Mutex::new(Inner { + scope, + seq: 0, + round: 0, + dropped: 0, + events: VecDeque::new(), + model_tools: HashMap::new(), + seen_model: HashSet::new(), + harness_tools: HashMap::new(), + seen_harness: HashSet::new(), + }), + } + } + pub fn reset(&self, scope: String) { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + inner.scope = scope; + inner.events.clear(); + inner.model_tools.clear(); + inner.seen_model.clear(); + inner.harness_tools.clear(); + inner.seen_harness.clear(); + inner.round = 0; + inner.dropped = 0; + } + fn push(inner: &mut Inner, mut event: Value) { + inner.seq = inner.seq.saturating_add(1); + event["seq"] = json!(inner.seq); + event["scope_id"] = json!(inner.scope); + inner.events.push_back(event); + if inner.events.len() > MAX_EVENTS { + inner.events.pop_front(); + inner.dropped = inner.dropped.saturating_add(1); + } + } + pub fn cursor(&self) -> (String, u64) { + let inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + (inner.scope.clone(), inner.seq) + } + pub fn snapshot(&self, scope: &str, since: u64) -> Option { + let inner = self.inner.lock().ok()?; + if inner.scope != scope { + return None; + } + let events = inner + .events + .iter() + .filter(|event| event["seq"].as_u64().is_some_and(|seq| seq > since)) + .take(256) + .cloned() + .collect::>(); + let cursor = events + .last() + .and_then(|event| event["seq"].as_u64()) + .unwrap_or(inner.seq); + Some(json!({ + "scope_id": scope, "cursor": cursor, "high_water": inner.seq, + "has_more": cursor, + path: &str, + provider: &str, + model: &str, + body: &[u8], + ) -> Option { + let shape = parse::Shape::for_path(path)?; + let mut inner = self.inner.lock().ok()?; + inner + .model_tools + .retain(|_, pending| pending.at.elapsed() < TOOL_TTL); + for (id, ok) in parse::request_results(body, shape) { + if let Some(pending) = inner.model_tools.remove(&id) { + Self::push( + &mut inner, + json!({"kind":"tool_result", "call_id":id, "name":pending.name, + "round":pending.round, "ok":ok, "source":"harness-reported"}), + ); + } + } + inner.round = inner.round.saturating_add(1); + Some(observe::Observation::new( + self.clone(), + inner.scope.clone(), + inner.round, + shape, + parse::identifier(provider, 64), + parse::identifier(model, 253), + )) + } + pub(crate) fn finish( + &self, + scope: &str, + round: u64, + event: Value, + tools: Vec<(String, String)>, + ) { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + if inner.scope != scope { + return; + } + Self::push(&mut inner, event); + for (id, name) in tools { + Self::push( + &mut inner, + json!({"kind":"tool_proposed", "round":round, "call_id":id, + "name":name, "ok":null, "source":"model-proposed"}), + ); + if !id.is_empty() + && inner.seen_model.len() < MAX_PENDING + && inner.seen_model.insert(id.clone()) + { + inner.model_tools.insert( + id, + Pending { + name, + at: Instant::now(), + round, + }, + ); + } else { + inner.model_tools.remove(&id); + inner.dropped = inner.dropped.saturating_add(1); + } + } + } + pub fn record_policy(&self, scope: &str, capability: &str, allowed: bool) { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + if inner.scope != scope { + return; + } + Self::push( + &mut inner, + json!({"kind":"policy", "capability":parse::identifier(capability,128), + "allowed":allowed, "source":"router-policy"}), + ); + } + pub fn record_gap(&self, scope: &str, source: &str) { + let mut inner = self.inner.lock().unwrap_or_else(|error| error.into_inner()); + if inner.scope != scope { + return; + } + inner.dropped = inner.dropped.saturating_add(1); + Self::push( + &mut inner, + json!({"kind":"observation_gap","source":parse::identifier(source,128)}), + ); + } + pub fn record_router_tool( + &self, + scope: &str, + name: &str, + ok: Option, + status: Option, + latency_ms: u64, + complete: bool, + ) { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + if inner.scope != scope { + return; + } + Self::push( + &mut inner, + json!({"kind":"tool_result", "name":parse::identifier(name,128), + "ok":ok, "http_status":status, "ms":latency_ms, "complete":complete, "source":"router"}), + ); + } + pub fn authorize_harness_tool(&self, scope: &str, id: &str, name: &str) -> bool { + if !crate::access_request::scope::identifier(id, 128) + || parse::identifier(name, 128).is_none() + { + return false; + } + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + if inner.scope != scope { + return false; + } + if inner.seen_harness.contains(id) || inner.seen_harness.len() >= MAX_PENDING { + inner.dropped = inner.dropped.saturating_add(1); + return false; + } + inner.seen_harness.insert(id.to_string()); + let round = inner.round; + inner.harness_tools.insert( + id.to_string(), + Pending { + name: name.to_string(), + at: Instant::now(), + round, + }, + ); + true + } + pub fn complete_harness_tool(&self, scope: &str, id: &str, ok: bool, latency_ms: u64) -> bool { + let mut inner = self.inner.lock().unwrap_or_else(|e| e.into_inner()); + if inner.scope != scope { + return false; + } + let Some(pending) = inner.harness_tools.remove(id) else { + return false; + }; + if pending.at.elapsed() >= TOOL_TTL { + return false; + } + Self::push( + &mut inner, + json!({"kind":"tool_result", "call_id":id, "name":pending.name, + "round":pending.round, "ok":ok, "ms":latency_ms.min(3_600_000), "source":"harness-reported"}), + ); + true + } +} diff --git a/inference-router/src/task_telemetry/mcp.rs b/inference-router/src/task_telemetry/mcp.rs new file mode 100644 index 000000000..f07250ff9 --- /dev/null +++ b/inference-router/src/task_telemetry/mcp.rs @@ -0,0 +1,79 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{TaskTelemetry, parse}; +use crate::mcp::pipeline::ProcessOutcome; +use serde_json::Value; + +pub fn record( + telemetry: &TaskTelemetry, + scope: &str, + request: &[u8], + outcome: &ProcessOutcome, + latency_ms: u64, +) { + if request.len() > parse::MAX_BODY { + telemetry.record_gap(scope, "mcp_request_metadata"); + return; + } + let Ok(request) = serde_json::from_slice::(request) else { + return; + }; + let (status, reply) = match outcome { + ProcessOutcome::JsonRpcResponse { body, .. } => { + if body.len() > parse::MAX_BODY { + telemetry.record_gap(scope, "mcp_response_metadata"); + (200, None) + } else { + (200, serde_json::from_slice::(body).ok()) + } + } + ProcessOutcome::Accepted => (202, None), + ProcessOutcome::PayloadTooLarge => (413, None), + ProcessOutcome::NotAcceptable(_) => (406, None), + }; + let requests = match &request { + Value::Array(values) => values.as_slice(), + value => std::slice::from_ref(value), + }; + if requests.len() > parse::MAX_TOOLS { + telemetry.record_gap(scope, "mcp_batch_metadata"); + } + for request in requests.iter().take(parse::MAX_TOOLS) { + if request["method"] != "tools/call" { + continue; + } + let Some(name) = request["params"]["name"] + .as_str() + .and_then(|name| parse::identifier(name, 128)) + else { + continue; + }; + let id = request.get("id"); + let response = reply.as_ref().and_then(|reply| { + let replies = match reply { + Value::Array(values) => values.as_slice(), + value => std::slice::from_ref(value), + }; + id.filter(|id| id.is_number() || id.as_str().is_some_and(|id| id.len() <= 128)) + .and_then(|id| replies.iter().find(|reply| reply.get("id") == Some(id))) + }); + let ok = response.and_then(|reply| { + if reply.get("error").is_some() { + return Some(false); + } + let result = reply.get("result")?; + result + .get("isError") + .and_then(Value::as_bool) + .map(|error| !error) + .or_else(|| { + result + .get("content") + .and_then(Value::as_array) + .map(|_| true) + }) + }); + telemetry.record_router_tool(scope, &name, ok, Some(status), latency_ms, ok.is_some()); + } +} diff --git a/inference-router/src/task_telemetry/observe.rs b/inference-router/src/task_telemetry/observe.rs new file mode 100644 index 000000000..2cb5dbaef --- /dev/null +++ b/inference-router/src/task_telemetry/observe.rs @@ -0,0 +1,182 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::{ + TaskTelemetry, + parse::{self, Parsed, Shape}, +}; +use bytes::Bytes; +use futures::{Stream, stream::BoxStream}; +use serde_json::json; +use std::{ + pin::Pin, + sync::Arc, + task::{Context, Poll}, + time::Instant, +}; + +mod stream; + +pub struct Observation { + telemetry: Arc, + scope: String, + round: u64, + shape: Shape, + provider: Option, + model: Option, + started: Instant, + status: Option, + finished: bool, +} +impl Observation { + pub(super) fn new( + telemetry: Arc, + scope: String, + round: u64, + shape: Shape, + provider: Option, + model: Option, + ) -> Self { + Self { + telemetry, + scope, + round, + shape, + provider, + model, + started: Instant::now(), + status: None, + finished: false, + } + } + pub fn headers(&mut self, status: u16) { + self.status = Some(status); + } + pub fn fail(&mut self, outcome: &str) { + self.finish_parsed( + outcome, + Parsed { + partial: true, + ..Default::default() + }, + ); + } + pub fn buffered(&mut self, status: u16, body: &[u8]) { + self.headers(status); + let parsed = if body.len() <= parse::MAX_BODY { + serde_json::from_slice(body) + .ok() + .map(|value| parse::response(&value, self.shape)) + } else { + None + }; + self.finish_parsed( + if status < 400 { + "complete" + } else { + "http_error" + }, + parsed.unwrap_or(Parsed { + partial: true, + ..Default::default() + }), + ); + } + fn finish_parsed(&mut self, outcome: &str, parsed: Parsed) { + if self.finished { + return; + } + self.finished = true; + let usage_state = match (parsed.usage.prompt_tokens, parsed.usage.completion_tokens) { + (Some(_), Some(_)) => "present", + (None, None) => "missing", + _ => "partial", + }; + let accepted = if matches!(outcome, "configuration_error" | "authentication_error") { + Some(false) + } else { + self.status.map(|status| (200..300).contains(&status)) + }; + let event = json!({"kind":"round", "round":self.round, "source":"router-upstream", + "provider":self.provider, "model":self.model, "http_status":self.status, + "accepted":accepted, + "outcome":outcome, "usage":parsed.usage, "usage_state":usage_state, + "finish_reason":parsed.finish, "partial_observation":parsed.partial, + "ms":self.started.elapsed().as_millis().min(u64::MAX as u128) as u64, + "tool_calls_observed":parsed.tools.len(), + }); + self.telemetry + .finish(&self.scope, self.round, event, parsed.tools); + } +} +impl Drop for Observation { + fn drop(&mut self) { + if !self.finished { + self.fail("cancelled"); + } + } +} + +pub fn wrap_stream( + inner: BoxStream<'static, Result>, + observation: Option, + is_sse: bool, +) -> BoxStream<'static, Result> { + match observation { + Some(observation) => { + let accumulator = stream::Accumulator::new(observation.shape, is_sse); + Box::pin(ObservedStream { + inner, + observation, + accumulator, + }) + } + None => inner, + } +} + +struct ObservedStream { + inner: BoxStream<'static, Result>, + observation: Observation, + accumulator: stream::Accumulator, +} +impl Drop for ObservedStream { + fn drop(&mut self) { + if !self.observation.finished { + let parsed = self.accumulator.finish(); + self.observation.finish_parsed("cancelled", parsed); + } + } +} +impl Stream for ObservedStream { + type Item = Result; + fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll> { + match self.inner.as_mut().poll_next(context) { + Poll::Ready(Some(Ok(bytes))) => { + self.accumulator.feed(&bytes); + Poll::Ready(Some(Ok(bytes))) + } + Poll::Ready(Some(Err(error))) => { + let parsed = self.accumulator.finish(); + self.observation + .finish_parsed("response_body_error", parsed); + Poll::Ready(Some(Err(error))) + } + Poll::Ready(None) => { + let parsed = self.accumulator.finish(); + let outcome = if self.observation.status.is_some_and(|status| status >= 400) { + "http_error" + } else if self.accumulator.failed { + "upstream_error" + } else if self.accumulator.complete() { + "complete" + } else { + "incomplete" + }; + self.observation.finish_parsed(outcome, parsed); + Poll::Ready(None) + } + Poll::Pending => Poll::Pending, + } + } +} diff --git a/inference-router/src/task_telemetry/observe/stream.rs b/inference-router/src/task_telemetry/observe/stream.rs new file mode 100644 index 000000000..04361ea50 --- /dev/null +++ b/inference-router/src/task_telemetry/observe/stream.rs @@ -0,0 +1,197 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use crate::task_telemetry::parse::{self, Parsed, Shape, add_tool, merge_usage}; +use serde_json::Value; +use std::collections::BTreeMap; + +const FRAME_LIMIT: usize = 64 * 1024; +pub(super) struct Accumulator { + shape: Shape, + sse: bool, + buffer: Vec, + discarding: bool, + tail: [u8; 4], + parsed: Parsed, + terminal: bool, + pub failed: bool, + tools: BTreeMap<(u64, u64), (String, String)>, +} +impl Accumulator { + pub(super) fn new(shape: Shape, sse: bool) -> Self { + Self { + shape, + sse, + buffer: Vec::new(), + discarding: false, + tail: [0; 4], + parsed: Parsed::default(), + terminal: false, + failed: false, + tools: BTreeMap::new(), + } + } + pub(super) fn feed(&mut self, bytes: &[u8]) { + if !self.sse { + if self.buffer.len().saturating_add(bytes.len()) <= FRAME_LIMIT && !self.discarding { + self.buffer.extend_from_slice(bytes); + } else { + self.buffer.clear(); + self.discarding = true; + self.parsed.partial = true; + } + return; + } + for byte in bytes { + self.tail.rotate_left(1); + self.tail[3] = *byte; + if !self.discarding { + if self.buffer.len() < FRAME_LIMIT { + self.buffer.push(*byte); + } else { + self.buffer.clear(); + self.discarding = true; + self.parsed.partial = true; + } + } + if self.tail[2..] == *b"\n\n" || self.tail == *b"\r\n\r\n" { + if !self.discarding { + let frame = std::mem::take(&mut self.buffer); + self.frame(&frame); + } + self.discarding = false; + } + } + } + fn frame(&mut self, bytes: &[u8]) { + let Ok(frame) = std::str::from_utf8(bytes) else { + self.parsed.partial = true; + return; + }; + let data = frame + .lines() + .filter_map(|line| line.strip_prefix("data:").map(str::trim)) + .collect::>() + .join("\n"); + if data.is_empty() { + return; + } + if data == "[DONE]" { + self.terminal = true; + return; + } + let Ok(value) = serde_json::from_str::(&data) else { + self.parsed.partial = true; + return; + }; + if let Some(usage) = value.get("usage") { + merge_usage(&mut self.parsed.usage, usage); + } + match self.shape { + Shape::OpenAi => { + if let Some(choices) = value["choices"].as_array() { + for choice in choices.iter().take(parse::MAX_TOOLS) { + if let Some(reason) = choice["finish_reason"] + .as_str() + .and_then(|value| parse::identifier(value, 64)) + { + self.parsed.finish = Some(reason); + } + if let Some(calls) = choice["delta"]["tool_calls"].as_array() { + for call in calls.iter().take(parse::MAX_TOOLS) { + let key = ( + choice["index"].as_u64().unwrap_or(0), + call["index"].as_u64().unwrap_or(0), + ); + if !self.tools.contains_key(&key) + && self.tools.len() >= parse::MAX_TOOLS + { + self.parsed.partial = true; + continue; + } + let tool = self.tools.entry(key).or_default(); + for (slot, value) in [ + (&mut tool.0, &call["id"]), + (&mut tool.1, &call["function"]["name"]), + ] { + if let Some(value) = value.as_str() { + if slot.as_str() == value { + continue; + } + if slot.len().saturating_add(value.len()) <= 128 { + slot.push_str(value); + } else { + self.parsed.partial = true; + } + } + } + } + } + } + } + } + Shape::Anthropic => match value["type"].as_str() { + Some("message_start") => { + if let Some(usage) = value["message"].get("usage") { + merge_usage(&mut self.parsed.usage, usage); + } + } + Some("content_block_start") if value["content_block"]["type"] == "tool_use" => { + add_tool( + &mut self.parsed, + &value["content_block"]["id"], + &value["content_block"]["name"], + ) + } + Some("message_delta") => { + self.parsed.finish = value["delta"]["stop_reason"] + .as_str() + .and_then(|value| parse::identifier(value, 64)); + } + Some("message_stop") => self.terminal = true, + Some("error") => { + self.terminal = true; + self.failed = true; + } + _ => {} + }, + Shape::Responses => match value["type"].as_str() { + Some("response.completed" | "response.failed" | "response.incomplete") => { + self.terminal = true; + self.failed = value["type"] != "response.completed"; + let parsed = parse::response(&value["response"], Shape::Responses); + self.parsed.usage = parsed.usage; + self.parsed.finish = parsed.finish; + self.parsed.tools = parsed.tools; + self.parsed.partial |= parsed.partial || self.failed; + } + Some("error") => { + self.terminal = true; + self.failed = true; + } + _ => {} + }, + } + } + pub(super) fn complete(&self) -> bool { + self.terminal && !self.failed + } + pub(super) fn finish(&mut self) -> Parsed { + if !self.sse && !self.discarding { + match serde_json::from_slice::(&self.buffer) { + Ok(value) => { + self.parsed = parse::response(&value, self.shape); + self.terminal = true; + } + Err(_) => self.parsed.partial = true, + } + } + if !self.terminal { + self.parsed.partial = true; + } + for (_, (id, name)) in std::mem::take(&mut self.tools) { + add_tool(&mut self.parsed, &Value::String(id), &Value::String(name)); + } + std::mem::take(&mut self.parsed) + } +} diff --git a/inference-router/src/task_telemetry/parse.rs b/inference-router/src/task_telemetry/parse.rs new file mode 100644 index 000000000..a3942cdd4 --- /dev/null +++ b/inference-router/src/task_telemetry/parse.rs @@ -0,0 +1,187 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use serde::Serialize; +use serde_json::Value; + +pub const MAX_BODY: usize = 64 * 1024; +pub const MAX_TOOLS: usize = 32; + +#[derive(Clone, Copy, Debug)] +pub enum Shape { + OpenAi, + Responses, + Anthropic, +} +impl Shape { + pub fn for_path(path: &str) -> Option { + match path.trim_matches('/').trim_start_matches("v1/") { + "chat/completions" => Some(Self::OpenAi), + "responses" => Some(Self::Responses), + "messages" => Some(Self::Anthropic), + _ => None, + } + } +} + +pub fn identifier(value: &str, max: usize) -> Option { + crate::access_request::scope::identifier(value, max).then(|| value.to_string()) +} +#[derive(Default, Clone, Serialize)] +pub struct Usage { + pub prompt_tokens: Option, + pub completion_tokens: Option, + pub total_tokens: Option, + pub cached_tokens: Option, +} +#[derive(Default)] +pub struct Parsed { + pub usage: Usage, + pub finish: Option, + pub tools: Vec<(String, String)>, + pub partial: bool, +} + +pub fn merge_usage(target: &mut Usage, usage: &Value) { + let get = |key: &str| usage.get(key).and_then(Value::as_u64); + if let Some(value) = get("prompt_tokens").or_else(|| get("input_tokens")) { + target.prompt_tokens = Some(value); + } + if let Some(value) = get("completion_tokens").or_else(|| get("output_tokens")) { + target.completion_tokens = Some(value); + } + if let Some(value) = get("total_tokens") { + target.total_tokens = Some(value); + } + if let Some(value) = get("cache_read_input_tokens").or_else(|| { + usage + .get("prompt_tokens_details") + .or_else(|| usage.get("input_tokens_details")) + .and_then(|value| value.get("cached_tokens")) + .and_then(Value::as_u64) + }) { + target.cached_tokens = Some(value); + } +} + +pub fn add_tool(parsed: &mut Parsed, id: &Value, name: &Value) { + let Some(name) = name.as_str().and_then(|value| identifier(value, 128)) else { + parsed.partial = true; + return; + }; + let id = id + .as_str() + .and_then(|value| identifier(value, 128)) + .unwrap_or_default(); + if parsed.tools.len() >= MAX_TOOLS { + parsed.partial = true; + return; + } + parsed.tools.push((id, name)); +} + +pub fn response(value: &Value, shape: Shape) -> Parsed { + let mut parsed = Parsed::default(); + if let Some(usage) = value.get("usage") { + merge_usage(&mut parsed.usage, usage); + } + match shape { + Shape::OpenAi => { + if let Some(choices) = value.get("choices").and_then(Value::as_array) { + parsed.partial |= choices.len() > MAX_TOOLS; + for choice in choices.iter().take(MAX_TOOLS) { + parsed.finish = choice + .get("finish_reason") + .and_then(Value::as_str) + .and_then(|value| identifier(value, 64)); + if let Some(tools) = choice + .pointer("/message/tool_calls") + .and_then(Value::as_array) + { + for tool in tools.iter().take(MAX_TOOLS + 1) { + add_tool(&mut parsed, &tool["id"], &tool["function"]["name"]); + } + } + } + } + } + Shape::Responses => { + parsed.finish = value + .get("status") + .and_then(Value::as_str) + .and_then(|value| identifier(value, 64)); + if let Some(items) = value.get("output").and_then(Value::as_array) { + parsed.partial |= items.len() > MAX_TOOLS; + for item in items.iter().take(MAX_TOOLS + 1) { + if item["type"] == "function_call" { + add_tool(&mut parsed, &item["call_id"], &item["name"]); + } + } + } + } + Shape::Anthropic => { + parsed.finish = value + .get("stop_reason") + .and_then(Value::as_str) + .and_then(|value| identifier(value, 64)); + if let Some(items) = value.get("content").and_then(Value::as_array) { + parsed.partial |= items.len() > MAX_TOOLS; + for item in items.iter().take(MAX_TOOLS + 1) { + if item["type"] == "tool_use" { + add_tool(&mut parsed, &item["id"], &item["name"]); + } + } + } + } + } + parsed +} + +pub fn request_results(body: &[u8], shape: Shape) -> Vec<(String, Option)> { + if body.len() > MAX_BODY { + return Vec::new(); + } + let Ok(value) = serde_json::from_slice::(body) else { + return Vec::new(); + }; + let mut results = Vec::new(); + let mut add = |id: &Value, ok: Option| { + if results.len() < MAX_TOOLS + && let Some(id) = id.as_str().and_then(|value| identifier(value, 128)) + { + results.push((id, ok)); + } + }; + if matches!(shape, Shape::Responses) { + if let Some(items) = value["input"].as_array() { + for item in items.iter().rev().take(128) { + if item["type"] == "function_call_output" { + add( + &item["call_id"], + item["is_error"].as_bool().map(|error| !error), + ); + } + } + } + } else if let Some(messages) = value["messages"].as_array() { + for message in messages.iter().rev().take(128) { + if message["role"] == "tool" { + add( + &message["tool_call_id"], + message["is_error"].as_bool().map(|error| !error), + ); + } + if let Some(content) = message["content"].as_array() { + for item in content.iter().take(MAX_TOOLS) { + if item["type"] == "tool_result" { + add( + &item["tool_use_id"], + item["is_error"].as_bool().map(|error| !error), + ); + } + } + } + } + } + results +} diff --git a/inference-router/tests/agt_governance_integration.rs b/inference-router/tests/agt_governance_integration.rs index 76c4b15e0..016ec6386 100644 --- a/inference-router/tests/agt_governance_integration.rs +++ b/inference-router/tests/agt_governance_integration.rs @@ -38,6 +38,7 @@ fn test_state(sandbox: &str, admin_token: Option<&str>) -> AppState { let policy_status = Arc::new(kars_inference_router::policy_status::PolicyStatusRegistry::new()); let governance = Arc::new(Governance::new_with_status(sandbox, policy_status.clone())); AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/anthropic_buffered_guardrail.rs b/inference-router/tests/anthropic_buffered_guardrail.rs index 750fe6909..2151fddc9 100644 --- a/inference-router/tests/anthropic_buffered_guardrail.rs +++ b/inference-router/tests/anthropic_buffered_guardrail.rs @@ -49,6 +49,7 @@ fn test_state(anthropic_endpoint: String, moderation_endpoint: String) -> AppSta policy_status.clone(), )); AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/chat_output_guardrail_nonjson.rs b/inference-router/tests/chat_output_guardrail_nonjson.rs index 99db094e3..ad80ae9cf 100644 --- a/inference-router/tests/chat_output_guardrail_nonjson.rs +++ b/inference-router/tests/chat_output_guardrail_nonjson.rs @@ -47,6 +47,7 @@ fn test_state(ollama_endpoint: String, moderation_endpoint: String) -> AppState policy_status.clone(), )); AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/common/governed_services.rs b/inference-router/tests/common/governed_services.rs new file mode 100644 index 000000000..a59804578 --- /dev/null +++ b/inference-router/tests/common/governed_services.rs @@ -0,0 +1,196 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#![allow(dead_code)] + +use axum::{ + Router, + body::Body, + extract::ConnectInfo, + http::{Request, StatusCode}, +}; +use kars_inference_router::{ + access_request::{Identity, scope::ResourceIdentity}, + auth::WorkloadIdentityAuth, + blocklist::Blocklist, + budget::TokenBudgetTracker, + config::{Config, RegistryMode}, + egress_blocked::BlockedBuffer, + governance::Governance, + governed_services::GovernedServices, + handoff::{DrainState, HandoffSession, HandoffTokenStore, PendingHandoffStore}, + mesh::{MeshInbox, MeshMetrics}, + policy_status::PolicyStatusRegistry, + routes::{self, AppState}, +}; +use serde_json::{Value, json}; +use std::sync::Arc; +use tower::ServiceExt; + +pub const CONTROL: &str = "operator-only-service-control-credential-0123456789abcdef"; +pub const AGENT_ADMIN: &str = "legacy-agent-admin-credential-0123456789abcdef0123456789"; + +pub fn state(workspace: &str, uid: &str) -> AppState { + let registry = Arc::new(PolicyStatusRegistry::new()); + let governance = Arc::new(Governance::new_with_status("test", registry.clone())); + let services = Arc::new(GovernedServices::new( + Identity { + sandbox: ResourceIdentity { + namespace: workspace.into(), + name: "test".into(), + uid: uid.into(), + }, + namespace_uid: format!("namespace-{uid}"), + task: Some(ResourceIdentity { + namespace: workspace.into(), + name: "task".into(), + uid: format!("task-{uid}"), + }), + task_authorization: Some(format!("sha256:{}", "a".repeat(64))), + task_generation: Some(1), + managed: true, + }, + Some(CONTROL.into()), + )); + let blocked = Arc::new(BlockedBuffer::with_defaults()); + blocked.bind_services(&services); + AppState { + services, + auth: Arc::new(WorkloadIdentityAuth::new()), + copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), + client: reqwest::Client::new(), + config: Arc::new(Config { + port: 0, + foundry_endpoint: None, + foundry_project_endpoint: None, + azure_openai_endpoint: None, + default_model: "model".into(), + content_safety_enabled: false, + prompt_shields_enabled: false, + content_safety_endpoint: None, + token_budget_daily: 1000, + token_budget_per_request: 100, + 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(), + providers: Default::default(), + }), + budget: TokenBudgetTracker::new(1000, 100), + policy_provider: governance.clone(), + audit_sink: governance.clone(), + signing_provider: governance.clone(), + governance, + blocklist: Blocklist::disabled(), + blocked_egress: blocked, + sandbox_name: Arc::new("test".into()), + inbox: Arc::new(MeshInbox::new()), + mesh_metrics: Arc::new(MeshMetrics::new()), + model_override: Default::default(), + admin_token: Some(Arc::new(AGENT_ADMIN.into())), + responses_only_models: Default::default(), + unavailable_models: Default::default(), + handoff_tokens: HandoffTokenStore::new(), + handoff_session: HandoffSession::new(), + drain_state: DrainState::new(), + pending_handoff: PendingHandoffStore::new(), + policy_status: registry, + inference_policy: kars_inference_router::inference_policy_loader::empty_handle(), + memory_binding: kars_inference_router::memory_binding_loader::empty_handle(), + egress_allowlist: kars_inference_router::egress_allowlist_loader::empty_handle(), + deployment_health: Arc::new( + kars_inference_router::deployment_health::DeploymentHealthRegistry::new(), + ), + } +} +pub fn app(state: AppState) -> Router { + routes::governed_service_routes(state.clone()) + .merge(routes::egress_routes()) + .merge(routes::sensitive_agt_routes()) + .with_state(state) +} +pub fn request( + method: &str, + path: &str, + body: Value, + token: Option<&str>, + scope: Option<&str>, + local: bool, +) -> Request { + let mut request = Request::builder() + .method(method) + .uri(path) + .header("content-type", "application/json"); + if let Some(token) = token { + request = request.header("authorization", format!("Bearer {token}")); + } + if let Some(scope) = scope { + request = request.header("x-kars-service-scope", scope); + } + let mut request = request.body(Body::from(body.to_string())).unwrap(); + request + .extensions_mut() + .insert(ConnectInfo::( + if local { + "127.0.0.1:32100" + } else { + "192.0.2.10:32100" + } + .parse() + .unwrap(), + )); + request +} +pub async fn send(app: &Router, request: Request) -> (StatusCode, Value) { + let response = app.clone().oneshot(request).await.unwrap(); + let status = response.status(); + let body = axum::body::to_bytes(response.into_body(), 2 * 1024 * 1024) + .await + .unwrap(); + ( + status, + serde_json::from_slice(&body) + .unwrap_or_else(|_| json!({"text":String::from_utf8_lossy(&body)})), + ) +} +pub async fn queue( + app: &Router, + scope: &str, + kind: &str, + target: &str, + port: Option, +) -> Value { + let (status, body) = send( + app, + request( + "POST", + "/v1/access-request", + json!({"scope_id":scope,"kind":kind,"target":target,"reason":"needed", "port":port}), + None, + None, + true, + ), + ) + .await; + assert_eq!(status, StatusCode::ACCEPTED, "{body}"); + body["request"].clone() +} +pub async fn decision(app: &Router, scope: &str, id: &str, verdict: &str) -> (StatusCode, Value) { + send( + app, + request( + "POST", + "/internal/access-requests/decision", + json!({"scope_id":scope,"request_id":id,"verdict":verdict}), + Some(CONTROL), + None, + true, + ), + ) + .await +} diff --git a/inference-router/tests/egress_blocked_endpoint.rs b/inference-router/tests/egress_blocked_endpoint.rs index 843a17772..6b8c8bfaf 100644 --- a/inference-router/tests/egress_blocked_endpoint.rs +++ b/inference-router/tests/egress_blocked_endpoint.rs @@ -36,6 +36,7 @@ fn test_state() -> AppState { policy_status.clone(), )); AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/failover_walk.rs b/inference-router/tests/failover_walk.rs index b499862c7..41fe902a5 100644 --- a/inference-router/tests/failover_walk.rs +++ b/inference-router/tests/failover_walk.rs @@ -131,6 +131,7 @@ async fn primary_503_falls_through_to_fallback_200() { let health = Arc::new(DeploymentHealthRegistry::new()); let upstream = UpstreamConfig { + telemetry: None, endpoint: base, deployment: "fallback-up".into(), sandbox_name: "sbx".into(), @@ -198,6 +199,7 @@ async fn unhealthy_primary_is_skipped_in_second_pass() { assert!(!health.is_healthy("Foundry::primary-down")); let upstream = UpstreamConfig { + telemetry: None, endpoint: base, deployment: "fallback-up".into(), sandbox_name: "sbx".into(), @@ -257,6 +259,7 @@ async fn all_unhealthy_still_punches_primary_for_last_resort() { assert!(!health.is_healthy("Foundry::fallback-up")); let upstream = UpstreamConfig { + telemetry: None, endpoint: base, deployment: "primary-down".into(), sandbox_name: "sbx".into(), diff --git a/inference-router/tests/foundry_route_guard.rs b/inference-router/tests/foundry_route_guard.rs index 76d5e3f06..b5913fecb 100644 --- a/inference-router/tests/foundry_route_guard.rs +++ b/inference-router/tests/foundry_route_guard.rs @@ -54,6 +54,7 @@ fn test_state() -> AppState { policy_status.clone(), )); AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/governed_access_services.rs b/inference-router/tests/governed_access_services.rs new file mode 100644 index 000000000..5c9bec61b --- /dev/null +++ b/inference-router/tests/governed_access_services.rs @@ -0,0 +1,533 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#[path = "common/governed_services.rs"] +mod support; +use axum::http::StatusCode; +use kars_inference_router::access_request::{ + AccessRequestBuffer, Error, Identity, Request, Status, +}; +use serde_json::json; +use std::{sync::Arc, time::Duration}; +use support::*; +use tower::ServiceExt; + +#[tokio::test] +async fn privileged_http_endpoints_reject_agent_admin_even_on_loopback() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state); + for token in [None, Some(AGENT_ADMIN)] { + for path in [ + "/internal/access-requests/reset", + "/internal/access-requests/decision", + ] { + let (status, _) = send( + &app, + request("POST", path, json!({"scope_id":scope}), token, None, true), + ) + .await; + assert_eq!(status, StatusCode::UNAUTHORIZED); + } + assert_eq!( + send( + &app, + request( + "GET", + "/internal/access-requests", + json!({}), + token, + None, + true + ) + ) + .await + .0, + StatusCode::UNAUTHORIZED + ); + } + assert_eq!( + send( + &app, + request("GET", "/v1/access-requests", json!({}), None, None, false) + ) + .await + .0, + StatusCode::UNAUTHORIZED + ); + assert_eq!( + send( + &app, + request( + "GET", + "/internal/access-requests", + json!({}), + Some(CONTROL), + None, + false + ) + ) + .await + .0, + StatusCode::OK + ); +} + +#[tokio::test] +async fn reset_fences_old_requests_and_preserves_uid_qualified_identity() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + state.budget.record_usage("test", 17).await; + let app = app(state.clone()); + let entry = queue(&app, &scope, "egress", "example.test", Some(443)).await; + let (status, reset) = send( + &app, + request( + "POST", + "/internal/access-requests/reset", + json!({"scope_id":scope,"assignment_id":"assignment-2"}), + Some(CONTROL), + None, + true, + ), + ) + .await; + assert_eq!(status, StatusCode::OK); + assert_eq!(reset["scope"]["identity"]["sandbox"]["uid"], "uid-a"); + assert_eq!(reset["scope"]["identity"]["task"]["uid"], "task-uid-a"); + assert_eq!(reset["enforcement_changed"], false); + assert_ne!(reset["scope"]["id"], scope); + assert_eq!( + decision( + &app, + &scope, + entry["request_id"].as_str().unwrap(), + "approved" + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!( + send( + &app, + request( + "POST", + "/v1/access-request", + json!({"scope_id":scope,"kind":"egress","target":"example.test"}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!( + send( + &app, + request( + "POST", + "/internal/access-requests/reset", + json!({"scope_id":scope}), + Some(CONTROL), + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!(state.budget.get_usage("test").await.0, 17); +} + +#[tokio::test] +async fn request_identity_port_tier_and_terminal_decisions_cannot_be_rebound() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state); + let first = queue(&app, &scope, "egress", "example.test", Some(443)).await; + let id = first["request_id"].as_str().unwrap(); + assert_eq!( + decision(&app, &scope, id, "approved").await.0, + StatusCode::OK + ); + assert_eq!( + decision(&app, &scope, id, "denied").await.0, + StatusCode::CONFLICT + ); + let other = queue(&app, &scope, "egress", "example.test", Some(80)).await; + assert_ne!(first["request_id"], other["request_id"]); + assert_eq!(other["status"], "pending"); + let (status,body)=send(&app,request("POST","/v1/access-request", + json!({"scope_id":scope,"kind":"egress","target":"example.test","reason":"changed","port":443}),None,None,true)).await; + assert_eq!(status, StatusCode::ACCEPTED); + assert_eq!(body["request"]["reason"], "needed"); + assert_eq!(body["request"]["status"], "approved"); + for (invalid, expected) in [ + ( + json!({"scope_id":scope,"kind":"tier","tier":6}), + StatusCode::BAD_REQUEST, + ), + ( + json!({"scope_id":scope,"kind":"egress","target":"https://example.test/token"}), + StatusCode::BAD_REQUEST, + ), + ( + json!({"scope_id":scope,"kind":"tool","target":"test","sandbox_uid":"other"}), + StatusCode::UNPROCESSABLE_ENTITY, + ), + ] { + assert_eq!( + send( + &app, + request("POST", "/v1/access-request", invalid, None, None, true) + ) + .await + .0, + expected + ); + } +} + +#[tokio::test] +async fn same_named_sandboxes_in_other_workspaces_cannot_share_request_scope() { + let a = state("workspace-a", "uid-a"); + let b = state("workspace-b", "uid-b"); + let scope = a.services.requests.scope().unwrap().id; + let app_b = app(b); + assert_eq!( + send( + &app_b, + request( + "POST", + "/v1/access-request", + json!({"scope_id":scope,"kind":"tool","target":"fetch"}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); +} + +#[tokio::test] +async fn decision_wait_wakes_and_agent_cancel_does_not_rewrite_an_operator_decision() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state.clone()); + let entry = queue(&app, &scope, "tool", "fetch", None).await; + let id = entry["request_id"].as_str().unwrap(); + let waiting = app.clone().oneshot(request( + "GET", + &format!("/v1/access-requests/{id}/wait?scope_id={scope}&timeout_ms=1000"), + json!({}), + None, + None, + true, + )); + let wait = tokio::spawn(waiting); + assert_eq!( + decision(&app, &scope, id, "approved").await.0, + StatusCode::OK + ); + assert_eq!(wait.await.unwrap().unwrap().status(), StatusCode::OK); + assert_eq!( + send( + &app, + request( + "POST", + &format!("/v1/access-requests/{id}/cancel"), + json!({"scope_id":scope}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::OK + ); + let cancelled = state.services.requests.entry(&scope, id).unwrap(); + assert_eq!(cancelled.status, Status::Cancelled); + assert_eq!(cancelled.decision, Some(Status::Approved)); +} + +#[tokio::test] +async fn storage_rate_and_expiry_limits_fail_closed() { + let buffer = AccessRequestBuffer::with_limits( + Identity::standalone("test"), + 1, + 128, + Duration::from_millis(10), + ); + let scope = buffer.scope().unwrap().id; + let request = |target: &str| Request { + scope_id: scope.clone(), + kind: "tool".into(), + target: target.into(), + reason: "reason".into(), + tier: None, + port: None, + }; + let (first, _) = buffer.record(request("first")).unwrap(); + assert!(matches!(buffer.record(request("second")), Err(Error::Full))); + tokio::time::sleep(Duration::from_millis(20)).await; + assert!(matches!( + buffer.transition(&scope, &first.request_id, Status::Approved), + Err(Error::Expired) + )); + let rate = AccessRequestBuffer::with_limits( + Identity::standalone("test"), + 64, + 1, + Duration::from_secs(30), + ); + let scope = rate.scope().unwrap().id; + let request = Request { + scope_id: scope, + kind: "tool".into(), + target: "first".into(), + reason: String::new(), + tier: None, + port: None, + }; + rate.record(request.clone()).unwrap(); + assert!(matches!(rate.record(request), Err(Error::RateLimited))); +} + +#[tokio::test] +async fn dropping_waiter_and_reset_release_bounded_wait_capacity() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state.clone()); + let entry = queue(&app, &scope, "tool", "fetch", None).await; + let id = entry["request_id"].as_str().unwrap().to_string(); + let services = state.services.clone(); + let scope_copy = scope.clone(); + let task = tokio::spawn(async move { + services + .wait_for_decision(&scope_copy, &id, Duration::from_secs(30)) + .await + }); + for _ in 0..30 { + if state.services.wait_slots() == 15 { + break; + } + tokio::task::yield_now().await; + } + assert_eq!(state.services.wait_slots(), 15); + task.abort(); + let _ = task.await; + assert_eq!(state.services.wait_slots(), 16); + let scope_copy = scope.clone(); + let id = entry["request_id"].as_str().unwrap().to_string(); + let services = state.services.clone(); + let task = tokio::spawn(async move { + services + .wait_for_decision(&scope_copy, &id, Duration::from_secs(30)) + .await + }); + state.services.reset(&scope, Some("next".into())).unwrap(); + assert!(matches!(task.await.unwrap(), Err(Error::StaleScope))); + assert_eq!(state.services.wait_slots(), 16); +} + +#[tokio::test] +async fn missing_control_configuration_never_falls_back_to_shared_agent_admin() { + let mut state = state("workspace-a", "uid-a"); + state.services = Arc::new( + kars_inference_router::governed_services::GovernedServices::new( + Identity::standalone("test"), + None, + ), + ); + let app = app(state); + assert_eq!( + send( + &app, + request( + "GET", + "/internal/access-requests", + json!({}), + Some(AGENT_ADMIN), + None, + true + ) + ) + .await + .0, + StatusCode::SERVICE_UNAVAILABLE + ); +} + +#[tokio::test] +async fn http_payload_rate_and_operator_origin_limits_are_enforced() { + let mut state = state("workspace-a", "uid-a"); + let identity = state.services.requests.scope().unwrap().identity; + let mut services = kars_inference_router::governed_services::GovernedServices::new( + identity, + Some(CONTROL.into()), + ); + services.allow_ips = Some(vec!["192.0.2.11".parse().unwrap()]); + state.services = Arc::new(services); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state.clone()); + assert_eq!( + send( + &app, + request( + "GET", + "/internal/access-requests", + json!({}), + Some(CONTROL), + None, + false + ) + ) + .await + .0, + StatusCode::FORBIDDEN + ); + assert_eq!( + send( + &app, + request( + "GET", + "/telemetry/cursor", + json!({}), + Some(CONTROL), + None, + false + ) + ) + .await + .0, + StatusCode::FORBIDDEN + ); + assert_eq!( + send( + &app, + request( + "POST", + "/v1/access-request", + json!({ + "scope_id":scope,"kind":"tool","target":"fetch","reason":"x".repeat(9000) + }), + None, + None, + true + ) + ) + .await + .0, + StatusCode::PAYLOAD_TOO_LARGE + ); + for _ in 0..32 { + assert_eq!( + send( + &app, + request( + "POST", + "/v1/access-request", + json!({ + "scope_id":scope,"kind":"tool","target":"fetch" + }), + None, + None, + true + ) + ) + .await + .0, + StatusCode::ACCEPTED + ); + } + assert_eq!( + send( + &app, + request( + "POST", + "/v1/access-request", + json!({ + "scope_id":scope,"kind":"tool","target":"fetch" + }), + None, + None, + true + ) + ) + .await + .0, + StatusCode::TOO_MANY_REQUESTS + ); +} + +#[tokio::test] +async fn unknown_decision_cannot_synthesize_an_approved_request() { + let state = state("workspace-a", "uid-a"); + let scope = state.services.requests.scope().unwrap().id; + let app = app(state.clone()); + assert_eq!( + decision(&app, &scope, "unknown-id", "approved").await.0, + StatusCode::NOT_FOUND + ); + assert!(state.services.requests.snapshot(&scope).unwrap().is_empty()); +} + +#[tokio::test] +async fn live_http_control_auth_and_reset_use_the_same_production_router() { + let state = state("workspace-a", "uid-a"); + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let shutdown = tokio_util::sync::CancellationToken::new(); + let stop = shutdown.clone(); + let server = tokio::spawn(async move { + axum::serve( + listener, + app(state).into_make_service_with_connect_info::(), + ) + .with_graceful_shutdown(stop.cancelled_owned()) + .await + .unwrap(); + }); + let client = reqwest::Client::builder().no_proxy().build().unwrap(); + let base = format!("http://{address}"); + let info = client + .get(format!("{base}/v1/access-requests")) + .send() + .await + .unwrap(); + assert_eq!(info.status(), StatusCode::OK); + let scope = info.json::().await.unwrap()["scope"]["id"] + .as_str() + .unwrap() + .to_string(); + let response = client + .post(format!("{base}/internal/access-requests/reset")) + .bearer_auth(AGENT_ADMIN) + .json(&json!({"scope_id":scope})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + let response = client + .post(format!("{base}/internal/access-requests/reset")) + .bearer_auth(CONTROL) + .json(&json!({"scope_id":scope,"assignment_id":"live-test"})) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_ne!( + response.json::().await.unwrap()["scope"]["id"], + scope + ); + shutdown.cancel(); + server.await.unwrap(); +} diff --git a/inference-router/tests/governed_egress_wait.rs b/inference-router/tests/governed_egress_wait.rs new file mode 100644 index 000000000..7a3aadb1f --- /dev/null +++ b/inference-router/tests/governed_egress_wait.rs @@ -0,0 +1,254 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#[path = "common/governed_services.rs"] +mod support; +use axum::http::StatusCode; +use kars_inference_router::{blocklist::Blocklist, routes::AppState}; +use serde_json::{Value, json}; +use std::time::Duration; +use support::*; +use tower::ServiceExt; +use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + +async fn setup() -> (AppState, axum::Router, MockServer, String, String) { + let upstream = MockServer::start().await; + Mock::given(method("GET")) + .respond_with(ResponseTemplate::new(200).set_body_string("response-secret-marker")) + .mount(&upstream) + .await; + let address = reqwest::Url::parse(&upstream.uri()) + .unwrap() + .socket_addrs(|| None) + .unwrap()[0]; + let mut state = state("workspace", "uid"); + state.blocklist = Blocklist::new(None).await; + state.client = reqwest::Client::builder() + .no_proxy() + .redirect(reqwest::redirect::Policy::none()) + .resolve("approved.example", address) + .build() + .unwrap(); + let scope = state.services.requests.scope().unwrap().id; + let url = format!( + "http://approved.example:{}/data?key=request-secret-marker", + address.port() + ); + (state.clone(), app(state), upstream, scope, url) +} + +async fn pending(app: &axum::Router) -> Value { + for _ in 0..100 { + let (_, body) = send( + app, + request( + "GET", + "/internal/access-requests", + json!({}), + Some(CONTROL), + None, + true, + ), + ) + .await; + if let Some(entry) = body["entries"] + .as_array() + .and_then(|entries| entries.first()) + { + return entry.clone(); + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + panic!("request was not queued"); +} + +#[tokio::test] +async fn approved_decision_alone_cannot_widen_egress_but_signed_policy_wakes_the_request() { + let (state, app, upstream, scope, url) = setup().await; + let waiting = tokio::spawn(app.clone().oneshot(request( + "POST", + "/egress/fetch", + json!({"url":url,"scope_id":scope,"wait_for_approval_ms":2000, + "headers":{"Authorization":"Bearer header-secret-marker"}}), + None, + None, + true, + ))); + let entry = pending(&app).await; + assert_eq!( + decision( + &app, + &scope, + entry["request_id"].as_str().unwrap(), + "approved" + ) + .await + .0, + StatusCode::OK + ); + tokio::time::sleep(Duration::from_millis(30)).await; + assert!(!waiting.is_finished()); + assert!(upstream.received_requests().await.unwrap().is_empty()); + // This is the primitive used by the existing signed allowlist loader; the + // access-request service has no API capable of writing it. + state + .blocklist + .replace_allowlist(vec!["approved.example".into()]) + .await; + let response = waiting.await.unwrap().unwrap(); + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(upstream.received_requests().await.unwrap().len(), 1); + let trace = state + .services + .telemetry + .snapshot(&scope, 0) + .unwrap() + .to_string(); + for secret in [ + "request-secret-marker", + "response-secret-marker", + "header-secret-marker", + "approved.example", + ] { + assert!(!trace.contains(secret), "{trace}"); + } + assert!(!state.blocklist.is_learn_mode()); +} + +#[tokio::test] +async fn denial_cancel_reset_and_timeout_do_not_send_unapproved_traffic() { + for action in ["deny", "cancel", "reset", "timeout"] { + let (state, app, upstream, scope, url) = setup().await; + let waiting=tokio::spawn(app.clone().oneshot(request("POST","/egress/fetch", + json!({"url":url,"scope_id":scope,"wait_for_approval_ms":if action=="timeout"{30}else{1000}}),None,None,true))); + let entry = pending(&app).await; + let id = entry["request_id"].as_str().unwrap(); + match action { + "deny" => { + assert_eq!(decision(&app, &scope, id, "denied").await.0, StatusCode::OK); + } + "cancel" => { + assert_eq!( + decision(&app, &scope, id, "approved").await.0, + StatusCode::OK + ); + assert_eq!( + send( + &app, + request( + "POST", + &format!("/v1/access-requests/{id}/cancel"), + json!({"scope_id":scope}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::OK + ); + } + "reset" => { + assert_eq!( + send( + &app, + request( + "POST", + "/internal/access-requests/reset", + json!({"scope_id":scope,"assignment_id":"next"}), + Some(CONTROL), + None, + true + ) + ) + .await + .0, + StatusCode::OK + ); + } + _ => {} + } + let response = waiting.await.unwrap().unwrap(); + assert_eq!( + response.status(), + if action == "timeout" { + StatusCode::REQUEST_TIMEOUT + } else { + StatusCode::CONFLICT + } + ); + assert!(upstream.received_requests().await.unwrap().is_empty()); + assert_eq!(state.services.wait_slots(), 16); + } +} + +#[tokio::test] +async fn ordinary_denied_fetch_is_still_immediate_and_private_targets_never_enter_approval_wait() { + let (_state, app, upstream, scope, url) = setup().await; + let response = tokio::time::timeout( + Duration::from_millis(500), + send( + &app, + request( + "POST", + "/egress/fetch", + json!({"url":url}), + None, + None, + true, + ), + ), + ) + .await + .unwrap(); + assert_eq!(response.0, StatusCode::FORBIDDEN); + let response = send( + &app, + request( + "POST", + "/egress/fetch", + json!({"url":upstream.uri(),"scope_id":scope,"wait_for_approval_ms":1000}), + None, + None, + true, + ), + ) + .await; + assert_eq!(response.0, StatusCode::FORBIDDEN); + assert!(upstream.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn wait_rejects_wrong_port_request_and_shutdown_cancels_waiters() { + let (state, app, _upstream, scope, url) = setup().await; + let entry = queue(&app, &scope, "egress", "approved.example", Some(443)).await; + let id = entry["request_id"].as_str().unwrap(); + assert!(matches!( + state + .services + .wait_for_egress( + &state.blocklist, + &scope, + id, + &url, + "test", + Duration::from_millis(20) + ) + .await, + Err(kars_inference_router::access_request::Error::Invalid) + )); + let services = state.services.clone(); + let id = id.to_string(); + let old = scope.clone(); + let waiter = tokio::spawn(async move { + services + .wait_for_decision(&old, &id, Duration::from_secs(30)) + .await + }); + state.services.shutdown.cancel(); + assert!(matches!( + waiter.await.unwrap(), + Err(kars_inference_router::access_request::Error::Unavailable) + )); +} diff --git a/inference-router/tests/governed_telemetry.rs b/inference-router/tests/governed_telemetry.rs new file mode 100644 index 000000000..2aac70446 --- /dev/null +++ b/inference-router/tests/governed_telemetry.rs @@ -0,0 +1,484 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +#[path = "common/governed_services.rs"] +mod support; +use axum::{ + body::Body, + http::{HeaderMap, Method, Request, StatusCode}, +}; +use bytes::Bytes; +use futures::{StreamExt, TryStreamExt}; +use kars_inference_router::{ + auth::WorkloadIdentityAuth, + config::ProviderEndpoint, + proxy::{self, AuthenticationProvenance, UpstreamConfig}, + routes::{self, McpRouteState}, + task_telemetry::{TaskTelemetry, observe::wrap_stream}, +}; +use serde_json::{Value, json}; +use std::sync::Arc; +use support::*; +use tower::ServiceExt; +use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + +fn telemetry() -> Arc { + Arc::new(TaskTelemetry::new("scope-a".into())) +} +fn upstream(endpoint: String, telemetry: Arc) -> UpstreamConfig { + let mut target = UpstreamConfig::azure(endpoint, "model".into(), "test".into()); + target.authentication = AuthenticationProvenance::Named { + provider_id: "qualified-provider".into(), + }; + target.telemetry = Some(telemetry); + target +} +fn events(telemetry: &TaskTelemetry) -> Vec { + telemetry.snapshot("scope-a", 0).unwrap()["events"] + .as_array() + .unwrap() + .clone() +} + +#[tokio::test] +async fn buffered_http_producer_observes_usage_but_never_api_bodies_or_tool_arguments() { + let server = MockServer::start().await; + Mock::given(method("POST")).respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "choices":[{"finish_reason":"tool_calls","message":{"content":"output-secret-marker","tool_calls":[ + {"id":"call-1","function":{"name":"fetch","arguments":"argument-secret-marker"}} + ]}}],"usage":{"prompt_tokens":11,"completion_tokens":7,"total_tokens":18} + }))).mount(&server).await; + let telemetry = telemetry(); + let target = upstream(server.uri(), telemetry.clone()); + let client = reqwest::Client::builder().no_proxy().build().unwrap(); + let (_, _, body) = proxy::forward( + &WorkloadIdentityAuth::new(), + None, + &client, + &target, + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from(r#"{"messages":[{"role":"user","content":"prompt-secret-marker"}]}"#), + ) + .await + .unwrap(); + assert!(String::from_utf8_lossy(&body).contains("output-secret-marker")); + let trace = events(&telemetry); + assert_eq!(trace[0]["usage"]["total_tokens"], 18); + assert_eq!(trace[0]["provider"], "qualified-provider"); + assert_eq!(trace[1]["kind"], "tool_proposed"); + assert!(trace[1]["ok"].is_null()); + proxy::forward(&WorkloadIdentityAuth::new(),None,&client,&target,Method::POST, + "chat/completions",&HeaderMap::new(),Bytes::from(r#"{"messages":[{"role":"tool","tool_call_id":"call-1","content":"result-secret-marker","is_error":true}]}"#)).await.unwrap(); + let trace = events(&telemetry); + assert!(trace.iter().any(|event| event["kind"] == "tool_result" + && event["ok"] == false + && event["source"] == "harness-reported")); + let serialized = serde_json::to_string(&trace).unwrap(); + for secret in [ + "prompt-secret-marker", + "output-secret-marker", + "argument-secret-marker", + "result-secret-marker", + ] { + assert!(!serialized.contains(secret)); + } +} + +#[tokio::test] +async fn fragmented_streams_preserve_bytes_and_report_each_supported_usage_shape_once() { + let cases = [ + ( + "chat/completions", + concat!( + "data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"🚀stream-secret-marker\"},\"finish_reason\":\"stop\"}]}\n\n", + "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":4,\"completion_tokens\":5,\"total_tokens\":9}}\n\n", + "data: [DONE]\n\n" + ), + 4, + 5, + Some(9), + ), + ( + "v1/messages", + concat!( + "data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":12,\"cache_read_input_tokens\":3}}}\n\n", + "data: {\"type\":\"message_delta\",\"usage\":{\"output_tokens\":8},\"delta\":{\"stop_reason\":\"end_turn\"}}\n\n", + "data: {\"type\":\"message_stop\"}\n\n" + ), + 12, + 8, + None, + ), + ( + "responses", + "data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"usage\":{\"input_tokens\":2,\"output_tokens\":3,\"total_tokens\":5},\"output\":[]}}\n\n", + 2, + 3, + Some(5), + ), + ]; + for (path, body, input, output, total) in cases { + let telemetry = telemetry(); + let mut observation = telemetry.begin(path, "provider", "model", b"{}").unwrap(); + observation.headers(200); + let chunks = body + .as_bytes() + .iter() + .map(|byte| Ok::<_, reqwest::Error>(Bytes::from(vec![*byte]))) + .collect::>(); + let bytes = wrap_stream( + futures::stream::iter(chunks).boxed(), + Some(observation), + true, + ) + .try_collect::>() + .await + .unwrap(); + assert_eq!(bytes.concat(), body.as_bytes()); + let events = events(&telemetry); + assert_eq!(events.len(), 1); + assert_eq!(events[0]["usage"]["prompt_tokens"], input); + assert_eq!(events[0]["usage"]["completion_tokens"], output); + assert_eq!(events[0]["usage"]["total_tokens"], json!(total)); + assert_eq!(events[0]["outcome"], "complete"); + assert!(!events[0].to_string().contains("stream-secret-marker")); + } +} + +#[tokio::test] +async fn streaming_http_producer_reports_missing_usage_and_no_replay_after_incomplete_stream() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_raw( + "data: {\"choices\":[{\"delta\":{\"content\":\"private\"}}]}\n\n", + "text/event-stream", + )) + .expect(1) + .mount(&server) + .await; + let telemetry = telemetry(); + let (_, _, stream) = proxy::forward_stream( + Arc::new(WorkloadIdentityAuth::new()), + None, + reqwest::Client::builder().no_proxy().build().unwrap(), + upstream(server.uri(), telemetry.clone()), + "chat/completions", + HeaderMap::new(), + Bytes::from(r#"{"stream":true}"#), + ) + .await + .unwrap(); + let _: Vec<_> = stream.try_collect().await.unwrap(); + let events = events(&telemetry); + assert_eq!(events.len(), 1); + assert_eq!(events[0]["outcome"], "incomplete"); + assert_eq!(events[0]["accepted"], true); + assert_eq!(events[0]["usage_state"], "missing"); + assert!(events[0]["usage"]["total_tokens"].is_null()); +} + +#[tokio::test] +async fn reset_discards_old_inflight_observations_and_ring_overflow_is_explicit() { + let telemetry = telemetry(); + let mut old = telemetry + .begin("responses", "provider", "model", b"{}") + .unwrap(); + telemetry.reset("scope-b".into()); + old.buffered( + 200, + br#"{"usage":{"input_tokens":10,"output_tokens":4,"total_tokens":14}}"#, + ); + assert!(telemetry.snapshot("scope-a", 0).is_none()); + assert_eq!( + telemetry.snapshot("scope-b", 0).unwrap()["events"], + json!([]) + ); + for _ in 0..1100 { + telemetry.record_policy("scope-b", "tool:fetch", false); + } + let snapshot = telemetry.snapshot("scope-b", 0).unwrap(); + assert_eq!(snapshot["dropped_events"], 76); + assert_eq!(snapshot["events"].as_array().unwrap().len(), 256); + assert_eq!(snapshot["has_more"], true); +} + +#[tokio::test] +async fn physical_failover_attempts_keep_actual_provider_identity_without_counting_false_successes() +{ + use kars_inference_router::{ + failover::forward_with_failover, + inference_policy_loader::{InferencePolicySnapshot, ModelPreference, ModelRef}, + }; + let first = MockServer::start().await; + let second = MockServer::start().await; + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(503) + .set_body_json(json!({"error":{"message":"secret-error-body"}})), + ) + .expect(1) + .mount(&first) + .await; + Mock::given(method("POST")).respond_with(ResponseTemplate::new(200).set_body_json(json!({"choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2,"total_tokens":3}}))).expect(1).mount(&second).await; + let state = state("workspace", "uid"); + let mut config = Arc::try_unwrap(state.config) + .ok() + .expect("fixture config is uniquely owned"); + for (id, endpoint) in [("first", first.uri()), ("second", second.uri())] { + config.providers.insert( + id.into(), + ProviderEndpoint { + tag: id.into(), + endpoint, + api_key: None, + }, + ); + } + let telemetry = telemetry(); + let base = upstream("http://127.0.0.1:1".into(), telemetry.clone()); + let policy = InferencePolicySnapshot { + model_preference: Some(ModelPreference { + primary: ModelRef { + provider: "first".into(), + deployment: "same-model".into(), + }, + fallback: vec![ModelRef { + provider: "second".into(), + deployment: "same-model".into(), + }], + }), + ..Default::default() + }; + let result = forward_with_failover( + &state.auth, + None, + &reqwest::Client::builder().no_proxy().build().unwrap(), + &state.deployment_health, + &base, + &config, + &policy, + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await + .unwrap(); + assert_eq!(result.0, StatusCode::OK); + let events = events(&telemetry); + assert_eq!(events.len(), 2); + assert_eq!(events[0]["provider"], "first"); + assert_eq!(events[0]["outcome"], "http_error"); + assert_eq!(events[1]["provider"], "second"); + assert_eq!(events[1]["usage"]["total_tokens"], 3); + assert!( + !serde_json::to_string(&events) + .unwrap() + .contains("secret-error-body") + ); +} + +#[tokio::test] +async fn telemetry_http_requires_scope_and_only_correlated_native_outcomes_are_accepted() { + let state = state("workspace", "uid"); + let scope = state.services.telemetry.cursor().0; + let app = app(state.clone()); + assert_eq!( + send( + &app, + request("GET", "/telemetry/trace", json!({}), None, None, true) + ) + .await + .0, + StatusCode::CONFLICT + ); + let body = json!({"scope_id":scope,"call_id":"call-a","ok":true,"latency_ms":5}); + assert_eq!( + send( + &app, + request("POST", "/telemetry/tool", body.clone(), None, None, true) + ) + .await + .0, + StatusCode::CONFLICT + ); + let (status,evaluation)=send(&app,request("POST","/agt/evaluate", + json!({"agent_id":"test","action":"tool:fetch","context":{ + "scope_id":scope,"tool_call_id":"call-a","tool_name":"spoofed-name","args_preview":"native-secret-marker" + }}),None,None,true)).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(evaluation["allowed"], true); + assert_eq!( + send( + &app, + request("POST", "/telemetry/tool", body.clone(), None, None, true) + ) + .await + .0, + StatusCode::OK + ); + assert_eq!( + send( + &app, + request("POST", "/telemetry/tool", body, None, None, true) + ) + .await + .0, + StatusCode::CONFLICT + ); + let (status, trace) = send( + &app, + request( + "GET", + "/telemetry/trace", + json!({}), + None, + Some(&scope), + true, + ), + ) + .await; + assert_eq!(status, StatusCode::OK); + assert!( + trace["events"] + .as_array() + .unwrap() + .iter() + .any(|event| event["source"] == "harness-reported" && event["name"] == "fetch") + ); + assert!(!trace.to_string().contains("native-secret-marker")); + assert!(!trace.to_string().contains("spoofed-name")); + assert_eq!(trace["durable"], false); +} + +#[tokio::test] +async fn actual_mcp_http_errors_are_not_mistaken_for_success_and_bodies_are_not_retained() { + let telemetry = telemetry(); + let mut state = McpRouteState::standard(); + state.task_telemetry = Some(telemetry.clone()); + let app = routes::mcp_route().with_state(state); + for (id, name) in [(1, "echo"), (2, "missing-tool")] { + let request = Request::post("/mcp") + .header("content-type", "application/json") + .header("accept", "application/json, text/event-stream") + .body(Body::from( + json!({"jsonrpc":"2.0","id":id,"method":"tools/call", + "params":{"name":name,"arguments":{"text":"mcp-secret-body"}}}) + .to_string(), + )) + .unwrap(); + assert_eq!( + app.clone().oneshot(request).await.unwrap().status(), + StatusCode::OK + ); + } + + let events = events(&telemetry); + assert_eq!(events.len(), 2); + assert_eq!(events[0]["ok"], true); + assert_eq!(events[1]["ok"], false); + assert!( + !serde_json::to_string(&events) + .unwrap() + .contains("mcp-secret-body") + ); +} + +#[tokio::test] +async fn mcp_is_error_and_dispatch_transport_failure_are_both_failure_observations() { + use kars_inference_router::mcp::tools::{ + DispatchError, SyncToAsync, ToolCallOutput, ToolCatalog, ToolDefinition, ToolDispatcher, + }; + struct Outcomes(ToolCatalog); + impl ToolDispatcher for Outcomes { + fn catalog(&self) -> &ToolCatalog { + &self.0 + } + fn invoke(&self, name: &str, _args: &Value) -> Result { + if name == "tool_error" { + Ok(ToolCallOutput { + content: vec![], + is_error: true, + }) + } else { + Err(DispatchError::ExecutionFailed { + tool: name.into(), + reason: "transport-secret-marker".into(), + }) + } + } + } + let catalog = ToolCatalog::new( + ["tool_error", "transport_error"] + .map(|name| ToolDefinition { + name: name.into(), + description: "test".into(), + input_schema: json!({"type":"object"}), + }) + .to_vec(), + ) + .unwrap(); + let telemetry = telemetry(); + let mut state = + McpRouteState::standard().with_tools(Arc::new(SyncToAsync::new(Outcomes(catalog)))); + state.task_telemetry = Some(telemetry.clone()); + let app = routes::mcp_route().with_state(state); + for (id, name) in [(1, "tool_error"), (2, "transport_error")] { + let request=Request::post("/mcp").header("content-type","application/json").header("accept","application/json, text/event-stream") + .body(Body::from(json!({"jsonrpc":"2.0","id":id,"method":"tools/call","params":{"name":name,"arguments":{}}}).to_string())).unwrap(); + assert_eq!( + app.clone().oneshot(request).await.unwrap().status(), + StatusCode::OK + ); + } + let events = events(&telemetry); + assert_eq!(events.len(), 2); + assert!( + events + .iter() + .all(|event| event["ok"] == false && event["http_status"] == 200) + ); + assert!( + !serde_json::to_string(&events) + .unwrap() + .contains("transport-secret-marker") + ); +} + +#[tokio::test] +async fn transport_failures_and_cancelled_partial_streams_never_become_success() { + let telemetry = telemetry(); + let result = proxy::forward( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::builder().no_proxy().build().unwrap(), + &upstream("http://127.0.0.1:1".into(), telemetry.clone()), + Method::POST, + "chat/completions", + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await; + assert!(result.is_err()); + assert_eq!(events(&telemetry)[0]["outcome"], "transport_error"); + let mut observation = telemetry + .begin("v1/messages", "provider", "model", b"{}") + .unwrap(); + observation.headers(200); + let chunks = futures::stream::once(async { + Ok::<_, reqwest::Error>(Bytes::from_static( + b"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":12}}}\n\n", + )) + }) + .chain(futures::stream::pending()) + .boxed(); + let mut wrapped = wrap_stream(chunks, Some(observation), true); + wrapped.next().await.unwrap().unwrap(); + drop(wrapped); + let trace = events(&telemetry); + assert_eq!(trace[1]["outcome"], "cancelled"); + assert_eq!(trace[1]["usage"]["prompt_tokens"], 12); + assert!(trace[1]["usage"]["completion_tokens"].is_null()); +} diff --git a/inference-router/tests/multi_provider_guardrails.rs b/inference-router/tests/multi_provider_guardrails.rs index 61601ecd2..6f18201d8 100644 --- a/inference-router/tests/multi_provider_guardrails.rs +++ b/inference-router/tests/multi_provider_guardrails.rs @@ -61,6 +61,7 @@ async fn ollama_provider_forwards_openai_compat_without_auth() { .await; let upstream = UpstreamConfig { + telemetry: None, endpoint: server.uri(), deployment: "llama3.1".into(), sandbox_name: "test-sandbox".into(), @@ -118,6 +119,7 @@ async fn anthropic_provider_forwards_messages_with_router_held_key() { .await; let upstream = UpstreamConfig { + telemetry: None, endpoint: server.uri(), deployment: "claude-sonnet-4-5".into(), sandbox_name: "test-sandbox".into(), diff --git a/inference-router/tests/policy_status_endpoint.rs b/inference-router/tests/policy_status_endpoint.rs index b40f633b3..8c1699b68 100644 --- a/inference-router/tests/policy_status_endpoint.rs +++ b/inference-router/tests/policy_status_endpoint.rs @@ -47,6 +47,7 @@ fn test_state() -> (AppState, Arc) { policy_status.clone(), )); let state = AppState { + services: Default::default(), auth: Arc::new(WorkloadIdentityAuth::new()), copilot: Arc::new(kars_inference_router::copilot_auth::CopilotTokenCache::from_env()), client: reqwest::Client::new(), diff --git a/inference-router/tests/proxy_fake_upstream.rs b/inference-router/tests/proxy_fake_upstream.rs index 16c878a38..4475036d1 100644 --- a/inference-router/tests/proxy_fake_upstream.rs +++ b/inference-router/tests/proxy_fake_upstream.rs @@ -82,6 +82,7 @@ async fn api_key_mode_proxies_chat_completion_with_filter_results() { let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { + telemetry: None, endpoint, deployment: "gpt-4o".to_string(), sandbox_name: "test-sandbox".to_string(), @@ -163,6 +164,7 @@ async fn wi_mode_falls_back_to_imds_and_proxies_embeddings() { let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { + telemetry: None, endpoint, deployment: "text-embedding-3-small".to_string(), sandbox_name: "test-sandbox-wi".to_string(), @@ -236,6 +238,7 @@ async fn upstream_error_status_is_propagated() { let auth = WorkloadIdentityAuth::new(); let (endpoint, client) = azure_endpoint(&azure.base_url()); let upstream = UpstreamConfig { + telemetry: None, endpoint, deployment: "gpt-4o".to_string(), sandbox_name: "test-sandbox-429".to_string(), diff --git a/tests/e2e/governed-services.sh b/tests/e2e/governed-services.sh new file mode 100644 index 000000000..2371013ef --- /dev/null +++ b/tests/e2e/governed-services.sh @@ -0,0 +1,134 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +# The subshell keeps private fixture credentials and its cleanup trap local. +test_governed_services() ( + set +x + local k=(kubectl --context kind-kars-e2e) + local scratch forward_pid="" port="" token agent_token scope request_id new_scope code + local sandbox_uid namespace_uid + scratch=$(mktemp -d) || return 1 + cleanup_governed_smoke() { + if [ -n "$forward_pid" ]; then + kill "$forward_pid" 2>/dev/null || true + wait "$forward_pid" 2>/dev/null || true + fi + rm -f "$scratch/forward.log" "$scratch/response.json" "$scratch/request.json" + rmdir "$scratch" + } + trap cleanup_governed_smoke EXIT + + # Values travel only through the test process and curl's stdin, not argv/logs. + token=$("${k[@]}" get secret router-services-admin -n kars-e2e-test \ + --request-timeout=20s -o go-template='{{index .data "control-token" | base64decode}}') || return 1 + agent_token=$("${k[@]}" get secret router-admin-token -n kars-e2e-test \ + --request-timeout=20s -o go-template='{{index .data "token" | base64decode}}') || return 1 + [ -n "$token" ] && [ -n "$agent_token" ] && [ "$token" != "$agent_token" ] || return 1 + sandbox_uid=$("${k[@]}" get karssandbox e2e-test -n kars-system \ + --request-timeout=20s -o jsonpath='{.metadata.uid}') || return 1 + namespace_uid=$("${k[@]}" get namespace kars-e2e-test \ + --request-timeout=20s -o jsonpath='{.metadata.uid}') || return 1 + [ -n "$sandbox_uid" ] && [ -n "$namespace_uid" ] || return 1 + "${k[@]}" get deployment e2e-test -n kars-e2e-test --request-timeout=20s -o json \ + >"$scratch/response.json" || return 1 + python3 - "$scratch/response.json" <<'PY' || return 1 +import json, sys +pod = json.load(open(sys.argv[1]))["spec"]["template"]["spec"] +volume = next(v for v in pod["volumes"] if v.get("secret", {}).get("secretName") == "router-services-admin") +for container in pod["containers"]: + mounts = [m for m in container.get("volumeMounts", []) if m["name"] == volume["name"]] + if container["name"] == "inference-router": + assert len(mounts) == 1 and mounts[0]["mountPath"] == "/etc/kars/services" + assert mounts[0]["readOnly"] is True + else: + assert not mounts +PY + + "${k[@]}" port-forward --address 127.0.0.1 service/e2e-test -n kars-e2e-test :8443 \ + >"$scratch/forward.log" 2>&1 & + forward_pid=$! + local deadline=$(($(date +%s) + 30)) + while [ "$(date +%s)" -lt "$deadline" ]; do + kill -0 "$forward_pid" 2>/dev/null || return 1 + port=$(sed -n 's/^Forwarding from 127\.0\.0\.1:\([0-9]*\) ->.*/\1/p' "$scratch/forward.log" | head -1) + [ -z "$port" ] || break + sleep 1 + done + [ -n "$port" ] || return 1 + + service_request() { + local expected="$1" method="$2" path="$3" bearer="${4:-}" + local args=(--disable --silent --show-error --noproxy 127.0.0.1 --connect-timeout 5 --max-time 15 + --config - --request "$method" --url "http://127.0.0.1:$port$path" + --output "$scratch/response.json" --write-out '%{http_code}') + if [ "$method" = POST ]; then + args+=(--header 'Content-Type: application/json' --data-binary "@$scratch/request.json") + fi + code=$( + { [ -z "$bearer" ] || printf 'header = "Authorization: Bearer %s"\n' "$bearer"; } \ + | curl "${args[@]}" + ) || return 1 + if [ "$code" != "$expected" ]; then + printf 'Governed service %s %s returned %s, expected %s\n' "$method" "$path" "$code" "$expected" >&2 + return 1 + fi + } + service_body() { + python3 - "$scratch/request.json" "$@" <<'PY' +import json, sys +args = sys.argv[2:] +assert len(args) % 2 == 0 +with open(sys.argv[1], "w") as output: + json.dump(dict(zip(args[::2], args[1::2])), output) +PY + } + response_field() { + python3 - "$scratch/response.json" "$1" <<'PY' +import json, sys +value = json.load(open(sys.argv[1])) +for field in sys.argv[2].split("."): + value = value[field] +assert isinstance(value, str) and value +print(value) +PY + } + + service_request 401 GET /internal/access-requests || return 1 + service_request 401 GET /internal/access-requests "$agent_token" || return 1 + service_request 200 GET /internal/access-requests "$token" || return 1 + python3 - "$scratch/response.json" "$sandbox_uid" "$namespace_uid" <<'PY' || return 1 +import json, sys +response = json.load(open(sys.argv[1])) +identity = response["scope"]["identity"] +assert identity["sandbox"] == {"namespace": "kars-system", "name": "e2e-test", "uid": sys.argv[2]} +assert identity["namespace_uid"] == sys.argv[3] +assert identity.get("task") is None +assert response["enforcement_changed"] is False +PY + scope=$(response_field scope.id) || return 1 + service_body scope_id "$scope" kind egress target example.invalid reason fixture || return 1 + service_request 202 POST /v1/access-request || return 1 + request_id=$(response_field request.request_id) || return 1 + service_body scope_id "$scope" request_id "$request_id" verdict approved || return 1 + service_request 401 POST /internal/access-requests/decision "$agent_token" || return 1 + service_request 200 POST /internal/access-requests/decision "$token" || return 1 + python3 - "$scratch/response.json" <<'PY' || return 1 +import json, sys +response = json.load(open(sys.argv[1])) +assert response["request"]["status"] == "approved" +assert response["enforcement_changed"] is False +PY + service_body scope_id "$scope" assignment_id fixture-assignment || return 1 + service_request 401 POST /internal/access-requests/reset "$agent_token" || return 1 + service_request 200 POST /internal/access-requests/reset "$token" || return 1 + new_scope=$(response_field scope.id) || return 1 + [ "$new_scope" != "$scope" ] || return 1 + [ "$(response_field scope.identity.sandbox.uid)" = "$sandbox_uid" ] || return 1 + service_body scope_id "$scope" request_id "$request_id" verdict approved || return 1 + service_request 409 POST /internal/access-requests/decision "$token" || return 1 + service_body scope_id "$scope" kind egress target example.invalid || return 1 + service_request 409 POST /v1/access-request || return 1 + service_request 200 GET /telemetry/cursor || return 1 + [ "$(response_field scope_id)" = "$new_scope" ] || return 1 +) diff --git a/tests/e2e/run.sh b/tests/e2e/run.sh index 9e9c58cce..e1b26dee7 100755 --- a/tests/e2e/run.sh +++ b/tests/e2e/run.sh @@ -3017,6 +3017,7 @@ EOF source "$SCRIPT_DIR/namespace-ownership.sh" source "$SCRIPT_DIR/credential-sources.sh" +source "$SCRIPT_DIR/governed-services.sh" main() { echo "" @@ -3046,6 +3047,11 @@ main() { test_sandbox_namespace_labels || true test_sandbox_deployment_exists || true test_sandbox_pod_starts || true + if test_governed_services; then + pass "Governed services isolate operator credentials, bind live identity, and reject stale scopes" + else + fail "Governed service authentication and scope lifecycle gate failed" + fi test_sandbox_networkpolicy_denies_ingress || true test_sandbox_suspended_lifecycle || true test_secondary_resource_watch || true From 11f4224d7c830b2c57878e356d1b4b20b13751e0 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Tue, 8 Sep 2026 06:01:22 +0200 Subject: [PATCH 2/4] fix(router): fence approved dispatch and classify semantic failures Coordinate dispatch claims with cancellation/reset after awaited policy checks and retain claims through response handling. Distinguish accepted Responses/OpenAI semantic errors from completed generations without replay or byte changes. Keep the separate HIGH SRE credential-privacy prerequisite explicitly blocked. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- docs/governed-services.md | 21 +- .../2026-09-08-governed-router-services.md | 24 ++ inference-router/src/access_request.rs | 108 ++++++++- .../src/access_request/dispatch_tests.rs | 225 ++++++++++++++++++ inference-router/src/governed_services.rs | 43 +++- inference-router/src/routes/egress.rs | 23 ++ .../src/task_telemetry/observe.rs | 10 + .../src/task_telemetry/observe/stream.rs | 31 ++- inference-router/src/task_telemetry/parse.rs | 34 ++- .../tests/governed_egress_wait.rs | 155 ++++++++++++ inference-router/tests/governed_telemetry.rs | 117 +++++++++ tests/e2e/run.sh | 2 +- 12 files changed, 784 insertions(+), 9 deletions(-) create mode 100644 inference-router/src/access_request/dispatch_tests.rs diff --git a/docs/governed-services.md b/docs/governed-services.md index 6e1fae172..4fa7c90d1 100644 --- a/docs/governed-services.md +++ b/docs/governed-services.md @@ -4,6 +4,12 @@ These APIs provide an in-process capability-request queue and bounded router telemetry. They do **not** deliver assignments, run agents, create approvals, grant capabilities, install resources, or provide a durable execution ledger. +**Publication blocker:** the legacy SRE agent's Kubernetes permissions can read +the new control Secret despite its router-only mount. This candidate is not +ready for deployment. A separate operator-authorized SRE identity/migration +prerequisite must close that API credential-access path; changing mounts alone +does not establish operator-only authority. + ## Identity and operator authentication Each router has a unique process-instance scope. On Kubernetes the controller @@ -60,7 +66,8 @@ operator decision separately; cancelling a wait does not revoke policy. Reset requires the current `scope_id` as a compare-and-swap condition and can include an `assignment_id` correlation label. It returns a fresh scope with the same trusted identity, clears request/telemetry state, and cancels old -waiters. An assignment label is not evidence that a runtime executed anything. +waiters. Reset returns conflict while an approved dispatch is already claimed. +An assignment label is not evidence that a runtime executed anything. Reset neither changes policy nor resets existing token-budget counters. ## Policy-gated egress waits @@ -77,6 +84,13 @@ existing signed-allowlist/controller path remains responsible for enforcement; private-address checks, normal egress filters, and redirect behavior still apply. Denial, cancellation, expiry, and reset never bypass those filters. +After the asynchronous policy check, the waiter revalidates the request. An +atomic dispatch claim then coordinates with cancellation and reset immediately +before sending. Cancellation acknowledged before that claim prevents dispatch; +after the claim it returns conflict rather than promising prevention. The claim +stays active through response handling and releases on completion or caller +cancellation. It does not assert that the upstream accepted the request. + Transparent proxy traffic records bounded denied requests and policy observations, but keeps its existing immediate-denial/retry behavior. It does not acquire a new implicit human-approval wait. @@ -101,6 +115,11 @@ are not successes. Metadata parsing is bounded, and `partial_observation` identifies incomplete observations. Streams are forwarded unchanged and are never replayed by telemetry after acceptance. +HTTP acceptance is distinct from semantic completion: buffered Responses with +`failed` or `incomplete` status and OpenAI error frames followed by `[DONE]` do +not become completed generations. Their accepted status, original response +bytes and available usage remain intact, without retaining provider error text. + Only model identifiers, tool names/IDs, status, timing, and reported usage are retained. Prompts, assistant text, URLs, headers, arguments, and result bodies are not retained in telemetry. A model-proposed tool has no successful outcome diff --git a/docs/security-audits/2026-09-08-governed-router-services.md b/docs/security-audits/2026-09-08-governed-router-services.md index 07afb2140..7609cd672 100644 --- a/docs/security-audits/2026-09-08-governed-router-services.md +++ b/docs/security-audits/2026-09-08-governed-router-services.md @@ -3,6 +3,16 @@ Date: 2026-09-08 Status: Implementation candidate; independent review and sign-offs pending. +## Known blocking finding + +The supported legacy SRE installation grants its agent-held Kubernetes +credential cluster-wide Secret reads. It can therefore obtain the purportedly +operator-only service token through the API. Router-only mounts do not close +that path. The HIGH remains unresolved in this candidate; a separately reviewed +operator-controlled UID registration/migration and SRE credential boundary is +required before deployment or merge. Legacy grants cannot be retired merely by +trusting names, labels or ownership-looking annotations. + ## Scope Bounded access-request services, operator-only reset/inspection/decisions, @@ -60,6 +70,20 @@ Local qualification completed before independent review: - Static A2A-isolation, null-provider, and tracked-source copyright checks passed; new candidate files are also checked for headers and module limits. +Two MEDIUM review findings have been repaired and source-reviewed separately: +the cancellation/expiry race across an asynchronous policy check, and semantic +model errors reported as completed generations. Dispatch now has a mutex- +coordinated claim boundary; cancellation/reset cannot acknowledge prevention +after that claim. HTTP acceptance remains distinct from semantic failure or +incompleteness, with no response rewriting or replay. + +Repair qualification passed 1,071 router unit tests and 29 governed-service +integration tests, including deterministic cancellation/reset/expiry/shutdown +interleavings and accepted failed/incomplete/error-plus-DONE responses. Strict +router all-target Clippy passed; the focused 12 telemetry integrations passed +again after a type-safe initializer adjustment. These results do not resolve +the separate HIGH SRE credential-access issue. + The existing disposable Kind harness now includes the actual deployed router: it checks that the control Secret is mounted only in the router, rejects the agent-visible token even through loopback port-forwarding, compares returned diff --git a/inference-router/src/access_request.rs b/inference-router/src/access_request.rs index ed2fa65ae..23b28582c 100644 --- a/inference-router/src/access_request.rs +++ b/inference-router/src/access_request.rs @@ -15,6 +15,10 @@ use tokio_util::sync::CancellationToken; pub mod scope; pub use scope::{Identity, Scope}; +#[cfg(test)] +#[path = "access_request/dispatch_tests.rs"] +mod dispatch_tests; + const CAPACITY: usize = 64; const REQUESTS_PER_MINUTE: u32 = 32; const TTL: Duration = Duration::from_secs(15 * 60); @@ -27,6 +31,7 @@ pub enum Status { Denied, Cancelled, Expired, + DispatchClaimed, } #[derive(Clone, Deserialize)] @@ -56,6 +61,7 @@ pub struct Entry { pub count: u32, pub first_seen_unix: u64, pub expires_at_unix: u64, + pub dispatch_active: bool, #[serde(skip)] expires: Instant, } @@ -81,6 +87,7 @@ struct Inner { window: Instant, rate: u32, cancellation: CancellationToken, + active_dispatches: usize, } pub struct AccessRequestBuffer { @@ -91,6 +98,31 @@ pub struct AccessRequestBuffer { ttl: Duration, } +/// The linearization boundary for an outgoing operation, not evidence that an +/// upstream accepted it. Cancel/reset cannot acknowledge prevention while held. +#[must_use] +pub struct DispatchClaim<'a> { + buffer: &'a AccessRequestBuffer, + request_id: Option, +} + +impl Drop for DispatchClaim<'_> { + fn drop(&mut self) { + if let Ok(mut inner) = self.buffer.inner.lock() { + inner.active_dispatches = inner.active_dispatches.saturating_sub(1); + if let Some(id) = &self.request_id + && let Some(entry) = inner + .requests + .iter_mut() + .find(|entry| &entry.request_id == id) + { + entry.dispatch_active = false; + } + } + self.buffer.notify(); + } +} + impl AccessRequestBuffer { pub fn new(identity: Identity) -> Self { Self::with_limits(identity, CAPACITY, REQUESTS_PER_MINUTE, TTL) @@ -116,6 +148,7 @@ impl AccessRequestBuffer { window: Instant::now(), rate: 0, cancellation: CancellationToken::new(), + active_dispatches: 0, }), changed: watch::channel(0).0, capacity: capacity.clamp(1, 256), @@ -185,7 +218,7 @@ impl AccessRequestBuffer { if let Some(index) = inner .requests .iter() - .position(|entry| entry.status != Status::Pending) + .position(|entry| entry.status != Status::Pending && !entry.dispatch_active) { inner.requests.remove(index); } else { @@ -210,6 +243,7 @@ impl AccessRequestBuffer { count: 1, first_seen_unix: now, expires_at_unix: now.saturating_add(self.ttl.as_secs()), + dispatch_active: false, expires: Instant::now() + self.ttl, }; inner.requests.push_back(entry.clone()); @@ -240,6 +274,9 @@ impl AccessRequestBuffer { .iter_mut() .find(|entry| entry.request_id == id) .ok_or(Error::Missing)?; + if entry.status == Status::DispatchClaimed { + return Err(Error::Terminal); + } if entry.expires <= Instant::now() { return Err(Error::Expired); } @@ -271,6 +308,9 @@ impl AccessRequestBuffer { } let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; Self::check(&inner, expected_scope)?; + if inner.active_dispatches > 0 { + return Err(Error::Terminal); + } inner.generation = inner.generation.checked_add(1).ok_or(Error::Unavailable)?; let instance = inner.scope.id.split(':').next().ok_or(Error::Unavailable)?; inner.scope.id = format!("{instance}:{}", inner.generation); @@ -297,6 +337,72 @@ impl AccessRequestBuffer { .find(|entry| entry.request_id == id) .ok_or(Error::Missing) } + + fn dispatch_ready( + inner: &mut Inner, + scope: &str, + id: Option<&str>, + shutdown: &CancellationToken, + ) -> Result<(), Error> { + Self::check(inner, scope)?; + if shutdown.is_cancelled() { + return Err(Error::Unavailable); + } + Self::expire(inner); + if let Some(id) = id { + let entry = inner + .requests + .iter() + .find(|entry| entry.request_id == id) + .ok_or(Error::Missing)?; + if entry.status == Status::Expired { + return Err(Error::Expired); + } + if entry.status != Status::Approved { + return Err(Error::Terminal); + } + } + Ok(()) + } + + /// Recheck after an asynchronous policy lookup. This is deliberately not a + /// dispatch permit: callers must still claim at the actual send boundary. + pub fn validate_dispatch( + &self, + scope: &str, + id: &str, + shutdown: &CancellationToken, + ) -> Result<(), Error> { + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::dispatch_ready(&mut inner, scope, Some(id), shutdown) + } + + pub fn claim_dispatch( + &self, + scope: &str, + id: Option<&str>, + shutdown: &CancellationToken, + ) -> Result, Error> { + let mut inner = self.inner.lock().map_err(|_| Error::Unavailable)?; + Self::dispatch_ready(&mut inner, scope, id, shutdown)?; + let active = inner.active_dispatches.checked_add(1).ok_or(Error::Full)?; + if let Some(id) = id { + let entry = inner + .requests + .iter_mut() + .find(|entry| entry.request_id == id) + .ok_or(Error::Missing)?; + entry.status = Status::DispatchClaimed; + entry.dispatch_active = true; + } + inner.active_dispatches = active; + drop(inner); + self.notify(); + Ok(DispatchClaim { + buffer: self, + request_id: id.map(str::to_string), + }) + } } fn validate(request: &mut Request) -> Result<(), Error> { diff --git a/inference-router/src/access_request/dispatch_tests.rs b/inference-router/src/access_request/dispatch_tests.rs new file mode 100644 index 000000000..bb73fff89 --- /dev/null +++ b/inference-router/src/access_request/dispatch_tests.rs @@ -0,0 +1,225 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +use super::*; +use crate::governed_services::GovernedServices; +use std::sync::Arc; +use tokio::sync::{RwLock, oneshot}; +use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; + +fn approved(port: u16) -> (Arc, String, String) { + let services = Arc::new(GovernedServices::new(Identity::standalone("test"), None)); + let scope = services.requests.scope().unwrap().id; + let (entry, _) = services + .requests + .record(Request { + scope_id: scope.clone(), + kind: "egress".into(), + target: "dispatch.example".into(), + reason: String::new(), + tier: None, + port: Some(port), + }) + .unwrap(); + services + .requests + .transition(&scope, &entry.request_id, Status::Approved) + .unwrap(); + (services, scope, entry.request_id) +} + +#[tokio::test] +async fn cancellation_expiry_reset_and_shutdown_during_policy_await_prevent_post_dispatch() { + for action in ["cancel", "expire", "reset", "shutdown"] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200)) + .expect(0) + .mount(&server) + .await; + let address = reqwest::Url::parse(&server.uri()) + .unwrap() + .socket_addrs(|| None) + .unwrap()[0]; + let target = format!("http://dispatch.example:{}/side-effect", address.port()); + let client = reqwest::Client::builder() + .no_proxy() + .resolve("dispatch.example", address) + .build() + .unwrap(); + let (services, scope, id) = approved(address.port()); + let policy = Arc::new(RwLock::new(false)); + let mut held = policy.write().await; + let (entered, observing) = oneshot::channel(); + let entered = Arc::new(Mutex::new(Some(entered))); + let worker = services.clone(); + let worker_scope = scope.clone(); + let worker_id = id.clone(); + let policy_for_worker = policy.clone(); + let task = tokio::spawn(async move { + worker + .wait_for_egress_check( + &worker_scope, + &worker_id, + &target, + "test", + Duration::from_secs(5), + || { + let policy = policy_for_worker.clone(); + let entered = entered.clone(); + async move { + if let Some(sender) = entered.lock().unwrap().take() { + let _ = sender.send(()); + } + if *policy.read().await { + Ok(()) + } else { + Err("blocked".to_string()) + } + } + }, + ) + .await?; + let _claim = worker.claim_egress_dispatch(&worker_scope, Some(&worker_id))?; + client + .post(&target) + .body("side-effect") + .send() + .await + .map_err(|_| Error::Unavailable)?; + Ok::<(), Error>(()) + }); + observing.await.unwrap(); + let expected = match action { + "cancel" => { + let cancelled = services + .requests + .transition(&scope, &id, Status::Cancelled) + .unwrap(); + assert_eq!(cancelled.status, Status::Cancelled); + Error::Terminal + } + "expire" => { + // Fixture mutation advances expiry only after the approved + // waiter is demonstrably blocked in its asynchronous check. + services + .requests + .inner + .lock() + .unwrap() + .requests + .iter_mut() + .find(|entry| entry.request_id == id) + .unwrap() + .expires = Instant::now(); + Error::Expired + } + "reset" => { + services + .reset(&scope, Some("next-assignment".into())) + .unwrap(); + Error::StaleScope + } + _ => { + services.shutdown.cancel(); + Error::Unavailable + } + }; + *held = true; + drop(held); + assert_eq!(task.await.unwrap(), Err(expected), "{action}"); + assert!( + server.received_requests().await.unwrap().is_empty(), + "{action}" + ); + } +} + +#[test] +fn acknowledged_cancel_or_reset_before_claim_wins_and_claims_are_single_use() { + let (services, scope, id) = approved(443); + services + .requests + .validate_dispatch(&scope, &id, &services.shutdown) + .unwrap(); + services + .requests + .transition(&scope, &id, Status::Cancelled) + .unwrap(); + assert!(matches!( + services.claim_egress_dispatch(&scope, Some(&id)), + Err(Error::Terminal) + )); + + let (services, scope, id) = approved(443); + services.reset(&scope, Some("next".into())).unwrap(); + assert!(matches!( + services.claim_egress_dispatch(&scope, Some(&id)), + Err(Error::StaleScope) + )); + + let (services, scope, id) = approved(443); + let claim = services.claim_egress_dispatch(&scope, Some(&id)).unwrap(); + assert_eq!( + services.requests.entry(&scope, &id).unwrap().status, + Status::DispatchClaimed + ); + assert!( + services + .requests + .entry(&scope, &id) + .unwrap() + .dispatch_active + ); + assert!(matches!( + services.requests.transition(&scope, &id, Status::Cancelled), + Err(Error::Terminal) + )); + assert!(matches!( + services.claim_egress_dispatch(&scope, Some(&id)), + Err(Error::Terminal) + )); + assert!(matches!(services.reset(&scope, None), Err(Error::Terminal))); + drop(claim); + assert!( + !services + .requests + .entry(&scope, &id) + .unwrap() + .dispatch_active + ); + assert!(matches!( + services.requests.transition(&scope, &id, Status::Cancelled), + Err(Error::Terminal) + )); + services.reset(&scope, None).unwrap(); +} + +#[test] +fn expiry_and_shutdown_between_readiness_and_claim_still_prevent_dispatch() { + for expired in [true, false] { + let (services, scope, id) = approved(443); + services + .requests + .validate_dispatch(&scope, &id, &services.shutdown) + .unwrap(); + if expired { + services.requests.inner.lock().unwrap().requests[0].expires = Instant::now(); + } else { + services.shutdown.cancel(); + } + assert!(matches!( + services.claim_egress_dispatch(&scope, Some(&id)), + Err(Error::Expired | Error::Unavailable) + )); + } +} + +#[test] +fn scope_only_dispatch_also_prevents_successful_reset_until_released() { + let (services, scope, _) = approved(443); + let claim = services.claim_egress_dispatch(&scope, None).unwrap(); + assert!(matches!(services.reset(&scope, None), Err(Error::Terminal))); + drop(claim); + services.reset(&scope, None).unwrap(); +} diff --git a/inference-router/src/governed_services.rs b/inference-router/src/governed_services.rs index 298923f5a..a723b0056 100644 --- a/inference-router/src/governed_services.rs +++ b/inference-router/src/governed_services.rs @@ -170,6 +170,25 @@ impl GovernedServices { sandbox: &str, timeout: Duration, ) -> Result<(), Error> { + self.wait_for_egress_check(scope, id, target, sandbox, timeout, || { + blocklist.check_egress(target, sandbox) + }) + .await + } + + pub(crate) async fn wait_for_egress_check( + &self, + scope: &str, + id: &str, + target: &str, + sandbox: &str, + timeout: Duration, + mut check_policy: F, + ) -> Result<(), Error> + where + F: FnMut() -> Fut, + Fut: std::future::Future>, + { let url = reqwest::Url::parse(target).map_err(|_| Error::Invalid)?; let host = url.host_str().ok_or(Error::Invalid)?; let port = url.port_or_known_default().ok_or(Error::Invalid)?; @@ -194,12 +213,21 @@ impl GovernedServices { return Err(Error::Invalid); } match entry.status { - Status::Denied | Status::Cancelled => return Err(Error::Terminal), + Status::Denied | Status::Cancelled | Status::DispatchClaimed => { + return Err(Error::Terminal); + } Status::Expired => return Err(Error::Expired), Status::Approved => { // The decision API never changes enforcement. Only the // normal signed-policy path can make this check pass. - if blocklist.check_egress(target, sandbox).await.is_ok() { + let allowed = tokio::select! { + result = check_policy() => result.is_ok(), + _ = cancellation.cancelled() => return Err(Error::StaleScope), + _ = self.shutdown.cancelled() => return Err(Error::Unavailable), + _ = changed.changed() => continue, + }; + if allowed { + self.requests.validate_dispatch(scope, id, &self.shutdown)?; return Ok(()); } } @@ -217,4 +245,15 @@ impl GovernedServices { .await .map_err(|_| Error::WaitTimeout)? } + + /// Claim immediately before creating/polling the outgoing send future. + /// This shares the request mutex with cancellation and scope reset. + pub fn claim_egress_dispatch( + &self, + scope: &str, + request_id: Option<&str>, + ) -> Result, Error> { + self.requests + .claim_dispatch(scope, request_id, &self.shutdown) + } } diff --git a/inference-router/src/routes/egress.rs b/inference-router/src/routes/egress.rs index 3178d6835..9680dbd65 100644 --- a/inference-router/src/routes/egress.rs +++ b/inference-router/src/routes/egress.rs @@ -176,6 +176,7 @@ async fn egress_fetch( } let sandbox: &str = &state.sandbox_name; + let mut dispatch_request = None; // Check egress access: blocklist → allowlist (Strict denies the rest). if let Err(reason) = state.blocklist.check_egress(url, sandbox).await { @@ -228,6 +229,7 @@ async fn egress_fetch( ); return super::access_request::error(error); } + dispatch_request = Some(entry.request_id); } else { state.services.telemetry.record_router_tool( &telemetry_scope, @@ -302,6 +304,27 @@ async fn egress_fetch( const MAX_RESPONSE_BYTES: usize = 2 * 1024 * 1024; // 2 MB + // No await separates this atomic claim from starting the send. A successful + // cancel/reset before the claim wins; afterwards prevention is not promised. + let _dispatch_claim = if wait_ms > 0 { + if state.blocklist.check_egress(url, sandbox).await.is_err() { + return errors::flat( + StatusCode::FORBIDDEN, + "Egress policy no longer permits this operation", + ) + .into_response(); + } + match state + .services + .claim_egress_dispatch(&telemetry_scope, dispatch_request.as_deref()) + { + Ok(claim) => Some(claim), + Err(error) => return super::access_request::error(error), + } + } else { + None + }; + match request .timeout(std::time::Duration::from_secs(30)) .send() diff --git a/inference-router/src/task_telemetry/observe.rs b/inference-router/src/task_telemetry/observe.rs index 2cb5dbaef..4267e89e4 100644 --- a/inference-router/src/task_telemetry/observe.rs +++ b/inference-router/src/task_telemetry/observe.rs @@ -87,6 +87,16 @@ impl Observation { return; } self.finished = true; + // HTTP acceptance and a complete body do not imply a successful model + // outcome (Responses may return failed/incomplete at HTTP 200). + let outcome = if outcome == "complete" { + match parsed.semantic { + Some(semantic) => semantic.label(), + None => outcome, + } + } else { + outcome + }; let usage_state = match (parsed.usage.prompt_tokens, parsed.usage.completion_tokens) { (Some(_), Some(_)) => "present", (None, None) => "missing", diff --git a/inference-router/src/task_telemetry/observe/stream.rs b/inference-router/src/task_telemetry/observe/stream.rs index 04361ea50..9433309fb 100644 --- a/inference-router/src/task_telemetry/observe/stream.rs +++ b/inference-router/src/task_telemetry/observe/stream.rs @@ -15,6 +15,7 @@ pub(super) struct Accumulator { parsed: Parsed, terminal: bool, pub failed: bool, + incomplete: bool, tools: BTreeMap<(u64, u64), (String, String)>, } impl Accumulator { @@ -28,6 +29,7 @@ impl Accumulator { parsed: Parsed::default(), terminal: false, failed: false, + incomplete: false, tools: BTreeMap::new(), } } @@ -68,6 +70,14 @@ impl Accumulator { self.parsed.partial = true; return; }; + if frame.lines().any(|line| { + line.strip_prefix("event:") + .is_some_and(|name| name.trim() == "error") + }) { + self.failed = true; + self.terminal = true; + self.parsed.semantic = Some(parse::SemanticOutcome::Failed); + } let data = frame .lines() .filter_map(|line| line.strip_prefix("data:").map(str::trim)) @@ -89,6 +99,12 @@ impl Accumulator { } match self.shape { Shape::OpenAi => { + if parse::semantic_outcome(&value, self.shape).is_some() { + self.failed = true; + self.terminal = true; + self.parsed.semantic = Some(parse::SemanticOutcome::Failed); + return; + } if let Some(choices) = value["choices"].as_array() { for choice in choices.iter().take(parse::MAX_TOOLS) { if let Some(reason) = choice["finish_reason"] @@ -158,12 +174,16 @@ impl Accumulator { Shape::Responses => match value["type"].as_str() { Some("response.completed" | "response.failed" | "response.incomplete") => { self.terminal = true; - self.failed = value["type"] != "response.completed"; + self.failed |= value["type"] == "response.failed"; + self.incomplete |= value["type"] == "response.incomplete"; let parsed = parse::response(&value["response"], Shape::Responses); + self.failed |= parsed.semantic == Some(parse::SemanticOutcome::Failed); + self.incomplete |= parsed.semantic == Some(parse::SemanticOutcome::Incomplete); self.parsed.usage = parsed.usage; self.parsed.finish = parsed.finish; self.parsed.tools = parsed.tools; - self.parsed.partial |= parsed.partial || self.failed; + self.parsed.semantic = parsed.semantic; + self.parsed.partial |= parsed.partial; } Some("error") => { self.terminal = true; @@ -174,7 +194,7 @@ impl Accumulator { } } pub(super) fn complete(&self) -> bool { - self.terminal && !self.failed + self.terminal && !self.failed && !self.incomplete } pub(super) fn finish(&mut self) -> Parsed { if !self.sse && !self.discarding { @@ -189,6 +209,11 @@ impl Accumulator { if !self.terminal { self.parsed.partial = true; } + if self.failed { + self.parsed.semantic = Some(parse::SemanticOutcome::Failed); + } else if self.incomplete { + self.parsed.semantic = Some(parse::SemanticOutcome::Incomplete); + } for (_, (id, name)) in std::mem::take(&mut self.tools) { add_tool(&mut self.parsed, &Value::String(id), &Value::String(name)); } diff --git a/inference-router/src/task_telemetry/parse.rs b/inference-router/src/task_telemetry/parse.rs index a3942cdd4..bbc854c45 100644 --- a/inference-router/src/task_telemetry/parse.rs +++ b/inference-router/src/task_telemetry/parse.rs @@ -40,6 +40,35 @@ pub struct Parsed { pub finish: Option, pub tools: Vec<(String, String)>, pub partial: bool, + pub semantic: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SemanticOutcome { + Failed, + Incomplete, +} + +impl SemanticOutcome { + pub fn label(self) -> &'static str { + match self { + Self::Failed => "upstream_error", + Self::Incomplete => "incomplete", + } + } +} + +pub fn semantic_outcome(value: &Value, shape: Shape) -> Option { + if matches!(shape, Shape::Responses) { + match value.get("status").and_then(Value::as_str) { + Some("failed") => return Some(SemanticOutcome::Failed), + Some("incomplete") => return Some(SemanticOutcome::Incomplete), + _ => {} + } + } + (value.get("error").is_some_and(|error| !error.is_null()) + || value.get("type").and_then(Value::as_str) == Some("error")) + .then_some(SemanticOutcome::Failed) } pub fn merge_usage(target: &mut Usage, usage: &Value) { @@ -81,7 +110,10 @@ pub fn add_tool(parsed: &mut Parsed, id: &Value, name: &Value) { } pub fn response(value: &Value, shape: Shape) -> Parsed { - let mut parsed = Parsed::default(); + let mut parsed = Parsed { + semantic: semantic_outcome(value, shape), + ..Default::default() + }; if let Some(usage) = value.get("usage") { merge_usage(&mut parsed.usage, usage); } diff --git a/inference-router/tests/governed_egress_wait.rs b/inference-router/tests/governed_egress_wait.rs index 7a3aadb1f..5d9fda648 100644 --- a/inference-router/tests/governed_egress_wait.rs +++ b/inference-router/tests/governed_egress_wait.rs @@ -219,6 +219,161 @@ async fn ordinary_denied_fetch_is_still_immediate_and_private_targets_never_ente assert!(upstream.received_requests().await.unwrap().is_empty()); } +#[tokio::test] +async fn cancel_after_readiness_but_before_dispatch_claim_prevents_a_post() { + let (state, app, upstream, scope, url) = setup().await; + let port = reqwest::Url::parse(&url) + .unwrap() + .port_or_known_default() + .unwrap(); + let entry = queue(&app, &scope, "egress", "approved.example", Some(port)).await; + let id = entry["request_id"].as_str().unwrap(); + assert_eq!( + decision(&app, &scope, id, "approved").await.0, + StatusCode::OK + ); + state + .blocklist + .replace_allowlist(vec!["approved.example".into()]) + .await; + state + .services + .wait_for_egress( + &state.blocklist, + &scope, + id, + &url, + "test", + Duration::from_secs(1), + ) + .await + .unwrap(); + let cancelled = send( + &app, + request( + "POST", + &format!("/v1/access-requests/{id}/cancel"), + json!({"scope_id":scope}), + None, + None, + true, + ), + ) + .await; + assert_eq!(cancelled.0, StatusCode::OK); + assert_eq!(cancelled.1["status"], "cancelled"); + assert!(matches!( + state.services.claim_egress_dispatch(&scope, Some(id)), + Err(kars_inference_router::access_request::Error::Terminal) + )); + assert!(upstream.received_requests().await.unwrap().is_empty()); +} + +#[tokio::test] +async fn cancel_and_reset_do_not_acknowledge_prevention_after_dispatch_has_started() { + let (state, app, upstream, scope, url) = setup().await; + Mock::given(method("POST")) + .respond_with( + ResponseTemplate::new(200) + .set_body_string("accepted") + .set_delay(Duration::from_millis(300)), + ) + .expect(1) + .mount(&upstream) + .await; + let waiting=tokio::spawn(app.clone().oneshot(request("POST","/egress/fetch", + json!({"url":url,"method":"POST","body":"side-effect","scope_id":scope,"wait_for_approval_ms":2000}), + None,None,true))); + let entry = pending(&app).await; + let id = entry["request_id"].as_str().unwrap(); + assert_eq!( + decision(&app, &scope, id, "approved").await.0, + StatusCode::OK + ); + state + .blocklist + .replace_allowlist(vec!["approved.example".into()]) + .await; + for _ in 0..100 { + if !upstream.received_requests().await.unwrap().is_empty() { + break; + } + tokio::time::sleep(Duration::from_millis(2)).await; + } + assert_eq!(upstream.received_requests().await.unwrap().len(), 1); + assert_eq!( + send( + &app, + request( + "POST", + &format!("/v1/access-requests/{id}/cancel"), + json!({"scope_id":scope}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!( + send( + &app, + request( + "POST", + "/internal/access-requests/reset", + json!({"scope_id":scope,"assignment_id":"next"}), + Some(CONTROL), + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!(waiting.await.unwrap().unwrap().status(), StatusCode::OK); + let completed = state.services.requests.entry(&scope, id).unwrap(); + assert_eq!( + completed.status, + kars_inference_router::access_request::Status::DispatchClaimed + ); + assert!(!completed.dispatch_active); + assert_eq!( + send( + &app, + request( + "POST", + &format!("/v1/access-requests/{id}/cancel"), + json!({"scope_id":scope}), + None, + None, + true + ) + ) + .await + .0, + StatusCode::CONFLICT + ); + assert_eq!( + send( + &app, + request( + "POST", + "/internal/access-requests/reset", + json!({"scope_id":scope,"assignment_id":"next"}), + Some(CONTROL), + None, + true + ) + ) + .await + .0, + StatusCode::OK + ); +} + #[tokio::test] async fn wait_rejects_wrong_port_request_and_shutdown_cancels_waiters() { let (state, app, _upstream, scope, url) = setup().await; diff --git a/inference-router/tests/governed_telemetry.rs b/inference-router/tests/governed_telemetry.rs index 2aac70446..6fb55ec7a 100644 --- a/inference-router/tests/governed_telemetry.rs +++ b/inference-router/tests/governed_telemetry.rs @@ -86,6 +86,123 @@ async fn buffered_http_producer_observes_usage_but_never_api_bodies_or_tool_argu } } +#[tokio::test] +async fn buffered_http_200_responses_distinguish_failed_and_incomplete_model_outcomes() { + for (status, outcome) in [("failed", "upstream_error"), ("incomplete", "incomplete")] { + let server = MockServer::start().await; + let body = json!({ + "status":status,"error":{"message":"credential-secret-marker"}, + "incomplete_details":{"reason":"private-details-marker"}, + "usage":if status=="failed"{json!({"input_tokens":7})}else{Value::Null}, + "output":[], + }); + Mock::given(method("POST")) + .respond_with(ResponseTemplate::new(200).set_body_json(body.clone())) + .expect(1) + .mount(&server) + .await; + let telemetry = telemetry(); + let target = upstream(server.uri(), telemetry.clone()); + let (http, _, bytes) = proxy::forward( + &WorkloadIdentityAuth::new(), + None, + &reqwest::Client::builder().no_proxy().build().unwrap(), + &target, + Method::POST, + "responses", + &HeaderMap::new(), + Bytes::from("{}"), + ) + .await + .unwrap(); + assert_eq!(http, StatusCode::OK); + assert_eq!(serde_json::from_slice::(&bytes).unwrap(), body); + let trace = events(&telemetry); + assert_eq!(trace.len(), 1); + assert_eq!(trace[0]["accepted"], true); + assert_eq!(trace[0]["outcome"], outcome); + assert_eq!(trace[0]["finish_reason"], status); + assert_eq!( + trace[0]["usage_state"], + if status == "failed" { + "partial" + } else { + "missing" + } + ); + assert!(trace[0]["usage"]["completion_tokens"].is_null()); + assert!(trace[0]["usage"]["total_tokens"].is_null()); + assert!(!trace[0].to_string().contains("secret-marker")); + assert!(!trace[0].to_string().contains("private-details-marker")); + } +} + +#[tokio::test] +async fn fragmented_openai_error_frames_cannot_be_overwritten_by_done() { + for error in [ + "data: {\"error\":{\"message\":\"credential-secret-marker🚫\"}}\n\n", + "event: error\ndata: {\"message\":\"credential-secret-marker🚫\"}\n\n", + ] { + let body = format!( + "data: {{\"choices\":[],\"usage\":{{\"prompt_tokens\":3}}}}\n\n{error}data: [DONE]\n\n" + ); + let telemetry = telemetry(); + let mut observation = telemetry + .begin("chat/completions", "provider", "model", b"{}") + .unwrap(); + observation.headers(200); + let chunks = body + .as_bytes() + .iter() + .map(|byte| Ok::<_, reqwest::Error>(Bytes::from(vec![*byte]))) + .collect::>(); + let returned = wrap_stream( + futures::stream::iter(chunks).boxed(), + Some(observation), + true, + ) + .try_collect::>() + .await + .unwrap(); + assert_eq!(returned.concat(), body.as_bytes()); + let trace = events(&telemetry); + assert_eq!(trace.len(), 1); + assert_eq!(trace[0]["accepted"], true); + assert_eq!(trace[0]["outcome"], "upstream_error"); + assert_eq!(trace[0]["usage"]["prompt_tokens"], 3); + assert!(trace[0]["usage"]["completion_tokens"].is_null()); + assert!(trace[0]["usage"]["total_tokens"].is_null()); + assert!(!trace[0].to_string().contains("credential-secret-marker")); + } +} + +#[tokio::test] +async fn responses_incomplete_stream_is_not_reported_as_a_complete_model_response() { + let telemetry = telemetry(); + let mut observation = telemetry + .begin("responses", "provider", "model", b"{}") + .unwrap(); + observation.headers(200); + let body = "data: {\"type\":\"response.incomplete\",\"response\":{\"status\":\"incomplete\",\"usage\":{\"input_tokens\":4,\"output_tokens\":2},\"output\":[]}}\n\n"; + let returned = wrap_stream( + futures::stream::once(async move { + Ok::<_, reqwest::Error>(Bytes::from_static(body.as_bytes())) + }) + .boxed(), + Some(observation), + true, + ) + .try_collect::>() + .await + .unwrap(); + assert_eq!(returned.concat(), body.as_bytes()); + let trace = events(&telemetry); + assert_eq!(trace[0]["accepted"], true); + assert_eq!(trace[0]["outcome"], "incomplete"); + assert_eq!(trace[0]["usage_state"], "present"); + assert!(trace[0]["usage"]["total_tokens"].is_null()); +} + #[tokio::test] async fn fragmented_streams_preserve_bytes_and_report_each_supported_usage_shape_once() { let cases = [ diff --git a/tests/e2e/run.sh b/tests/e2e/run.sh index e1b26dee7..d57b47b47 100755 --- a/tests/e2e/run.sh +++ b/tests/e2e/run.sh @@ -3048,7 +3048,7 @@ main() { test_sandbox_deployment_exists || true test_sandbox_pod_starts || true if test_governed_services; then - pass "Governed services isolate operator credentials, bind live identity, and reject stale scopes" + pass "Service API smoke: router-only token mount, credential checks, and scope reset" else fail "Governed service authentication and scope lifecycle gate failed" fi From fb8035fe956de5d771c914e3f638260dfa9654b2 Mon Sep 17 00:00:00 2001 From: pallakatos Date: Tue, 8 Sep 2026 20:05:49 +0200 Subject: [PATCH 3/4] fix(controller): satisfy strict governed privacy integration lint Paired offline/default-feature controller/router qualification passed: 43 controller tests, 11 router unit tests, and 29 governed HTTP integration tests. Strict all-target Clippy passed after these mechanical fixes. Hosted SRE Kind and independent review remain pending; no deployment or public publication approval is asserted. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- controller/src/reconciler/governed_services/credentials.rs | 2 +- controller/src/reconciler/mod.rs | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/controller/src/reconciler/governed_services/credentials.rs b/controller/src/reconciler/governed_services/credentials.rs index b9df854ab..5f09fb384 100644 --- a/controller/src/reconciler/governed_services/credentials.rs +++ b/controller/src/reconciler/governed_services/credentials.rs @@ -144,7 +144,7 @@ async fn review_consumer( .as_deref() .is_none_or(str::is_empty) || deployment.metadata.deletion_timestamp.is_some() - || !crate::sre_authority::controller_managed(&deployment, name) + || !crate::sre_authority::controller_managed(deployment, name) { return Err( "Governed service credential consumer is not a live controller-owned Deployment" diff --git a/controller/src/reconciler/mod.rs b/controller/src/reconciler/mod.rs index 57e1418c3..9aeba3e9c 100644 --- a/controller/src/reconciler/mod.rs +++ b/controller/src/reconciler/mod.rs @@ -3013,7 +3013,6 @@ async fn reconcile(sandbox: Arc, ctx: Arc) -> Result Date: Tue, 8 Sep 2026 20:50:10 +0200 Subject: [PATCH 4/4] fix(controller): complete privacy qualification before SRE readiness Remove workload availability from the credential privacy barrier. Fresh denial, v2/UID/epoch/template verification and termination of old cached credential Pods remain required. The continuity regression retains zero SRE/normal available replicas, removes the actual old Pod instead of mutating its version, and proves Ready unlocks SRE authorization. 39 targeted tests and strict paired all-target Clippy pass; full hosted qualification and review remain pending. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: eb3654cd-f1e0-445a-8734-430800af1903 --- .../governed_services/continuity_tests.rs | 58 +++++++++++++++---- controller/src/sre_authority/migration.rs | 18 ++---- docs/governed-services.md | 8 ++- 3 files changed, 57 insertions(+), 27 deletions(-) diff --git a/controller/src/reconciler/governed_services/continuity_tests.rs b/controller/src/reconciler/governed_services/continuity_tests.rs index 983c62786..8cb580eed 100644 --- a/controller/src/reconciler/governed_services/continuity_tests.rs +++ b/controller/src/reconciler/governed_services/continuity_tests.rs @@ -398,14 +398,22 @@ async fn interleaved_full_sre_and_sandbox_credential_reconcile_keep_qualified_ep } #[tokio::test] -async fn unfinished_old_token_retirement_blocks_issuance_without_stopping_current_rollout() { +async fn credential_transition_qualifies_before_authority_dependent_sre_availability() { let (_server, client, state) = fixture().await; { let mut state = state.lock().unwrap(); + let mut old = state.objects["/api/v1/namespaces/kars-normal/pods"]["items"][0].clone(); + old["metadata"]["name"] = "old-cached-router".into(); + old["metadata"]["uid"] = "old-cached-router".into(); + old["metadata"]["deletionTimestamp"] = "2026-09-08T00:00:00Z".into(); + old["metadata"]["annotations"][CONTROL_VERSION] = "old-token".into(); state .objects .get_mut("/api/v1/namespaces/kars-normal/pods") - .unwrap()["items"][0]["metadata"]["annotations"][CONTROL_VERSION] = "old-token".into(); + .unwrap()["items"] + .as_array_mut() + .unwrap() + .push(old); } let reg: KarsSRERegistration = object(&state, REG); assert!( @@ -415,7 +423,14 @@ async fn unfinished_old_token_retirement_blocks_issuance_without_stopping_curren ); let sandbox: KarsSandbox = object(&state, NORMAL); let namespace: Namespace = object(&state, "/api/v1/namespaces/kars-normal"); + let sre: KarsSandbox = object(&state, SRE); + let sre_namespace: Namespace = object(&state, "/api/v1/namespaces/kars-sre"); assert!(super::ensure(&client, &sandbox, &namespace).await.is_err()); + assert!( + crate::sre_authority::pod::authorize(&client, &sre, &sre_namespace) + .await + .is_err() + ); { let mut state = state.lock().unwrap(); assert_eq!(state.objects[REG]["status"]["phase"], "Migrating"); @@ -432,15 +447,16 @@ async fn unfinished_old_token_retirement_blocks_issuance_without_stopping_curren state .objects .get_mut("/api/v1/namespaces/kars-normal/pods") - .unwrap()["items"][0]["metadata"]["annotations"][CONTROL_VERSION] = - "control-normal:1".into(); + .unwrap()["items"] + .as_array_mut() + .unwrap() + .retain(|pod| pod["metadata"]["uid"] != "old-cached-router"); for name in ["sre", "normal"] { - state - .objects - .get_mut(&format!( - "/apis/apps/v1/namespaces/kars-{name}/deployments/{name}" - )) - .unwrap()["status"]["availableReplicas"] = 1.into(); + assert_eq!( + state.objects[&format!("/apis/apps/v1/namespaces/kars-{name}/deployments/{name}")] + ["status"]["availableReplicas"], + 0 + ); } } let reg: KarsSRERegistration = object(&state, REG); @@ -448,4 +464,26 @@ async fn unfinished_old_token_retirement_blocks_issuance_without_stopping_curren .await .unwrap(); super::ensure(&client, &sandbox, &namespace).await.unwrap(); + assert!( + crate::sre_authority::pod::authorize(&client, &sre, &sre_namespace) + .await + .unwrap() + .is_some() + ); + let state = state.lock().unwrap(); + assert_eq!(state.objects[REG]["status"]["phase"], "Ready"); + assert_eq!( + state.objects["/apis/apps/v1/namespaces/kars-sre/deployments/sre"]["status"]["availableReplicas"], + 0 + ); + assert_eq!( + state.objects["/apis/apps/v1/namespaces/kars-normal/deployments/normal"]["status"]["availableReplicas"], + 0 + ); + assert!( + state + .calls + .iter() + .all(|(_, _, body)| body.get("stringData").is_none()) + ); } diff --git a/controller/src/sre_authority/migration.rs b/controller/src/sre_authority/migration.rs index 5c4e4b5a0..112c48d12 100644 --- a/controller/src/sre_authority/migration.rs +++ b/controller/src/sre_authority/migration.rs @@ -22,7 +22,7 @@ pub(super) const WAITING_FOR_CONSUMERS: &str = const WAITING_FOR_ROTATION: &str = "Waiting for owned control credential consumers to restart on the new privacy epoch"; const WAITING_FOR_ROLLOUT: &str = - "Owned control credential consumer has not completed its privacy-epoch rollout"; + "Waiting for prior-epoch or prior-control-version consumers to terminate"; pub(super) fn is_waiting(detail: &str) -> bool { matches!( @@ -331,16 +331,6 @@ pub(super) async fn rotate_owned_control_credentials( }))).await.map_err(|e|api_error("Restart owned control credential consumer",e))?; return Err(WAITING_FOR_ROTATION.into()); } - let desired = deployment - .spec - .as_ref() - .and_then(|spec| spec.replicas) - .unwrap_or(1); - let ready = deployment.status.as_ref().is_some_and(|status| { - status.observed_generation == deployment.metadata.generation - && status.updated_replicas.unwrap_or(0) == desired - && status.available_replicas.unwrap_or(0) == desired - }); let selector = deployment .spec .as_ref() @@ -366,9 +356,9 @@ pub(super) async fn rotate_owned_control_credentials( // router which continues accepting its startup-cached control token. return Err(WAITING_FOR_ROLLOUT.into()); } - if !ready && !super::live::currently_qualified(reg) { - return Err(WAITING_FOR_ROLLOUT.into()); - } + // Privacy completion is credential/template identity plus termination + // of old caches, not workload availability. In particular, the SRE + // readiness endpoint itself requires this authority to become Ready. } Ok(()) } diff --git a/docs/governed-services.md b/docs/governed-services.md index e71727559..81808dddc 100644 --- a/docs/governed-services.md +++ b/docs/governed-services.md @@ -44,9 +44,11 @@ to zero. Restoring authority requires a new token, not reuse of the potentially exposed cache. Foreign consumers remain untouched and block automatic recovery. Initial qualification still pending is not permission to issue or reuse a credential, but it is not itself evidence of privacy loss and does not stop the -rollout being qualified. An already-qualified current v2 epoch survives ordinary -availability failures only while fresh privacy checks pass and no old-epoch or -old-control-version Pods remain. Canonical SRE's early authorization failure +rollout being qualified. Privacy qualification completes only after fresh +denial/v2/UID/epoch/template checks and termination of old-epoch or +old-control-version Pods. Workload availability is separate: the SRE readiness +endpoint itself depends on qualified authority, so it cannot be an input to +that qualification. Canonical SRE's early authorization failure also performs the ownership-fenced, no-issuance privacy-loss quarantine. These checks occur during reconciliation; they are not a claim of instantaneous cluster-wide revocation or cancellation of already accepted upstream work.