1use std::{
42 collections::HashMap,
43 sync::{Arc, Mutex},
44};
45
46use agent_client_protocol::schema::v1::{SessionMode, SessionModeId, SessionModeState};
47
48use basis::approval::{ApprovalAnswer, ApprovalDecision, ApprovalRequest, Approver};
49
50#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
56pub enum ApprovalMode {
57 Always,
59 #[default]
63 Prompt,
64 Never,
67}
68
69const ALWAYS: &str = "always";
73const PROMPT: &str = "prompt";
74const NEVER: &str = "never";
75
76fn mode_for(id: &str) -> Option<ApprovalMode> {
78 match id {
79 ALWAYS => Some(ApprovalMode::Always),
80 PROMPT => Some(ApprovalMode::Prompt),
81 NEVER => Some(ApprovalMode::Never),
82 _ => None,
83 }
84}
85
86fn mode_id(mode: ApprovalMode) -> SessionModeId {
87 SessionModeId::new(match mode {
88 ApprovalMode::Always => ALWAYS,
89 ApprovalMode::Prompt => PROMPT,
90 ApprovalMode::Never => NEVER,
91 })
92}
93
94fn describe(mode: ApprovalMode) -> SessionMode {
96 let (name, description) = match mode {
97 ApprovalMode::Always => (
98 "Always allow",
99 "Act without asking. What a confined or unattended session wants.",
100 ),
101 ApprovalMode::Prompt => (
102 "Ask each time",
103 "Ask before anything that changes state outside this process.",
104 ),
105 ApprovalMode::Never => (
106 "Read only",
107 "Refuse anything that changes state outside this process.",
108 ),
109 };
110
111 SessionMode::new(mode_id(mode), name).description(description)
112}
113
114#[derive(Debug, Clone, Copy, PartialEq, Eq)]
116pub enum ModeError {
117 Unknown,
119 NotOffered,
121}
122
123impl std::fmt::Display for ModeError {
124 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
125 match self {
126 Self::Unknown => f.write_str("unknown mode"),
127 Self::NotOffered => {
128 f.write_str("this session was opened read-only and cannot change mode")
129 }
130 }
131 }
132}
133
134#[derive(Clone)]
141pub struct SessionModes {
142 inner: Arc<Mutex<State>>,
143}
144
145struct State {
146 current: ApprovalMode,
147 switchable: bool,
149 remembered: HashMap<String, bool>,
151}
152
153impl SessionModes {
154 pub fn new(initial: ApprovalMode) -> Self {
163 Self {
164 inner: Arc::new(Mutex::new(State {
165 current: initial,
166 switchable: !matches!(initial, ApprovalMode::Never),
167 remembered: HashMap::new(),
168 })),
169 }
170 }
171
172 pub fn current(&self) -> ApprovalMode {
173 self.lock().current
174 }
175
176 pub fn state(&self) -> SessionModeState {
179 let state = self.lock();
180 let available = if state.switchable {
181 vec![
182 describe(ApprovalMode::Always),
183 describe(ApprovalMode::Prompt),
184 describe(ApprovalMode::Never),
185 ]
186 } else {
187 vec![describe(state.current)]
188 };
189
190 SessionModeState::new(mode_id(state.current), available)
191 }
192
193 pub fn set(&self, id: &SessionModeId) -> Result<ApprovalMode, ModeError> {
200 let mode = mode_for(&id.0).ok_or(ModeError::Unknown)?;
201
202 let mut state = self.lock();
203 if !state.switchable && mode != state.current {
204 return Err(ModeError::NotOffered);
205 }
206
207 state.current = mode;
208 state.remembered.clear();
209 Ok(mode)
210 }
211
212 fn remember(&self, tool_name: &str, allow: bool) {
213 self.lock().remembered.insert(tool_name.to_string(), allow);
214 }
215
216 fn remembered(&self, tool_name: &str) -> Option<bool> {
217 self.lock().remembered.get(tool_name).copied()
218 }
219
220 fn lock(&self) -> std::sync::MutexGuard<'_, State> {
221 self.inner
225 .lock()
226 .unwrap_or_else(|poisoned| poisoned.into_inner())
227 }
228}
229
230pub struct ModedApprover<A> {
233 modes: SessionModes,
234 inner: A,
235}
236
237impl<A> ModedApprover<A> {
238 pub fn new(modes: SessionModes, inner: A) -> Self {
239 Self { modes, inner }
240 }
241}
242
243#[async_trait::async_trait]
244impl<A: Approver> Approver for ModedApprover<A> {
245 async fn approve(&mut self, request: &ApprovalRequest) -> ApprovalAnswer {
246 match self.modes.current() {
250 ApprovalMode::Always => ApprovalDecision::Allow.into(),
251 ApprovalMode::Never => ApprovalAnswer::new(ApprovalDecision::Deny).because(format!(
252 "{} changes state outside this process, and this session is set to refuse that",
253 request.tool_name
254 )),
255 ApprovalMode::Prompt => self.ask(request).await,
256 }
257 }
258}
259
260impl<A: Approver> ModedApprover<A> {
261 async fn ask(&mut self, request: &ApprovalRequest) -> ApprovalAnswer {
262 if let Some(allow) = self.modes.remembered(&request.tool_name) {
263 return if allow {
264 ApprovalDecision::Allow.into()
265 } else {
266 ApprovalAnswer::new(ApprovalDecision::Deny).because(format!(
267 "{} was refused earlier in this session, and that answer still stands",
268 request.tool_name
269 ))
270 };
271 }
272
273 let answer = self.inner.approve(request).await;
278 match answer.decision {
279 ApprovalDecision::AllowForSession => {
280 self.modes.remember(&request.tool_name, true);
281 ApprovalAnswer {
282 decision: ApprovalDecision::Allow,
283 ..answer
284 }
285 }
286 ApprovalDecision::DenyForSession => {
287 self.modes.remember(&request.tool_name, false);
288 ApprovalAnswer {
289 decision: ApprovalDecision::Deny,
290 ..answer
291 }
292 }
293 _ => answer,
294 }
295 }
296}
297
298#[cfg(test)]
299mod tests {
300 use super::*;
301 use serde_json::json;
302 use std::sync::atomic::{AtomicUsize, Ordering};
303
304 struct Counting {
307 asked: Arc<AtomicUsize>,
308 answer: ApprovalDecision,
309 }
310
311 #[async_trait::async_trait]
312 impl Approver for Counting {
313 async fn approve(&mut self, _request: &ApprovalRequest) -> ApprovalAnswer {
314 self.asked.fetch_add(1, Ordering::SeqCst);
315 self.answer.into()
316 }
317 }
318
319 fn request(tool_name: &str) -> ApprovalRequest {
320 ApprovalRequest {
321 request_id: "r1".to_string(),
322 tool_call_id: "c1".to_string(),
323 tool_name: tool_name.to_string(),
324 description: "wants to write".to_string(),
325 input: json!({}),
326 }
327 }
328
329 fn gate(
330 initial: ApprovalMode,
331 answer: ApprovalDecision,
332 ) -> (SessionModes, ModedApprover<Counting>, Arc<AtomicUsize>) {
333 let modes = SessionModes::new(initial);
334 let asked = Arc::new(AtomicUsize::new(0));
335 let approver = ModedApprover::new(
336 modes.clone(),
337 Counting {
338 asked: Arc::clone(&asked),
339 answer,
340 },
341 );
342 (modes, approver, asked)
343 }
344
345 #[test]
346 fn every_offered_mode_maps_back_to_one_lan_can_read() {
347 for mode in SessionModes::new(ApprovalMode::Prompt)
350 .state()
351 .available_modes
352 {
353 assert!(
354 mode_for(&mode.id.0).is_some(),
355 "offered {} but cannot read it back",
356 mode.id.0
357 );
358 }
359 }
360
361 #[test]
362 fn the_state_reports_the_current_mode_and_all_three() {
363 let state = SessionModes::new(ApprovalMode::Prompt).state();
364
365 assert_eq!(&*state.current_mode_id.0, PROMPT);
366 assert_eq!(state.available_modes.len(), 3);
367 }
368
369 #[test]
370 fn a_read_only_session_offers_nothing_else() {
371 let modes = SessionModes::new(ApprovalMode::Never);
372 let state = modes.state();
373
374 assert_eq!(state.available_modes.len(), 1);
375 assert_eq!(&*state.current_mode_id.0, NEVER);
376 assert_eq!(
377 modes.set(&SessionModeId::new(ALWAYS)),
378 Err(ModeError::NotOffered),
379 "a client cannot lift a prohibition it was never given"
380 );
381 }
382
383 #[test]
384 fn switching_reports_the_new_mode() {
385 let modes = SessionModes::new(ApprovalMode::Prompt);
386
387 assert_eq!(
388 modes.set(&SessionModeId::new(ALWAYS)),
389 Ok(ApprovalMode::Always)
390 );
391 assert_eq!(modes.current(), ApprovalMode::Always);
392 }
393
394 #[test]
395 fn an_unknown_mode_is_refused() {
396 let modes = SessionModes::new(ApprovalMode::Prompt);
397
398 assert_eq!(
399 modes.set(&SessionModeId::new("architect")),
400 Err(ModeError::Unknown)
401 );
402 assert_eq!(
403 modes.current(),
404 ApprovalMode::Prompt,
405 "a refused switch must leave the session where it was"
406 );
407 }
408
409 #[tokio::test]
410 async fn allow_and_refuse_answer_without_asking() {
411 for (mode, expected) in [
412 (ApprovalMode::Always, ApprovalDecision::Allow),
413 (ApprovalMode::Never, ApprovalDecision::Deny),
414 ] {
415 let (_modes, mut approver, asked) = gate(mode, ApprovalDecision::Allow);
416
417 assert_eq!(approver.approve(&request("shell")).await.decision, expected);
418 assert_eq!(
419 asked.load(Ordering::SeqCst),
420 0,
421 "{mode:?} has nothing to ask about"
422 );
423 }
424 }
425
426 #[tokio::test]
427 async fn a_read_only_session_says_so_when_it_refuses() {
428 let (_modes, mut approver, _asked) = gate(ApprovalMode::Never, ApprovalDecision::Allow);
431
432 assert_eq!(
433 approver.approve(&request("shell")).await.reason.as_deref(),
434 Some(
435 "shell changes state outside this process, \
436 and this session is set to refuse that"
437 )
438 );
439 }
440
441 #[tokio::test]
442 async fn asking_puts_the_request_to_the_client() {
443 let (_modes, mut approver, asked) = gate(ApprovalMode::Prompt, ApprovalDecision::Allow);
444
445 assert_eq!(
446 approver.approve(&request("shell")).await.decision,
447 ApprovalDecision::Allow
448 );
449 assert_eq!(asked.load(Ordering::SeqCst), 1);
450 }
451
452 #[tokio::test]
453 async fn an_answer_for_the_session_is_not_asked_twice() {
454 let (_modes, mut approver, asked) =
455 gate(ApprovalMode::Prompt, ApprovalDecision::AllowForSession);
456
457 assert_eq!(
460 approver.approve(&request("shell")).await.decision,
461 ApprovalDecision::Allow
462 );
463 assert_eq!(
464 approver.approve(&request("shell")).await.decision,
465 ApprovalDecision::Allow
466 );
467 assert_eq!(asked.load(Ordering::SeqCst), 1);
468
469 assert_eq!(
471 approver.approve(&request("files")).await.decision,
472 ApprovalDecision::Allow
473 );
474 assert_eq!(asked.load(Ordering::SeqCst), 2);
475 }
476
477 #[tokio::test]
478 async fn changing_mode_forgets_what_was_allowed_for_the_session() {
479 let (modes, mut approver, _asked) =
480 gate(ApprovalMode::Prompt, ApprovalDecision::AllowForSession);
481
482 approver.approve(&request("shell")).await;
483 modes.set(&SessionModeId::new(NEVER)).expect("switches");
484
485 assert_eq!(
486 approver.approve(&request("shell")).await.decision,
487 ApprovalDecision::Deny,
488 "a stale allow must not survive the mode that replaced it"
489 );
490
491 modes.set(&SessionModeId::new(PROMPT)).expect("switches");
493 assert_eq!(
494 approver.approve(&request("shell")).await.decision,
495 ApprovalDecision::Allow,
496 "the client is asked again, and answered again"
497 );
498 }
499
500 #[tokio::test]
501 async fn a_refusal_for_the_session_is_also_remembered() {
502 let (_modes, mut approver, asked) =
503 gate(ApprovalMode::Prompt, ApprovalDecision::DenyForSession);
504
505 assert_eq!(
506 approver.approve(&request("shell")).await.decision,
507 ApprovalDecision::Deny
508 );
509
510 let repeated = approver.approve(&request("shell")).await;
511 assert_eq!(repeated.decision, ApprovalDecision::Deny);
512 assert_eq!(
513 repeated.reason.as_deref(),
514 Some("shell was refused earlier in this session, and that answer still stands"),
515 "a remembered refusal still owes the model a reason"
516 );
517 assert_eq!(asked.load(Ordering::SeqCst), 1);
518 }
519}