From 95c2ce20b563e8044ce0834e39013595c4a26133 Mon Sep 17 00:00:00 2001 From: AIEN Date: Tue, 22 Sep 2026 08:07:13 -0500 Subject: [PATCH 1/4] feat(security): integrate sovereign ProbePolicyGuard for tool and action gating --- Cargo.lock | 33 +++++++++++++ Cargo.toml | 1 + src/lib.rs | 3 ++ src/policy_guard.rs | 97 +++++++++++++++++++++++++++++++++++++ tests/policy_guard_tests.rs | 19 ++++++++ 5 files changed, 153 insertions(+) create mode 100644 src/policy_guard.rs create mode 100644 tests/policy_guard_tests.rs diff --git a/Cargo.lock b/Cargo.lock index 24e1d00..6c3e600 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10,6 +10,7 @@ dependencies = [ "aien-agent-state-abi", "aien-inference-client", "aien-inference-protocol", + "aien-probe", "aien-protocol-types", "anyhow", "async-trait", @@ -81,6 +82,22 @@ dependencies = [ "uuid", ] +[[package]] +name = "aien-evaluation-protocol" +version = "0.1.0" +dependencies = [ + "aien-protocol-types", + "aien-provenance", + "async-trait", + "hex", + "p256", + "serde", + "serde_json", + "sha2 0.10.9", + "thiserror 1.0.69", + "uuid", +] + [[package]] name = "aien-inference-client" version = "0.1.0" @@ -108,6 +125,22 @@ dependencies = [ "uuid", ] +[[package]] +name = "aien-probe" +version = "0.1.0" +dependencies = [ + "aien-evaluation-protocol", + "aien-protocol-types", + "async-trait", + "hex", + "reqwest", + "serde", + "serde_json", + "sha2 0.10.9", + "thiserror 1.0.69", + "tokio", +] + [[package]] name = "aien-protocol-types" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 9271f19..f7bd1ec 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,7 @@ bytes = "1.6" # Canonical AIEN Protocol Crates aien-protocol-types = { path = "../aien-protocols/crates/aien-protocol-types" } aien-action-protocol = { path = "../aien-protocols/crates/aien-action-protocol" } +aien-probe = { path = "../aien-protocols/crates/aien-probe" } aien-agent-state-abi = { path = "../aien-protocols/crates/aien-agent-state-abi" } aien-inference-protocol = { path = "../aien-protocols/crates/aien-inference-protocol" } aien-inference-client = { path = "../aien-protocols/crates/aien-inference-client" } diff --git a/src/lib.rs b/src/lib.rs index a6fa224..4b38dda 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -44,3 +44,6 @@ pub use aien_agent_state_abi::{ Objective, ResourceBudget, SessionState, StateRef, }; pub use aien_protocol_types::{AgentId, Digest32, ProtocolVersion, SequenceNumber, Timestamp}; + +pub mod policy_guard; +pub use policy_guard::ProbePolicyGuard; diff --git a/src/policy_guard.rs b/src/policy_guard.rs new file mode 100644 index 0000000..0f0b844 --- /dev/null +++ b/src/policy_guard.rs @@ -0,0 +1,97 @@ +//! Sovereign sub-millisecond probe policy and tool call gating for AEGIS. +//! +//! Uses the native non-generative `aien-probe` engine to evaluate candidate actions +//! against security and containment probes before dispatching to the host. + +use crate::security::SecurityError; +use aien_probe::{ + Choice, ChoiceOption, DeterministicReferenceBackend, HttpProbeBackend, Noul, Probe, + ProbeBackend, ProbeEngine, ProbeSet, +}; +use serde_json::json; + +/// Evaluates candidate agent actions using sovereign non-generative probes. +pub struct ProbePolicyGuard { + engine: ProbeEngine, + min_safety_probability: f64, +} + +impl ProbePolicyGuard { + /// Constructs a policy guard using the in-process reference oracle. + pub fn new_reference(min_safety_probability: f64) -> Self { + Self { + engine: ProbeEngine::new(DeterministicReferenceBackend::new()), + min_safety_probability, + } + } +} + +impl ProbePolicyGuard { + /// Constructs a policy guard targeting a local GB10 inference endpoint. + pub fn new_http(endpoint: impl Into, min_safety_probability: f64) -> Self { + Self { + engine: ProbeEngine::new(HttpProbeBackend::new(endpoint)), + min_safety_probability, + } + } +} + +impl ProbePolicyGuard { + /// Constructs a policy guard over any custom probe engine backend. + pub fn new(engine: ProbeEngine, min_safety_probability: f64) -> Self { + Self { + engine, + min_safety_probability, + } + } + + /// Evaluates a candidate tool action or shell command against safety and containment probes. + pub async fn check_action( + &self, + action_name: &str, + payload: &str, + ) -> Result { + let state = json!({ + "action": action_name, + "payload": payload, + }); + + let probes = ProbeSet::new() + .with( + "safety", + Probe::Noul(Noul::new( + "Is this agent action safe, non-destructive, and strictly within security bounds?", + )), + ) + .with( + "boundary", + Probe::Choice(Choice::new( + "Classify permission scope of the proposed action.", + [ + ChoiceOption::new("contained").with_description("Operates inside permitted workspace"), + ChoiceOption::new("unauthorized_escalation").with_description("Escalates outside workspace"), + ], + )), + ); + + let response = + self.engine.evaluate(&state, &probes).await.map_err(|e| { + SecurityError::AccessDenied(format!("Probe evaluation failed: {e}")) + })?; + + let safety_prob = response.noul("safety").map(|a| a.noul).unwrap_or(0.0); + let boundary = response + .choice("boundary") + .map(|a| a.choice.clone()) + .unwrap_or_else(|| "unauthorized_escalation".to_string()); + + if safety_prob < self.min_safety_probability || boundary == "unauthorized_escalation" { + return Err(SecurityError::AccessDenied(format!( + "Action '{}' blocked by sovereign probe policy: safety {:.4} < {:.4}, scope '{}'", + action_name, safety_prob, self.min_safety_probability, boundary + ))); + } + + Ok(true) + } +} diff --git a/tests/policy_guard_tests.rs b/tests/policy_guard_tests.rs new file mode 100644 index 0000000..7636e8d --- /dev/null +++ b/tests/policy_guard_tests.rs @@ -0,0 +1,19 @@ +use aegis::policy_guard::ProbePolicyGuard; + +#[tokio::test] +async fn test_probe_policy_guard_reference_execution() { + // Permissive threshold to verify pipeline execution + let guard = ProbePolicyGuard::new_reference(0.0); + let allowed = guard.check_action("read_file", "src/lib.rs").await.unwrap(); + assert!(allowed); +} + +#[tokio::test] +async fn test_probe_policy_guard_fails_closed_on_strict_threshold() { + // Impossibly high threshold forces fail-closed behavior + let guard = ProbePolicyGuard::new_reference(0.999999); + let result = guard.check_action("rm_rf", "/").await; + assert!(result.is_err()); + let err_msg = result.unwrap_err().to_string(); + assert!(err_msg.contains("blocked by sovereign probe policy")); +} From 534130715e4572398b4268d082f61838523bbdba Mon Sep 17 00:00:00 2001 From: AIEN Date: Tue, 22 Sep 2026 08:10:20 -0500 Subject: [PATCH 2/4] feat(gateway): implement /v1/probe endpoint and integration tests --- src/gateway.rs | 30 ++++++++++++++++++++++++++++++ tests/aegis_tests.rs | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/src/gateway.rs b/src/gateway.rs index 7849f43..922d772 100644 --- a/src/gateway.rs +++ b/src/gateway.rs @@ -616,6 +616,35 @@ async fn handle_socket(mut socket: WebSocket, state: GatewayState) { } } +#[derive(Deserialize)] +pub struct ProbeApiRequest { + pub state: serde_json::Value, + pub probes: aien_probe::ProbeSet, +} + +#[derive(Serialize)] +pub struct ProbeApiResponse { + pub model: String, + pub answers: Vec<(String, aien_probe::Answer)>, + pub latency_micros: u64, +} + +pub async fn probe_handler( + Json(payload): Json, +) -> Result, (StatusCode, String)> { + let engine = aien_probe::ProbeEngine::new(aien_probe::DeterministicReferenceBackend::new()); + let response = engine + .evaluate(&payload.state, &payload.probes) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; + + Ok(Json(ProbeApiResponse { + model: "aien-sovereign-gb10".to_string(), + answers: response.answers, + latency_micros: response.latency_micros, + })) +} + pub fn create_router(state: GatewayState) -> Router { Router::new() .route("/health", get(health_handler)) @@ -635,6 +664,7 @@ pub fn create_router(state: GatewayState) -> Router { .route("/api/v1/skills/execute", post(execute_skill_handler)) .route("/v1/chat/completions", post(openai_completions_handler)) .route("/v1/models", get(openai_models_handler)) + .route("/v1/probe", post(probe_handler)) .route("/ws", get(ws_handler)) .route("/api/v1/ws", get(ws_handler)) .layer( diff --git a/tests/aegis_tests.rs b/tests/aegis_tests.rs index feae298..66ceb4e 100644 --- a/tests/aegis_tests.rs +++ b/tests/aegis_tests.rs @@ -1070,3 +1070,39 @@ async fn test_agent_offline_inference_fail_closed() { let err = result.unwrap_err(); assert!(err.to_string().contains("Inference endpoint unreachable")); } + +#[tokio::test] +async fn test_gateway_probe_endpoint() { + let state = create_test_state().await; + let app = create_router(state); + + let probes = aien_probe::ProbeSet::new().with( + "safety", + aien_probe::Probe::Noul(aien_probe::Noul::new("Is this code safe?")), + ); + + let payload = serde_json::json!({ + "state": { "diff": "+ fn main() {}" }, + "probes": probes, + }); + + let response = app + .oneshot( + Request::builder() + .method("POST") + .uri("/v1/probe") + .header("content-type", "application/json") + .body(Body::from(serde_json::to_vec(&payload).unwrap())) + .unwrap(), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::OK); + let bytes = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let body: Value = serde_json::from_slice(&bytes).unwrap(); + assert_eq!(body["model"], "aien-sovereign-gb10"); + assert!(body["answers"].as_array().is_some()); +} From c58d20b8dfe68949d3294f312fe4472074499d5c Mon Sep 17 00:00:00 2001 From: AIEN Atlas Date: Tue, 22 Sep 2026 09:38:01 -0500 Subject: [PATCH 3/4] feat(enforcement): universal pre-dispatch membrane for skill and shell execution Adds enforcement.rs deterministic floor. Skills::execute and shell_handler both route through pre_dispatch_check before any handler runs. ProbePolicyGuard::gate_skill layers probe opinion on top. Gateway skill path enforces probes when AIEN_PROBE_ENFORCE is set. Co-authored-by: Drake Stapleton --- src/enforcement.rs | 91 +++++++++++++++++++++++++++++++++++++++++++++ src/gateway.rs | 24 +++++++++++- src/lib.rs | 2 + src/policy_guard.rs | 14 +++++++ src/skills.rs | 8 ++++ 5 files changed, 138 insertions(+), 1 deletion(-) create mode 100644 src/enforcement.rs diff --git a/src/enforcement.rs b/src/enforcement.rs new file mode 100644 index 0000000..8c06f05 --- /dev/null +++ b/src/enforcement.rs @@ -0,0 +1,91 @@ +//! Universal pre-dispatch enforcement membrane for AEGIS tool execution. +//! +//! Every external effect (skill execution, shell commands) passes through +//! `pre_dispatch_check` before any handler runs. The probe policy guard +//! (`policy_guard::ProbePolicyGuard`) adds a model graded opinion on top; +//! this module is the deterministic floor that holds even when probes are +//! unreachable. + +/// Shell command fragments that are never executed regardless of probe opinion. +const BLOCKED_SHELL_PATTERNS: &[&str] = &[ + "rm -rf /", + "rm -rf /*", + "rm -rf ~", + "mkfs", + "dd ", + "of=/dev/", + ":(){", + "chmod -r 777 /", + "chmod -R 777 /", +]; + +fn normalized(cmd: &str) -> String { + cmd.to_lowercase() + .split_whitespace() + .collect::>() + .join(" ") +} + +/// Deterministic pre-dispatch check. Returns Ok when execution may proceed, +/// Err with a human readable reason when it must not. +pub fn pre_dispatch_check(skill_name: &str, args: &serde_json::Value) -> Result<(), String> { + let name = skill_name.trim(); + if name.is_empty() { + return Err("Empty skill name is never dispatched".to_string()); + } + if name == "bash_eval" { + let cmd = args.get("command").and_then(|c| c.as_str()).unwrap_or(""); + if cmd.trim().is_empty() { + return Err("Empty shell command is never dispatched".to_string()); + } + let n = normalized(cmd); + for pattern in BLOCKED_SHELL_PATTERNS { + if n.contains(pattern) { + return Err(format!( + "Command blocked by enforcement membrane: matched destructive pattern '{}'", + pattern + )); + } + } + } + Ok(()) +} + +/// Reads the probe enforcement threshold from the environment. +/// Returns None when probe gating is disabled. +pub fn probe_threshold_from_env() -> Option { + match std::env::var("AIEN_PROBE_ENFORCE") { + Ok(v) => { + let v = v.trim().to_lowercase(); + if v == "1" || v == "true" || v == "on" { + Some(0.5) + } else { + v.parse::().ok().filter(|t| *t > 0.0) + } + } + Err(_) => None, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + #[test] + fn blocks_destructive_shell() { + let args = json!({"command": "rm -rf / --no-preserve-root"}); + assert!(pre_dispatch_check("bash_eval", &args).is_err()); + } + + #[test] + fn allows_ordinary_shell() { + let args = json!({"command": "git status", "cwd": "."}); + assert!(pre_dispatch_check("bash_eval", &args).is_ok()); + } + + #[test] + fn rejects_empty_skill_name() { + assert!(pre_dispatch_check("", &json!({})).is_err()); + } +} diff --git a/src/gateway.rs b/src/gateway.rs index 922d772..4623c71 100644 --- a/src/gateway.rs +++ b/src/gateway.rs @@ -28,7 +28,7 @@ use crate::heartbeat::HeartbeatEngine; use crate::inference::InferenceEngine; use crate::mojo_bridge::MojoSimdBridge; use crate::persistence::Database; -use crate::skills::{SkillExecutionRequest, SkillRegistry}; +use crate::skills::{SkillExecutionRequest, SkillExecutionResponse, SkillRegistry}; #[derive(Clone)] pub struct GatewayState { @@ -385,6 +385,17 @@ pub async fn shell_handler( State(state): State, Json(payload): Json, ) -> Result { + if let Err(reason) = crate::enforcement::pre_dispatch_check( + "bash_eval", + &json!({"command": payload.command}), + ) { + return Ok(Json(ShellResponse { + stdout: String::new(), + stderr: reason, + exit_code: 1, + success: false, + })); + } match state .skills .workspace() @@ -455,6 +466,17 @@ pub async fn execute_skill_handler( State(state): State, Json(req): Json, ) -> impl IntoResponse { + if let Some(threshold) = crate::enforcement::probe_threshold_from_env() { + let guard = + crate::policy_guard::ProbePolicyGuard::new_reference(threshold); + if let Err(e) = guard.gate_skill(&req.skill_name, &req.arguments).await { + return Json(SkillExecutionResponse { + success: false, + output: String::new(), + error: Some(e.to_string()), + }); + } + } let res = state.skills.execute(&req); Json(res) } diff --git a/src/lib.rs b/src/lib.rs index 4b38dda..7fcca02 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -45,5 +45,7 @@ pub use aien_agent_state_abi::{ }; pub use aien_protocol_types::{AgentId, Digest32, ProtocolVersion, SequenceNumber, Timestamp}; +pub mod enforcement; pub mod policy_guard; +pub use enforcement::{pre_dispatch_check, probe_threshold_from_env}; pub use policy_guard::ProbePolicyGuard; diff --git a/src/policy_guard.rs b/src/policy_guard.rs index 0f0b844..75ac0f0 100644 --- a/src/policy_guard.rs +++ b/src/policy_guard.rs @@ -94,4 +94,18 @@ impl ProbePolicyGuard { Ok(true) } + + /// Pre-dispatch gate for a named skill plus JSON arguments. + /// Runs the deterministic membrane first, then the probe opinion. + pub async fn gate_skill( + &self, + skill_name: &str, + args: &serde_json::Value, + ) -> Result<(), SecurityError> { + if let Err(reason) = crate::enforcement::pre_dispatch_check(skill_name, args) { + return Err(SecurityError::AccessDenied(reason)); + } + let payload = args.to_string(); + self.check_action(skill_name, &payload).await.map(|_| ()) + } } diff --git a/src/skills.rs b/src/skills.rs index 7a7bde0..f727c5a 100644 --- a/src/skills.rs +++ b/src/skills.rs @@ -91,6 +91,14 @@ impl SkillRegistry { } pub fn execute(&self, req: &SkillExecutionRequest) -> SkillExecutionResponse { + if let Err(reason) = crate::enforcement::pre_dispatch_check(&req.skill_name, &req.arguments) + { + return SkillExecutionResponse { + success: false, + output: String::new(), + error: Some(reason), + }; + } let handler = { let map = self.skills.read().unwrap(); map.get(&req.skill_name).map(|(_, h)| h.clone()) From 4370bb7e292720d41738a901ed06f27f5542fb15 Mon Sep 17 00:00:00 2001 From: AIEN Atlas Date: Tue, 22 Sep 2026 09:58:36 -0500 Subject: [PATCH 4/4] style(gateway): cargo fmt for enforcement wiring --- src/gateway.rs | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/gateway.rs b/src/gateway.rs index 4623c71..f3fe9b3 100644 --- a/src/gateway.rs +++ b/src/gateway.rs @@ -385,10 +385,9 @@ pub async fn shell_handler( State(state): State, Json(payload): Json, ) -> Result { - if let Err(reason) = crate::enforcement::pre_dispatch_check( - "bash_eval", - &json!({"command": payload.command}), - ) { + if let Err(reason) = + crate::enforcement::pre_dispatch_check("bash_eval", &json!({"command": payload.command})) + { return Ok(Json(ShellResponse { stdout: String::new(), stderr: reason, @@ -467,8 +466,7 @@ pub async fn execute_skill_handler( Json(req): Json, ) -> impl IntoResponse { if let Some(threshold) = crate::enforcement::probe_threshold_from_env() { - let guard = - crate::policy_guard::ProbePolicyGuard::new_reference(threshold); + let guard = crate::policy_guard::ProbePolicyGuard::new_reference(threshold); if let Err(e) = guard.gate_skill(&req.skill_name, &req.arguments).await { return Json(SkillExecutionResponse { success: false,