telos_agent/agent/policies/
mod.rs

1//! Stable, semantic policy extension points for the agent runtime.
2
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use serde::{Deserialize, Serialize};
7
8use crate::error::AgentError;
9use crate::model::message::{Message, ToolCall, ToolResult};
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
12#[serde(rename_all = "snake_case")]
13pub enum SessionMode {
14    Create,
15    Resume,
16}
17
18#[derive(Debug, Clone, PartialEq, Eq)]
19pub enum PolicyPoint {
20    SessionStart { mode: Option<SessionMode> },
21    ModelResponse,
22    ToolBeforeInvoke { matcher: Option<String> },
23    ToolAfterInvoke { matcher: Option<String> },
24    TurnBeforeFinish,
25}
26
27#[derive(Debug, Clone, Serialize)]
28#[serde(tag = "point", rename_all = "snake_case")]
29pub enum PolicyContext {
30    SessionStart { session_id: String, mode: SessionMode, message_count: usize },
31    ModelResponse { session_id: String, turn_id: u64, iteration: usize, message: Message },
32    ToolBeforeInvoke { session_id: String, turn_id: u64, call: ToolCall },
33    ToolAfterInvoke { session_id: String, turn_id: u64, call: ToolCall, result: ToolResult },
34    TurnBeforeFinish { session_id: String, turn_id: u64, message: Message },
35}
36
37#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
38#[serde(tag = "decision", rename_all = "snake_case")]
39pub enum PolicyDecision {
40    Continue,
41    Reject { reason: String },
42}
43
44#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
45pub struct PolicyOutcome {
46    #[serde(flatten)]
47    pub decision: PolicyDecision,
48    #[serde(default)]
49    pub feedback: Vec<String>,
50}
51
52impl PolicyOutcome {
53    pub fn continue_() -> Self {
54        Self { decision: PolicyDecision::Continue, feedback: Vec::new() }
55    }
56}
57
58#[async_trait]
59pub trait Policy: Send + Sync {
60    fn name(&self) -> &str;
61    async fn evaluate(&self, context: &PolicyContext) -> Result<PolicyOutcome, AgentError>;
62}
63
64#[derive(Clone)]
65pub struct PolicyEntry {
66    pub point: PolicyPoint,
67    pub policy: Arc<dyn Policy>,
68}
69
70impl std::fmt::Debug for PolicyEntry {
71    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
72        f.debug_struct("PolicyEntry")
73            .field("point", &self.point)
74            .field("policy", &self.policy.name())
75            .finish()
76    }
77}
78
79#[derive(Clone, Default)]
80pub struct PolicyRegistry {
81    entries: Vec<PolicyEntry>,
82}
83
84impl std::fmt::Debug for PolicyRegistry {
85    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
86        f.debug_struct("PolicyRegistry").field("count", &self.entries.len()).finish()
87    }
88}
89
90impl PolicyRegistry {
91    pub fn new() -> Self {
92        Self::default()
93    }
94    pub fn register(&mut self, entry: PolicyEntry) {
95        self.entries.push(entry);
96    }
97
98    pub fn session_start(&self, mode: SessionMode) -> Vec<Arc<dyn Policy>> {
99        self.entries
100            .iter()
101            .filter_map(|entry| match entry.point {
102                PolicyPoint::SessionStart { mode: configured }
103                    if configured.is_none_or(|value| value == mode) =>
104                {
105                    Some(entry.policy.clone())
106                }
107                _ => None,
108            })
109            .collect()
110    }
111
112    pub fn model_response(&self) -> Vec<Arc<dyn Policy>> {
113        self.by_point(|point| matches!(point, PolicyPoint::ModelResponse))
114    }
115
116    pub fn turn_before_finish(&self) -> Vec<Arc<dyn Policy>> {
117        self.by_point(|point| matches!(point, PolicyPoint::TurnBeforeFinish))
118    }
119
120    pub fn tool_before(&self, name: &str) -> Vec<Arc<dyn Policy>> {
121        self.tool_policies(name, false)
122    }
123
124    pub fn tool_after(&self, name: &str) -> Vec<Arc<dyn Policy>> {
125        self.tool_policies(name, true)
126    }
127
128    fn by_point(&self, predicate: impl Fn(&PolicyPoint) -> bool) -> Vec<Arc<dyn Policy>> {
129        self.entries
130            .iter()
131            .filter(|entry| predicate(&entry.point))
132            .map(|entry| entry.policy.clone())
133            .collect()
134    }
135
136    fn tool_policies(&self, name: &str, after: bool) -> Vec<Arc<dyn Policy>> {
137        self.entries
138            .iter()
139            .filter_map(|entry| {
140                let matcher = match &entry.point {
141                    PolicyPoint::ToolBeforeInvoke { matcher } if !after => matcher,
142                    PolicyPoint::ToolAfterInvoke { matcher } if after => matcher,
143                    _ => return None,
144                };
145                matcher
146                    .as_ref()
147                    .is_none_or(|pattern| {
148                        glob::Pattern::new(pattern).is_ok_and(|glob| glob.matches(name))
149                    })
150                    .then(|| entry.policy.clone())
151            })
152            .collect()
153    }
154}
155
156#[cfg(test)]
157mod tests {
158    use super::*;
159
160    struct Named(&'static str);
161    #[async_trait]
162    impl Policy for Named {
163        fn name(&self) -> &str {
164            self.0
165        }
166        async fn evaluate(&self, _: &PolicyContext) -> Result<PolicyOutcome, AgentError> {
167            Ok(PolicyOutcome::continue_())
168        }
169    }
170
171    #[test]
172    fn semantic_points_filter_without_exposing_runtime_phases() {
173        let mut registry = PolicyRegistry::new();
174        registry.register(PolicyEntry {
175            point: PolicyPoint::ToolBeforeInvoke { matcher: Some("Bash*".into()) },
176            policy: Arc::new(Named("specific")),
177        });
178        registry.register(PolicyEntry {
179            point: PolicyPoint::ToolBeforeInvoke { matcher: None },
180            policy: Arc::new(Named("all")),
181        });
182        let names: Vec<_> = registry
183            .tool_before("Bash")
184            .into_iter()
185            .map(|policy| policy.name().to_string())
186            .collect();
187        assert_eq!(names, ["specific", "all"]);
188        assert_eq!(registry.tool_before("Grep").len(), 1);
189        assert!(registry.tool_after("Bash").is_empty());
190    }
191
192    #[test]
193    fn outcome_uses_flat_command_protocol() {
194        let outcome: PolicyOutcome = serde_json::from_value(serde_json::json!({
195            "decision": "reject", "reason": "blocked", "feedback": ["note"]
196        }))
197        .unwrap();
198        assert_eq!(outcome.feedback, ["note"]);
199        assert_eq!(outcome.decision, PolicyDecision::Reject { reason: "blocked".into() });
200    }
201}
telos_agent/agent/policies/mod.rs

telos_agent/agent/policies/
mod.rs
