use async_trait::async_trait;
use thiserror::Error;
use crate::receipt::RedactionRecord;
use crate::types::{Brief, BriefContext};
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GovernanceVerdict {
Allow,
Redacted(Vec<RedactionRecord>),
Rejected(String),
}
impl GovernanceVerdict {
pub fn is_redacted(&self) -> bool {
matches!(self, GovernanceVerdict::Redacted(_))
}
pub fn is_rejected(&self) -> bool {
matches!(self, GovernanceVerdict::Rejected(_))
}
}
#[derive(Debug, Error)]
pub enum GovernanceError {
#[error("governance backend error: {0}")]
Backend(String),
#[error("governance error: {0}")]
Other(String),
}
#[async_trait]
pub trait Governance: Send + Sync {
async fn review(
&self,
draft: &mut Brief,
ctx: &BriefContext,
) -> Result<GovernanceVerdict, GovernanceError>;
}
#[derive(Debug, Default, Clone, Copy)]
pub struct NoOpGovernance;
impl NoOpGovernance {
pub fn new() -> Self {
NoOpGovernance
}
}
#[async_trait]
impl Governance for NoOpGovernance {
async fn review(
&self,
_draft: &mut Brief,
_ctx: &BriefContext,
) -> Result<GovernanceVerdict, GovernanceError> {
Ok(GovernanceVerdict::Allow)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::receipt::BriefReceipt;
use crate::types::{BriefMessage, Role, SourceId, TokenBudget};
fn empty_draft() -> Brief {
Brief {
system: None,
messages: Vec::new(),
tools: Vec::new(),
receipt: BriefReceipt::empty(),
}
}
#[tokio::test]
async fn no_op_always_allows() {
let mut draft = empty_draft();
let ctx = BriefContext::new(TokenBudget::default());
let v = NoOpGovernance.review(&mut draft, &ctx).await.expect("ok");
assert_eq!(v, GovernanceVerdict::Allow);
}
struct SecretScrubber;
#[async_trait]
impl Governance for SecretScrubber {
async fn review(
&self,
draft: &mut Brief,
_ctx: &BriefContext,
) -> Result<GovernanceVerdict, GovernanceError> {
let mut records = Vec::new();
for msg in &mut draft.messages {
if let BriefMessage::Text {
role: Role::User,
content,
source,
} = msg
{
if content.contains("secret") {
records.push(RedactionRecord {
source: source.clone(),
rule: "rule:no_user_secrets".into(),
});
*content = content.replace("secret", "[REDACTED]");
}
}
}
Ok(if records.is_empty() {
GovernanceVerdict::Allow
} else {
GovernanceVerdict::Redacted(records)
})
}
}
#[tokio::test]
async fn redacted_verdict_carries_records() {
let mut draft = empty_draft();
draft.messages.push(BriefMessage::Text {
role: Role::User,
content: "my secret is hunter2".into(),
source: SourceId::new("user_message"),
});
let ctx = BriefContext::new(TokenBudget::default());
let v = SecretScrubber.review(&mut draft, &ctx).await.expect("ok");
match v {
GovernanceVerdict::Redacted(records) => {
assert_eq!(records.len(), 1);
assert_eq!(records[0].source, SourceId::new("user_message"));
assert_eq!(records[0].rule, "rule:no_user_secrets");
}
other => panic!("expected Redacted, got {other:?}"),
}
match &draft.messages[0] {
BriefMessage::Text { content, .. } => {
assert!(content.contains("[REDACTED]"));
assert!(!content.contains("secret"));
}
_ => panic!(),
}
}
#[tokio::test]
async fn allow_returned_when_nothing_to_redact() {
let mut draft = empty_draft();
draft.messages.push(BriefMessage::Text {
role: Role::User,
content: "hi".into(),
source: SourceId::new("user_message"),
});
let ctx = BriefContext::new(TokenBudget::default());
let v = SecretScrubber.review(&mut draft, &ctx).await.expect("ok");
assert_eq!(v, GovernanceVerdict::Allow);
}
struct AlwaysReject;
#[async_trait]
impl Governance for AlwaysReject {
async fn review(
&self,
_draft: &mut Brief,
_ctx: &BriefContext,
) -> Result<GovernanceVerdict, GovernanceError> {
Ok(GovernanceVerdict::Rejected("locked down".into()))
}
}
#[tokio::test]
async fn rejected_verdict_carries_reason() {
let mut draft = empty_draft();
let ctx = BriefContext::new(TokenBudget::default());
let v = AlwaysReject.review(&mut draft, &ctx).await.expect("ok");
match &v {
GovernanceVerdict::Rejected(reason) => assert_eq!(reason, "locked down"),
other => panic!("expected Rejected, got {other:?}"),
}
assert!(v.is_rejected());
assert!(!v.is_redacted());
}
#[tokio::test]
async fn backend_error_propagates() {
struct Broken;
#[async_trait]
impl Governance for Broken {
async fn review(
&self,
_draft: &mut Brief,
_ctx: &BriefContext,
) -> Result<GovernanceVerdict, GovernanceError> {
Err(GovernanceError::Backend("policy db unreachable".into()))
}
}
let mut draft = empty_draft();
let ctx = BriefContext::new(TokenBudget::default());
let err = Broken.review(&mut draft, &ctx).await.expect_err("err");
match err {
GovernanceError::Backend(msg) => assert!(msg.contains("policy db")),
other => panic!("expected Backend, got {other:?}"),
}
}
#[test]
fn verdict_helpers() {
assert!(!GovernanceVerdict::Allow.is_redacted());
assert!(!GovernanceVerdict::Allow.is_rejected());
assert!(GovernanceVerdict::Redacted(Vec::new()).is_redacted());
assert!(GovernanceVerdict::Rejected("x".into()).is_rejected());
}
}