diff --git a/crates/core/src/inspection.rs b/crates/core/src/inspection.rs new file mode 100644 index 00000000..84ecdb2b --- /dev/null +++ b/crates/core/src/inspection.rs @@ -0,0 +1,136 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +//! Internal neutral inspection contract primitives and Relay adapters for the +//! cross-repo POC. + +#![cfg_attr(not(test), allow(dead_code))] + +use crate::api::llm::LlmRequest; +use crate::error::{FlowError, Result}; +use crate::json::Json; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub(crate) enum InspectionTarget { + LlmRequest { + provider: String, + request: Json, + }, + ToolRequest { + tool_name: String, + input: Json, + }, + HttpRequest { + method: String, + path: String, + headers: Vec<(String, String)>, + body: Vec, + }, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)] +pub(crate) struct InspectionContext { + pub sandbox_id: Option, + pub scope_id: Option, + pub provider: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub(crate) struct Finding { + pub code: String, + pub message: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub(crate) enum InspectionDecision { + Allow, + Deny { + reason: String, + findings: Vec, + }, + Mutate { + target: InspectionTarget, + findings: Vec, + }, +} + +pub(crate) trait Inspector: Send + Sync { + fn inspect( + &self, + target: InspectionTarget, + ctx: &InspectionContext, + ) -> Result; +} + +pub(crate) struct RelayInspectionAdapter { + inspector: I, +} + +impl RelayInspectionAdapter +where + I: Inspector, +{ + pub(crate) fn new(inspector: I) -> Self { + Self { inspector } + } + + pub(crate) fn inspect_llm_request( + &self, + provider: &str, + request: LlmRequest, + ctx: &InspectionContext, + ) -> Result { + let request_value = serde_json::to_value(&request) + .map_err(|error| FlowError::InvalidArgument(error.to_string()))?; + let decision = self.inspector.inspect( + InspectionTarget::LlmRequest { + provider: provider.to_string(), + request: request_value, + }, + ctx, + )?; + + match decision { + InspectionDecision::Allow => Ok(request), + InspectionDecision::Deny { reason, .. } => Err(FlowError::GuardrailRejected(reason)), + InspectionDecision::Mutate { target, .. } => match target { + InspectionTarget::LlmRequest { request, .. } => serde_json::from_value(request) + .map_err(|error| FlowError::InvalidArgument(error.to_string())), + other => Err(FlowError::InvalidArgument(format!( + "expected mutated LlmRequest target, got {other:?}" + ))), + }, + } + } + + pub(crate) fn inspect_tool_request( + &self, + tool_name: &str, + input: Json, + ctx: &InspectionContext, + ) -> Result { + let decision = self.inspector.inspect( + InspectionTarget::ToolRequest { + tool_name: tool_name.to_string(), + input: input.clone(), + }, + ctx, + )?; + + match decision { + InspectionDecision::Allow => Ok(input), + InspectionDecision::Deny { reason, .. } => Err(FlowError::GuardrailRejected(reason)), + InspectionDecision::Mutate { target, .. } => match target { + InspectionTarget::ToolRequest { input, .. } => Ok(input), + other => Err(FlowError::InvalidArgument(format!( + "expected mutated ToolRequest target, got {other:?}" + ))), + }, + } + } +} + +#[cfg(test)] +#[path = "../tests/unit/inspection_tests.rs"] +mod tests; diff --git a/crates/core/src/lib.rs b/crates/core/src/lib.rs index ffd03a4d..9039beb0 100644 --- a/crates/core/src/lib.rs +++ b/crates/core/src/lib.rs @@ -58,6 +58,7 @@ pub mod codec; pub mod config_editor; mod context; pub mod error; +mod inspection; pub mod json; pub mod observability; pub mod plugin; diff --git a/crates/core/tests/unit/inspection_tests.rs b/crates/core/tests/unit/inspection_tests.rs new file mode 100644 index 00000000..18743e7c --- /dev/null +++ b/crates/core/tests/unit/inspection_tests.rs @@ -0,0 +1,157 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +use crate::api::llm::LlmRequest; +use crate::error::FlowError; +use crate::inspection::{ + Finding, InspectionContext, InspectionDecision, InspectionTarget, Inspector, + RelayInspectionAdapter, +}; +use crate::json::Json; +use serde_json::{Map, json}; + +struct TestInspector; + +impl Inspector for TestInspector { + fn inspect( + &self, + target: InspectionTarget, + ctx: &InspectionContext, + ) -> crate::error::Result { + match target { + InspectionTarget::LlmRequest { + provider, + mut request, + } => { + if ctx.provider.as_deref() == Some("deny-llm") { + return Ok(InspectionDecision::Deny { + reason: format!("blocked provider {provider}"), + findings: vec![Finding { + code: "llm_denied".to_string(), + message: "provider rejected".to_string(), + }], + }); + } + + request["headers"]["x-inspected"] = json!("true"); + Ok(InspectionDecision::Mutate { + target: InspectionTarget::LlmRequest { provider, request }, + findings: vec![Finding { + code: "llm_mutated".to_string(), + message: "request annotated".to_string(), + }], + }) + } + InspectionTarget::ToolRequest { + tool_name, + mut input, + } => { + if ctx.scope_id.as_deref() == Some("deny-tool") { + return Ok(InspectionDecision::Deny { + reason: format!("blocked tool {tool_name}"), + findings: vec![Finding { + code: "tool_denied".to_string(), + message: "tool rejected".to_string(), + }], + }); + } + + input["relay_inspected"] = json!(true); + Ok(InspectionDecision::Mutate { + target: InspectionTarget::ToolRequest { tool_name, input }, + findings: vec![Finding { + code: "tool_mutated".to_string(), + message: "tool args annotated".to_string(), + }], + }) + } + InspectionTarget::HttpRequest { .. } => Ok(InspectionDecision::Allow), + } + } +} + +fn make_request() -> LlmRequest { + LlmRequest { + headers: Map::new(), + content: json!({ + "messages": [{"role": "user", "content": "hello"}] + }), + } +} + +#[test] +fn relay_inspection_adapter_mutates_llm_requests() { + let adapter = RelayInspectionAdapter::new(TestInspector); + let request = make_request(); + let ctx = InspectionContext { + provider: Some("allow-llm".to_string()), + ..InspectionContext::default() + }; + + let mutated = adapter + .inspect_llm_request("openai", request, &ctx) + .expect("inspection should succeed"); + + assert_eq!( + mutated.headers.get("x-inspected"), + Some(&Json::String("true".to_string())) + ); +} + +#[test] +fn relay_inspection_adapter_denies_llm_requests() { + let adapter = RelayInspectionAdapter::new(TestInspector); + let error = adapter + .inspect_llm_request( + "openai", + make_request(), + &InspectionContext { + provider: Some("deny-llm".to_string()), + ..InspectionContext::default() + }, + ) + .expect_err("inspection should deny"); + + match error { + FlowError::GuardrailRejected(reason) => { + assert!(reason.contains("blocked provider openai")); + } + other => panic!("expected guardrail rejection, got {other}"), + } +} + +#[test] +fn relay_inspection_adapter_mutates_tool_requests() { + let adapter = RelayInspectionAdapter::new(TestInspector); + let mutated = adapter + .inspect_tool_request( + "search", + json!({"query": "books"}), + &InspectionContext::default(), + ) + .expect("inspection should succeed"); + + assert_eq!(mutated["relay_inspected"], json!(true)); +} + +#[test] +fn relay_inspection_adapter_denies_tool_requests() { + let adapter = RelayInspectionAdapter::new(TestInspector); + let error = adapter + .inspect_tool_request( + "search", + json!({"query": "books"}), + &InspectionContext { + scope_id: Some("deny-tool".to_string()), + ..InspectionContext::default() + }, + ) + .expect_err("inspection should deny"); + + match error { + FlowError::GuardrailRejected(reason) => { + assert!(reason.contains("blocked tool search")); + } + other => panic!("expected guardrail rejection, got {other}"), + } +}