Skip to main content

aimds_core/
gate.rs

1//! `SafetyGate` — composing trait for the AIMDS 3-gate pipeline.
2//!
3//! This trait formalizes the canonical 3-gate flow that AIMDS
4//! consumers expect:
5//!
6//!   1. Pre-storage PII detection (`aimds-detection`)
7//!   2. Sanitization for cookies / tokens / high-entropy blobs
8//!      (`aimds-detection` + `aimds-analysis`)
9//!   3. Prompt-injection / role-hijack / jailbreak check
10//!      (`aimds-detection` + `aimds-response`)
11//!
12//! Downstream consumers (notably `ruvnet/ruflo`'s
13//! `ruflo-federation-peer` — see ADR-120 Step 3) embed this trait
14//! to run the 3-gate inspection in-process on every federation
15//! message hop. Composing the three crates' surfaces through one
16//! trait lets the embedder swap or stub the gate (e.g. for tests)
17//! without binding to a concrete pipeline type.
18//!
19//! `aimds-detection`, `-analysis`, and `-response` (v0.2.0+) provide
20//! concrete implementations of this trait via a `ComposedGate` type;
21//! this crate exports the trait shape only so `aimds-core` stays
22//! dep-light.
23
24use crate::types::{PromptInput, SanitizedOutput};
25use crate::AimdsError;
26use async_trait::async_trait;
27
28/// The three gates surfaced as a single composing operation.
29///
30/// Implementors decide whether to short-circuit (the canonical
31/// `aimds-response::ComposedGate` does — Block in any gate stops the
32/// pipeline) or run all three in parallel for telemetry. Either is
33/// valid; the trait contract only specifies the input/output shapes.
34#[async_trait]
35pub trait SafetyGate: Send + Sync {
36    /// Run the 3-gate inspection on a prompt-shaped input. Returns a
37    /// [`SafetyVerdict`] describing whether to forward, block, or
38    /// forward a redacted variant.
39    ///
40    /// All three gates inspect the same input. PII detection (gate 1)
41    /// and prompt-injection detection (gate 3) are pure functions of
42    /// the content; sanitization (gate 2) can mutate the input into
43    /// the `Redact` payload.
44    async fn inspect(&self, input: &PromptInput) -> Result<SafetyVerdict, AimdsError>;
45}
46
47/// Result of running the 3-gate pipeline on one input.
48#[derive(Debug, Clone)]
49pub enum SafetyVerdict {
50    /// All three gates passed. Forward the input as-is.
51    Pass,
52
53    /// At least one gate flagged the input as unsafe. The carried
54    /// string is a human-readable reason describing which gate fired
55    /// and (when known) what pattern triggered.
56    Block(String),
57
58    /// PII detection or sanitization rewrote the input. The carried
59    /// `SanitizedOutput` holds the cleaned content + a list of
60    /// redactions that were applied; forward the sanitized variant.
61    Redact(SanitizedOutput),
62}
63
64impl SafetyVerdict {
65    /// True when the verdict allows the message to be forwarded
66    /// (either as-is or after redaction).
67    pub fn is_forwardable(&self) -> bool {
68        matches!(self, SafetyVerdict::Pass | SafetyVerdict::Redact(_))
69    }
70
71    /// True when the verdict requires the message to be quarantined.
72    pub fn is_blocked(&self) -> bool {
73        matches!(self, SafetyVerdict::Block(_))
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80    use crate::types::SanitizedOutput;
81    use chrono::Utc;
82    use uuid::Uuid;
83
84    fn fake_sanitized() -> SanitizedOutput {
85        SanitizedOutput {
86            original_id: Uuid::nil(),
87            timestamp: Utc::now(),
88            sanitized_content: String::new(),
89            modifications: vec![],
90            is_safe: true,
91        }
92    }
93
94    #[test]
95    fn pass_is_forwardable_and_not_blocked() {
96        let v = SafetyVerdict::Pass;
97        assert!(v.is_forwardable());
98        assert!(!v.is_blocked());
99    }
100
101    #[test]
102    fn redact_is_forwardable_and_not_blocked() {
103        let v = SafetyVerdict::Redact(fake_sanitized());
104        assert!(v.is_forwardable());
105        assert!(!v.is_blocked());
106    }
107
108    #[test]
109    fn block_is_blocked_and_not_forwardable() {
110        let v = SafetyVerdict::Block("test rule".into());
111        assert!(!v.is_forwardable());
112        assert!(v.is_blocked());
113    }
114
115    /// Confirms a downstream `T: SafetyGate` bound accepts a
116    /// trivially-implemented gate. Doctest-equivalent.
117    #[test]
118    fn safety_gate_is_object_safe() {
119        fn requires_gate<T: SafetyGate>() {}
120        struct AlwaysPass;
121        #[async_trait]
122        impl SafetyGate for AlwaysPass {
123            async fn inspect(
124                &self,
125                _input: &PromptInput,
126            ) -> Result<SafetyVerdict, AimdsError> {
127                Ok(SafetyVerdict::Pass)
128            }
129        }
130        requires_gate::<AlwaysPass>();
131    }
132}