mermaid_cli/providers/
auto_classifier.rs

1//! LLM-backed safety vetting for `SafetyMode::Auto`.
2//!
3//! Under Auto mode the rule engine (`mermaid-runtime`) classifies a
4//! borderline action as [`PolicyDecision::Classify`] and defers the
5//! allow/escalate call to a model. This module is that model call. It lives
6//! in `mermaid-cli` (not the runtime crate) because the runtime is
7//! deliberately model-free — the policy gate injects an
8//! `Arc<dyn AutoClassifier>` into [`ExecContext`] and awaits [`AutoClassifier::vet`]
9//! before letting a borderline action run.
10//!
11//! Authority is **allow-or-escalate only** — the classifier never hard-blocks
12//! (destructive patterns are already denied by the rule engine), and any
13//! error / timeout / unparseable reply **fails safe** to "escalate to human".
14//!
15//! [`PolicyDecision::Classify`]: crate::runtime::PolicyDecision
16//! [`ExecContext`]: crate::providers::ctx::ExecContext
17
18use std::sync::Arc;
19use std::time::Duration;
20
21use async_trait::async_trait;
22use tokio_util::sync::CancellationToken;
23
24use crate::domain::{ChatRequest, TurnId};
25use crate::models::{ChatMessage, ReasoningLevel};
26use crate::providers::factory::ProviderFactory;
27
28/// How long to wait for the classifier before failing safe (escalating).
29const VET_TIMEOUT: Duration = Duration::from_secs(10);
30/// Verdicts are tiny — `ALLOW` or `ESCALATE: <short reason>`.
31const VET_MAX_TOKENS: usize = 150;
32
33const SYSTEM_PROMPT: &str = "You are a safety reviewer for an AI coding agent running in \"auto\" mode. \
34The agent has already decided to take an action; your job is to wave through the routine, aligned ones \
35and stop only the genuinely risky or off-task ones. Bias strongly toward ALLOW: most actions an engineer \
36would expect while pursuing the stated goal should pass. ESCALATE only when an action is destructive, \
37leaks secrets or credentials, reaches untrusted network endpoints, modifies shared/production \
38infrastructure, or clearly does not serve the user's goal. When in doubt about real risk, ESCALATE. \
39Reply with EXACTLY one line: `ALLOW` or `ESCALATE: <short reason>`.";
40
41/// One action to vet, expressed in the classifier's terms.
42#[derive(Debug, Clone)]
43pub struct VetRequest {
44    pub tool: String,
45    pub summary: String,
46    pub command: Option<String>,
47    pub path: Option<String>,
48    /// The user's stated intent for the turn (latest user message), if known.
49    pub intent: Option<String>,
50    /// Absolute working directory, for context.
51    pub workdir: String,
52    pub turn: TurnId,
53    /// Turn cancellation — a Ctrl+C aborts the vet (which then fails safe).
54    pub token: CancellationToken,
55}
56
57/// The classifier's verdict. `allow == false` means "escalate to a human".
58#[derive(Debug, Clone, PartialEq, Eq)]
59pub struct VetVerdict {
60    pub allow: bool,
61    pub reason: String,
62}
63
64impl VetVerdict {
65    pub fn allow() -> Self {
66        Self {
67            allow: true,
68            reason: String::new(),
69        }
70    }
71    pub fn escalate(reason: impl Into<String>) -> Self {
72        Self {
73            allow: false,
74            reason: reason.into(),
75        }
76    }
77}
78
79/// Vets a borderline action against the user's intent. Implementors must be
80/// cheap to clone-share (`Arc`) and safe to call concurrently.
81#[async_trait]
82pub trait AutoClassifier: Send + Sync {
83    async fn vet(&self, req: &VetRequest) -> VetVerdict;
84}
85
86/// Production classifier: builds a focused one-shot prompt and runs it through
87/// a provider (by default the session's own model).
88pub struct ModelAutoClassifier {
89    providers: Arc<ProviderFactory>,
90    model_id: String,
91}
92
93impl ModelAutoClassifier {
94    pub fn new(providers: Arc<ProviderFactory>, model_id: String) -> Self {
95        Self {
96            providers,
97            model_id,
98        }
99    }
100
101    fn build_request(&self, req: &VetRequest) -> ChatRequest {
102        let action = describe_action(req);
103        let intent = req
104            .intent
105            .as_deref()
106            .map(str::trim)
107            .filter(|s| !s.is_empty())
108            .unwrap_or("(no explicit goal stated this turn)");
109        let user = format!(
110            "Working directory: {wd}\n\nUser's current goal:\n{intent}\n\nProposed action:\n{action}\n\n\
111             Does this action plausibly serve the user's goal and look safe to run automatically?",
112            wd = req.workdir,
113            intent = intent,
114            action = action,
115        );
116        ChatRequest {
117            model_id: self.model_id.clone(),
118            messages: vec![ChatMessage::user(user)],
119            system_prompt: SYSTEM_PROMPT.to_string(),
120            instructions: None,
121            // The judgment is simple and we want it fast/cheap — no
122            // extended thinking.
123            reasoning: ReasoningLevel::None,
124            temperature: 0.0,
125            max_tokens: VET_MAX_TOKENS,
126            tools: Vec::new(),
127        }
128    }
129}
130
131#[async_trait]
132impl AutoClassifier for ModelAutoClassifier {
133    async fn vet(&self, req: &VetRequest) -> VetVerdict {
134        let request = self.build_request(req);
135        let providers = Arc::clone(&self.providers);
136        let model_id = self.model_id.clone();
137        let turn = req.turn;
138        let token = req.token.clone();
139
140        let call = async move {
141            let provider = providers.resolve(&model_id).await?;
142            let (text, _usage) =
143                crate::providers::model::collect_text(provider, turn, request, token).await?;
144            Ok::<String, crate::models::ModelError>(text)
145        };
146
147        match tokio::time::timeout(VET_TIMEOUT, call).await {
148            Ok(Ok(text)) => parse_verdict(&text),
149            Ok(Err(err)) => VetVerdict::escalate(format!("classifier unavailable: {err}")),
150            Err(_) => VetVerdict::escalate("classifier timed out"),
151        }
152    }
153}
154
155fn describe_action(req: &VetRequest) -> String {
156    if let Some(cmd) = &req.command {
157        format!("Tool `{}` will run a shell command:\n  {}", req.tool, cmd)
158    } else if let Some(path) = &req.path {
159        format!(
160            "Tool `{}` will act on path `{}` ({})",
161            req.tool, path, req.summary
162        )
163    } else {
164        format!("Tool `{}`: {}", req.tool, req.summary)
165    }
166}
167
168/// Parse the classifier's reply. Tolerant: matches a leading `ALLOW` /
169/// `ESCALATE` case-insensitively; anything unrecognized fails safe (escalate).
170fn parse_verdict(text: &str) -> VetVerdict {
171    let trimmed = text.trim();
172    if trimmed.is_empty() {
173        return VetVerdict::escalate("classifier returned an empty response");
174    }
175    let upper = trimmed.to_ascii_uppercase();
176    if upper.starts_with("ALLOW") {
177        return VetVerdict::allow();
178    }
179    if upper.starts_with("ESCALATE") {
180        let reason = trimmed
181            .split_once(':')
182            .map(|(_, r)| r.trim())
183            .filter(|r| !r.is_empty())
184            .map(clip)
185            .unwrap_or_else(|| "flagged by the safety classifier".to_string());
186        return VetVerdict::escalate(reason);
187    }
188    VetVerdict::escalate(format!("unrecognized classifier reply: {}", clip(trimmed)))
189}
190
191/// Cap a reason string at a sane length on a char boundary.
192fn clip(s: &str) -> String {
193    const MAX: usize = 160;
194    if s.len() <= MAX {
195        return s.to_string();
196    }
197    let cut = s.floor_char_boundary(MAX);
198    format!("{}…", &s[..cut])
199}
200
201#[cfg(test)]
202mod tests {
203    use super::*;
204
205    #[test]
206    fn allow_parses() {
207        assert!(parse_verdict("ALLOW").allow);
208        assert!(parse_verdict("  allow\n").allow);
209        assert!(parse_verdict("Allow — looks fine").allow);
210    }
211
212    #[test]
213    fn escalate_parses_with_reason() {
214        let v = parse_verdict("ESCALATE: pipes a remote script into sh");
215        assert!(!v.allow);
216        assert_eq!(v.reason, "pipes a remote script into sh");
217    }
218
219    #[test]
220    fn escalate_without_reason_has_default() {
221        let v = parse_verdict("escalate");
222        assert!(!v.allow);
223        assert!(!v.reason.is_empty());
224    }
225
226    #[test]
227    fn garbage_and_empty_fail_safe() {
228        // Anything we can't read is treated as "escalate", never "allow".
229        for reply in ["", "   ", "maybe?", "yes", "no", "I think it's fine"] {
230            assert!(
231                !parse_verdict(reply).allow,
232                "expected escalate (fail-safe) for {reply:?}",
233            );
234        }
235    }
236}
mermaid_cli/providers/auto_classifier.rs

mermaid_cli/providers/
auto_classifier.rs
