use std::sync::Arc;
use std::time::Duration;
use crate::{
decision::{BlockCode, Decision},
error::GuardrailError,
request::GuardrailRequest,
};
pub const STAGE_ERROR_CODE: &str = "stage_error";
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum FailMode {
#[default]
Open,
Closed,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum StageOutcome {
Completed,
ErrorAllowed,
ErrorBlocked,
}
#[derive(Debug, Clone, Copy)]
pub struct StageEvent {
pub stage: &'static str,
pub elapsed: Duration,
pub outcome: StageOutcome,
}
#[allow(clippy::double_must_use)]
#[async_trait::async_trait]
pub trait Stage: Send + Sync + 'static {
fn name(&self) -> &'static str;
async fn evaluate(&self, req: &GuardrailRequest) -> Result<Decision, GuardrailError>;
}
pub struct Pipeline {
stages: Vec<(Arc<dyn Stage>, FailMode)>,
}
impl Pipeline {
pub fn new(stages: Vec<Arc<dyn Stage>>) -> Self {
Self {
stages: stages.into_iter().map(|s| (s, FailMode::Open)).collect(),
}
}
pub fn builder() -> PipelineBuilder {
PipelineBuilder::default()
}
pub fn len(&self) -> usize {
self.stages.len()
}
pub fn is_empty(&self) -> bool {
self.stages.is_empty()
}
pub async fn run(
&self,
req: GuardrailRequest,
) -> Result<(Decision, GuardrailRequest), GuardrailError> {
self.run_with_observer(req, |_stage, _elapsed| {}).await
}
pub async fn run_with_observer<F>(
&self,
req: GuardrailRequest,
mut observer: F,
) -> Result<(Decision, GuardrailRequest), GuardrailError>
where
F: FnMut(&str, std::time::Duration),
{
self.run_with_events(req, |event| observer(event.stage, event.elapsed))
.await
}
pub async fn run_with_events<F>(
&self,
mut req: GuardrailRequest,
mut observer: F,
) -> Result<(Decision, GuardrailRequest), GuardrailError>
where
F: FnMut(&StageEvent),
{
let mut redaction_reasons: Vec<String> = Vec::new();
let mut redaction_entities: Vec<String> = Vec::new();
for (stage, fail_mode) in &self.stages {
let stage_name = stage.name();
tracing::debug!(stage = stage_name, "evaluating stage");
let stage_start = std::time::Instant::now();
let result = stage.evaluate(&req).await;
let elapsed = stage_start.elapsed();
let decision = match result {
Ok(d) => {
observer(&StageEvent {
stage: stage_name,
elapsed,
outcome: StageOutcome::Completed,
});
d
}
Err(e) => match fail_mode {
FailMode::Open => {
tracing::warn!(
stage = stage_name,
error = %e,
"stage failed; failing open (allow)"
);
observer(&StageEvent {
stage: stage_name,
elapsed,
outcome: StageOutcome::ErrorAllowed,
});
Decision::Allow
}
FailMode::Closed => {
tracing::error!(
stage = stage_name,
error = %e,
"stage failed; failing closed (block)"
);
observer(&StageEvent {
stage: stage_name,
elapsed,
outcome: StageOutcome::ErrorBlocked,
});
Decision::Block {
reason: format!(
"stage '{stage_name}' failed and is configured to fail closed"
),
code: BlockCode::Custom(STAGE_ERROR_CODE.into()),
}
}
},
};
tracing::debug!(
stage = stage_name,
decision = decision.name(),
"stage decision"
);
match decision {
Decision::Allow => continue,
Decision::Redact {
reason,
mutated,
entities,
} => {
tracing::info!(stage = stage_name, reason = %reason, "request redacted");
req = mutated;
redaction_reasons.push(reason);
for entity in entities {
if !redaction_entities.contains(&entity) {
redaction_entities.push(entity);
}
}
}
Decision::Block { reason, code } => {
tracing::info!(
stage = stage_name,
reason = %reason,
code = %code,
"request blocked"
);
return Ok((Decision::Block { reason, code }, req));
}
}
}
if redaction_reasons.is_empty() {
Ok((Decision::Allow, req))
} else {
Ok((
Decision::Redact {
reason: redaction_reasons.join("; "),
mutated: req.clone(),
entities: redaction_entities,
},
req,
))
}
}
}
#[derive(Default)]
pub struct PipelineBuilder {
stages: Vec<(Arc<dyn Stage>, FailMode)>,
}
impl PipelineBuilder {
pub fn stage(self, s: impl Stage) -> Self {
self.stage_with_fail_mode(s, FailMode::Open)
}
pub fn arc_stage(self, s: Arc<dyn Stage>) -> Self {
self.arc_stage_with_fail_mode(s, FailMode::Open)
}
pub fn stage_with_fail_mode(self, s: impl Stage, fail_mode: FailMode) -> Self {
self.arc_stage_with_fail_mode(Arc::new(s), fail_mode)
}
pub fn arc_stage_with_fail_mode(mut self, s: Arc<dyn Stage>, fail_mode: FailMode) -> Self {
self.stages.push((s, fail_mode));
self
}
pub fn build(self) -> Pipeline {
Pipeline {
stages: self.stages,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::test_helpers::{clean_request, BlockingStage, PassthroughStage};
#[tokio::test]
async fn test_empty_pipeline_allows() {
let pipeline = PipelineBuilder::default().build();
let req = clean_request();
let (decision, _) = pipeline.run(req).await.unwrap();
assert_eq!(decision, Decision::Allow);
}
#[tokio::test]
async fn test_single_pass_stage_allows() {
let pipeline = PipelineBuilder::default().stage(PassthroughStage).build();
let req = clean_request();
let (decision, _) = pipeline.run(req).await.unwrap();
assert_eq!(decision, Decision::Allow);
}
#[tokio::test]
async fn test_blocking_stage_short_circuits() {
let pipeline = PipelineBuilder::default()
.stage(BlockingStage)
.stage(PassthroughStage)
.build();
let req = clean_request();
let (decision, _) = pipeline.run(req).await.unwrap();
assert!(matches!(decision, Decision::Block { .. }));
}
#[tokio::test]
async fn test_single_redacting_stage_returns_redact_not_allow() {
use crate::test_helpers::RedactingStage;
let pipeline = PipelineBuilder::default()
.stage(RedactingStage::new(
"PII detected and redacted: Email",
vec!["email".to_string()],
))
.build();
let req = clean_request();
let (decision, _) = pipeline.run(req).await.unwrap();
match decision {
Decision::Redact {
reason, entities, ..
} => {
assert_eq!(reason, "PII detected and redacted: Email");
assert_eq!(entities, vec!["email".to_string()]);
}
other => panic!("expected Decision::Redact, got {other:?}"),
}
}
#[tokio::test]
async fn test_multiple_redacting_stages_accumulate_reasons_and_entities() {
use crate::test_helpers::RedactingStage;
let pipeline = PipelineBuilder::default()
.stage(RedactingStage::with_name(
"stage_a",
"redacted email",
vec!["email".to_string()],
))
.stage(RedactingStage::with_name(
"stage_b",
"redacted phone",
vec!["phone".to_string()],
))
.build();
let req = clean_request();
let (decision, _) = pipeline.run(req).await.unwrap();
match decision {
Decision::Redact {
reason, entities, ..
} => {
assert_eq!(reason, "redacted email; redacted phone");
assert_eq!(entities, vec!["email".to_string(), "phone".to_string()]);
}
other => panic!("expected Decision::Redact, got {other:?}"),
}
}
#[tokio::test]
async fn test_duplicate_entities_across_stages_are_deduplicated() {
use crate::test_helpers::RedactingStage;
let pipeline = PipelineBuilder::default()
.stage(RedactingStage::with_name(
"stage_a",
"first pass",
vec!["email".to_string()],
))
.stage(RedactingStage::with_name(
"stage_b",
"second pass",
vec!["email".to_string(), "phone".to_string()],
))
.build();
let req = clean_request();
let (decision, _) = pipeline.run(req).await.unwrap();
match decision {
Decision::Redact { entities, .. } => {
assert_eq!(entities, vec!["email".to_string(), "phone".to_string()]);
}
other => panic!("expected Decision::Redact, got {other:?}"),
}
}
#[tokio::test]
async fn test_redact_then_block_returns_block_not_redact() {
use crate::test_helpers::RedactingStage;
let pipeline = PipelineBuilder::default()
.stage(RedactingStage::new(
"redacted first",
vec!["email".to_string()],
))
.stage(BlockingStage)
.build();
let req = clean_request();
let (decision, _) = pipeline.run(req).await.unwrap();
assert!(matches!(decision, Decision::Block { .. }));
}
#[tokio::test]
async fn test_redacted_request_is_passed_to_subsequent_stages() {
use crate::test_helpers::RedactingStage;
let pipeline = PipelineBuilder::default()
.stage(RedactingStage::new("redact", vec!["email".to_string()]))
.stage(PassthroughStage)
.build();
let req = clean_request();
let (decision, final_req) = pipeline.run(req.clone()).await.unwrap();
assert!(matches!(decision, Decision::Redact { .. }));
assert_eq!(final_req.user_text(), req.user_text());
}
#[tokio::test]
async fn test_pipeline_len() {
let pipeline = PipelineBuilder::default()
.stage(PassthroughStage)
.stage(PassthroughStage)
.build();
assert_eq!(pipeline.len(), 2);
assert!(!pipeline.is_empty());
}
#[tokio::test]
async fn test_run_with_observer_records_each_stage() {
let pipeline = PipelineBuilder::default()
.stage(PassthroughStage)
.stage(PassthroughStage)
.build();
let mut observed = Vec::new();
let (decision, _) = pipeline
.run_with_observer(clean_request(), |stage, elapsed| {
observed.push((stage.to_string(), elapsed));
})
.await
.unwrap();
assert_eq!(decision, Decision::Allow);
assert_eq!(observed.len(), 2);
assert!(observed.iter().all(|(name, _)| name == "passthrough"));
}
#[tokio::test]
async fn test_erroring_stage_fails_open_by_default_and_continues() {
use crate::test_helpers::ErrorStage;
let pipeline = PipelineBuilder::default()
.stage(ErrorStage)
.stage(BlockingStage)
.build();
let mut outcomes = Vec::new();
let (decision, _) = pipeline
.run_with_events(clean_request(), |e| outcomes.push((e.stage, e.outcome)))
.await
.unwrap();
assert!(matches!(decision, Decision::Block { .. }));
assert_eq!(
outcomes,
vec![
("error_stage", StageOutcome::ErrorAllowed),
("blocking", StageOutcome::Completed),
]
);
}
#[tokio::test]
async fn test_erroring_stage_with_fail_closed_blocks_and_short_circuits() {
use crate::test_helpers::ErrorStage;
let pipeline = PipelineBuilder::default()
.stage_with_fail_mode(ErrorStage, FailMode::Closed)
.stage(PassthroughStage)
.build();
let mut outcomes = Vec::new();
let (decision, _) = pipeline
.run_with_events(clean_request(), |e| outcomes.push((e.stage, e.outcome)))
.await
.unwrap();
match decision {
Decision::Block { code, reason } => {
assert_eq!(code.as_str(), STAGE_ERROR_CODE);
assert!(reason.contains("error_stage"));
assert!(!reason.contains("simulated stage failure"));
}
other => panic!("expected Block, got {other:?}"),
}
assert_eq!(outcomes, vec![("error_stage", StageOutcome::ErrorBlocked)]);
}
#[tokio::test]
async fn test_fail_mode_is_per_stage() {
use crate::test_helpers::ErrorStage;
let pipeline = PipelineBuilder::default()
.stage_with_fail_mode(ErrorStage, FailMode::Open)
.stage_with_fail_mode(PassthroughStage, FailMode::Closed)
.build();
let (decision, _) = pipeline.run(clean_request()).await.unwrap();
assert_eq!(decision, Decision::Allow);
}
#[tokio::test]
async fn test_pipeline_new_uses_fail_open() {
use crate::test_helpers::ErrorStage;
let pipeline = Pipeline::new(vec![Arc::new(ErrorStage)]);
let (decision, _) = pipeline.run(clean_request()).await.unwrap();
assert_eq!(decision, Decision::Allow);
}
#[tokio::test]
async fn test_run_with_observer_stops_after_block() {
let pipeline = PipelineBuilder::default()
.stage(BlockingStage)
.stage(PassthroughStage)
.build();
let mut observed = Vec::new();
let (decision, _) = pipeline
.run_with_observer(clean_request(), |stage, elapsed| {
observed.push((stage.to_string(), elapsed));
})
.await
.unwrap();
assert!(matches!(decision, Decision::Block { .. }));
assert_eq!(observed.len(), 1);
assert_eq!(observed[0].0, "blocking");
}
}