use crate::boundaries::{Guardrail, WriteAccess, WriteGovernance};
use crate::error::PeError;
use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionMetrics {
pub output_tokens: u32,
pub tool_calls_made: u32,
pub output_text: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WriteRecord {
pub destination: String,
pub key: String,
pub has_grant: bool,
}
impl Guardrail {
#[must_use = "this returns a Result that must be checked"]
pub fn validate(&self, metrics: &ExecutionMetrics) -> Result<(), PeError> {
match self {
Self::MaxOutputTokens(max) => {
if metrics.output_tokens > *max {
return Err(PeError::GuardrailViolation {
guardrail: format!("MaxOutputTokens({})", max),
details: format!("output was {} tokens", metrics.output_tokens),
});
}
}
Self::MaxToolCallsPerTurn(max) => {
if metrics.tool_calls_made > *max {
return Err(PeError::GuardrailViolation {
guardrail: format!("MaxToolCallsPerTurn({})", max),
details: format!("{} calls made", metrics.tool_calls_made),
});
}
}
Self::MustCiteSources | Self::NoCodeExecution => {}
Self::Custom { .. } => {}
}
Ok(())
}
}
impl WriteGovernance {
#[must_use = "this returns a Result that must be checked"]
pub fn validate_writes(&self, writes: &[WriteRecord]) -> Result<(), PeError> {
for write in writes {
let access = match write.destination.as_str() {
"own_memory" => &self.own_memory,
"collective" => &self.collective,
"vault" => &self.vault,
"task_store" => &self.task_store,
_ => continue, };
match access {
WriteAccess::ReadOnly => {
return Err(PeError::WriteGovernanceViolation {
destination: write.destination.clone(),
reason: "ReadOnly — no writes permitted".into(),
});
}
WriteAccess::RequiresGrant if !write.has_grant => {
return Err(PeError::WriteGovernanceViolation {
destination: write.destination.clone(),
reason: "RequiresGrant — no grant obtained".into(),
});
}
_ => {} }
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::boundaries::WriteGovernance;
fn metrics(output_tokens: u32, tool_calls: u32) -> ExecutionMetrics {
ExecutionMetrics {
output_tokens,
tool_calls_made: tool_calls,
output_text: String::new(),
}
}
#[test]
fn max_output_tokens_within_limit() {
let g = Guardrail::MaxOutputTokens(2000);
assert!(g.validate(&metrics(1500, 0)).is_ok());
}
#[test]
fn max_output_tokens_violation() {
let g = Guardrail::MaxOutputTokens(2000);
let err = g.validate(&metrics(2500, 0)).unwrap_err();
assert!(matches!(err, PeError::GuardrailViolation { .. }));
}
#[test]
fn max_tool_calls_within_limit() {
let g = Guardrail::MaxToolCallsPerTurn(5);
assert!(g.validate(&metrics(0, 3)).is_ok());
}
#[test]
fn max_tool_calls_violation() {
let g = Guardrail::MaxToolCallsPerTurn(5);
let err = g.validate(&metrics(0, 8)).unwrap_err();
assert!(matches!(err, PeError::GuardrailViolation { .. }));
}
#[test]
fn write_governance_free_allows_all() {
let gov = WriteGovernance::default(); let writes = vec![WriteRecord {
destination: "own_memory".into(),
key: "test".into(),
has_grant: false,
}];
assert!(gov.validate_writes(&writes).is_ok());
}
#[test]
fn write_governance_read_only_rejects() {
let gov = WriteGovernance {
vault: WriteAccess::ReadOnly,
..Default::default()
};
let writes = vec![WriteRecord {
destination: "vault".into(),
key: "secret".into(),
has_grant: false,
}];
let err = gov.validate_writes(&writes).unwrap_err();
assert!(matches!(err, PeError::WriteGovernanceViolation { .. }));
}
#[test]
fn write_governance_requires_grant_without_grant_rejects() {
let gov = WriteGovernance {
collective: WriteAccess::RequiresGrant,
..Default::default()
};
let writes = vec![WriteRecord {
destination: "collective".into(),
key: "notes".into(),
has_grant: false,
}];
assert!(gov.validate_writes(&writes).is_err());
let writes_granted = vec![WriteRecord {
destination: "collective".into(),
key: "notes".into(),
has_grant: true,
}];
assert!(gov.validate_writes(&writes_granted).is_ok());
}
}