agent_engine/tools/
secret_prompt.rs

1use std::sync::{Arc, Mutex};
2
3pub use crate::session::{ConfirmScope, PromptKind};
4
5/// UI-only secret prompt plumbing for interactive tools.
6///
7/// Secrets sent through this channel are never part of tool parameters, tool
8/// results, chat messages, or API messages. The TUI owns the input UI and sends
9/// only the final secret bytes back to the waiting tool.
10#[derive(Clone)]
11pub struct SecretPromptHandle {
12    tx: tokio::sync::mpsc::UnboundedSender<SecretPromptRequest>,
13}
14
15impl SecretPromptHandle {
16    pub fn new(tx: tokio::sync::mpsc::UnboundedSender<SecretPromptRequest>) -> Self {
17        Self { tx }
18    }
19
20    pub async fn prompt(&self, title: String, prompt: String) -> Option<String> {
21        self.prompt_scoped(title, prompt, None).await
22    }
23
24    /// Like [`Self::prompt`], carrying the extension tool-call confirm scope
25    /// so the dialog can name the source extension(s) and the reach of
26    /// "Allow all" before the user answers.
27    pub async fn prompt_scoped(
28        &self,
29        title: String,
30        prompt: String,
31        confirm_scope: Option<ConfirmScope>,
32    ) -> Option<String> {
33        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
34        let request = SecretPromptRequest {
35            kind: PromptKind::from_title(&title),
36            title,
37            prompt,
38            confirm_scope,
39            response_tx,
40        };
41        self.tx.send(request).ok()?;
42        response_rx.await.ok().flatten()
43    }
44}
45
46pub struct SecretPromptRequest {
47    /// `Confirm` renders as a y/n dialog (body visible), `Secret` as a masked
48    /// field. Derived from the title by `SecretPromptHandle::prompt`; carried
49    /// verbatim from the wire `PromptRequest` by the daemon PromptBridge.
50    pub kind: PromptKind,
51    pub title: String,
52    pub prompt: String,
53    /// Extension tool-call confirms only (see [`ConfirmScope`]).
54    pub confirm_scope: Option<ConfirmScope>,
55    pub response_tx: tokio::sync::oneshot::Sender<Option<String>>,
56}
57
58/// Which button of a `Confirm` dialog is focused / was chosen.
59#[derive(Clone, Copy, Debug, PartialEq, Eq)]
60pub enum ConfirmChoice {
61    /// Answer "y": allow this one call.
62    Allow,
63    /// Answer [`CONFIRM_ANSWER_ALLOW_ALL`]: allow this call and latch the
64    /// source extension(s) so their later confirms are auto-approved. Only
65    /// offered for `Confirm tool call` prompts that carry a [`ConfirmScope`]
66    /// (not `activate_tools`), and only reachable deliberately: focus the
67    /// button (Tab / arrows) and press Enter. No letter key selects it.
68    AllowAll,
69    /// Answer `None`: deny (the default focus, so a bare Enter fails closed).
70    Deny,
71}
72
73/// The answer string [`ConfirmChoice::AllowAll`] sends. The runtime's
74/// `resolve_before_tool_call_result` matches on it (re-exported there as
75/// `runtime::CONFIRM_ANSWER_ALLOW_ALL`); defined here so the tools module
76/// does not depend on the runtime module.
77pub const CONFIRM_ANSWER_ALLOW_ALL: &str = "always";
78
79/// A `Confirm` dialog ignores every key for this long after it becomes
80/// active, so keystrokes already in flight (typing a steer when the dialog
81/// pops up) can never answer it.
82pub const CONFIRM_ARM_DELAY: std::time::Duration = std::time::Duration::from_millis(400);
83
84pub struct PendingSecretPrompt {
85    pub kind: PromptKind,
86    pub title: String,
87    pub prompt: String,
88    pub confirm_scope: Option<ConfirmScope>,
89    pub buffer: String,
90    /// `Confirm` only: the focused button. Starts on `Deny`.
91    pub confirm_focus: ConfirmChoice,
92    /// Start of the current input guard (see [`CONFIRM_ARM_DELAY`]): set
93    /// when the prompt becomes active, and restarted by every key ignored
94    /// while unarmed and when "Allow all" asks for its confirmation.
95    pub activated_at: std::time::Instant,
96    /// `Confirm` only: Enter was pressed on "Allow all" once; a second
97    /// Enter (after a fresh guard) arms it, Esc or moving focus cancels.
98    pub allow_all_pending: bool,
99    pub response_tx: tokio::sync::oneshot::Sender<Option<String>>,
100}
101
102impl PendingSecretPrompt {
103    /// `Confirm` only: whether keys are accepted yet. `Secret` prompts are
104    /// always armed (they are typed into, not answered by a single key).
105    pub fn input_armed(&self) -> bool {
106        self.kind != PromptKind::Confirm || self.activated_at.elapsed() >= CONFIRM_ARM_DELAY
107    }
108
109    /// Whether the "Allow all" button is offered: only the runtime's
110    /// per-tool-call gate (`Confirm tool call`) with a known source scope,
111    /// never the `activate_tools` gate (its runtime only accepts y/yes) and
112    /// never an unattributed confirm (the latch could not be scoped).
113    pub fn offers_allow_all(&self) -> bool {
114        self.kind == PromptKind::Confirm
115            && self.title == crate::session::types::CONFIRM_PROMPT_TITLE
116            && self
117                .confirm_scope
118                .as_ref()
119                .is_some_and(|s| !s.extensions.is_empty())
120    }
121
122    /// Buttons in display order.
123    pub fn confirm_choices(&self) -> &'static [ConfirmChoice] {
124        if self.offers_allow_all() {
125            &[ConfirmChoice::Allow, ConfirmChoice::AllowAll, ConfirmChoice::Deny]
126        } else {
127            &[ConfirmChoice::Allow, ConfirmChoice::Deny]
128        }
129    }
130}
131
132pub struct SecretPromptQueue {
133    active: Option<PendingSecretPrompt>,
134    pending: std::collections::VecDeque<SecretPromptRequest>,
135}
136
137impl Default for SecretPromptQueue {
138    fn default() -> Self {
139        Self::new()
140    }
141}
142
143impl SecretPromptQueue {
144    pub fn new() -> Self {
145        Self {
146            active: None,
147            pending: std::collections::VecDeque::new(),
148        }
149    }
150
151    pub fn poll_requests(
152        &mut self,
153        rx: &Arc<Mutex<tokio::sync::mpsc::UnboundedReceiver<SecretPromptRequest>>>,
154    ) {
155        if let Ok(mut rx) = rx.lock() {
156            while let Ok(req) = rx.try_recv() {
157                self.pending.push_back(req);
158            }
159        }
160        self.activate_next();
161    }
162
163    fn activate_next(&mut self) {
164        if self.active.is_some() {
165            return;
166        }
167        if let Some(req) = self.pending.pop_front() {
168            self.active = Some(PendingSecretPrompt {
169                kind: req.kind,
170                title: req.title,
171                prompt: req.prompt,
172                confirm_scope: req.confirm_scope,
173                buffer: String::new(),
174                confirm_focus: ConfirmChoice::Deny,
175                activated_at: std::time::Instant::now(),
176                allow_all_pending: false,
177                response_tx: req.response_tx,
178            });
179        }
180    }
181
182    pub fn is_active(&self) -> bool {
183        self.active.is_some()
184    }
185
186    pub fn active(&self) -> Option<&PendingSecretPrompt> {
187        self.active.as_ref()
188    }
189
190    pub fn push_char(&mut self, ch: char) {
191        if let Some(active) = self.active.as_mut() {
192            active.buffer.push(ch);
193        }
194    }
195
196    pub fn backspace(&mut self) {
197        if let Some(active) = self.active.as_mut() {
198            active.buffer.pop();
199        }
200    }
201
202    /// `Confirm` dialog: move focus to the next / previous button. Focus
203    /// does NOT wrap (from the default Deny, "forward" stays put). Moving
204    /// cancels a pending "Allow all" confirmation. No-op for `Secret`.
205    pub fn move_confirm_focus(&mut self, forward: bool) {
206        if let Some(active) = self.active.as_mut() {
207            if active.kind != PromptKind::Confirm {
208                return;
209            }
210            let choices = active.confirm_choices();
211            let i = choices
212                .iter()
213                .position(|c| *c == active.confirm_focus)
214                .unwrap_or(choices.len() - 1);
215            let j = if forward {
216                (i + 1).min(choices.len() - 1)
217            } else {
218                i.saturating_sub(1)
219            };
220            active.confirm_focus = choices[j];
221            active.allow_all_pending = false;
222        }
223    }
224
225    /// Restart the input guard: the dialog ignores keys for another
226    /// [`CONFIRM_ARM_DELAY`] (called for every key ignored while unarmed, so
227    /// typing that runs past the window keeps the dialog dead until a pause).
228    pub fn restart_input_guard(&mut self) {
229        if let Some(a) = self.active.as_mut() {
230            a.activated_at = std::time::Instant::now();
231        }
232    }
233
234    /// `Confirm` dialog: whether "Allow all" is waiting for its second Enter.
235    pub fn allow_all_pending(&self) -> bool {
236        self.active.as_ref().is_some_and(|a| a.allow_all_pending)
237    }
238
239    /// Esc while "Allow all" awaits confirmation: cancel just that step.
240    /// Returns whether there was one to cancel.
241    pub fn cancel_allow_all_pending(&mut self) -> bool {
242        match self.active.as_mut() {
243            Some(a) if a.allow_all_pending => {
244                a.allow_all_pending = false;
245                true
246            }
247            _ => false,
248        }
249    }
250
251    /// Input ownership moved away: remove the "Allow all" button from the
252    /// prompt on screen AND from every prompt already queued behind it
253    /// (only the owner may arm it).
254    pub fn drop_allow_all(&mut self) {
255        for queued in self.pending.iter_mut() {
256            queued.confirm_scope = None;
257        }
258        if let Some(a) = self.active.as_mut() {
259            a.confirm_scope = None;
260            a.allow_all_pending = false;
261            if a.confirm_focus == ConfirmChoice::AllowAll {
262                a.confirm_focus = ConfirmChoice::Deny;
263            }
264        }
265    }
266
267    /// `Confirm` dialog: the focused button (`Deny` when no confirm is active).
268    pub fn confirm_focus(&self) -> ConfirmChoice {
269        self.active
270            .as_ref()
271            .filter(|a| a.kind == PromptKind::Confirm)
272            .map_or(ConfirmChoice::Deny, |a| a.confirm_focus)
273    }
274
275    /// Whether the active prompt accepts keys yet (see [`CONFIRM_ARM_DELAY`]).
276    pub fn input_armed(&self) -> bool {
277        self.active.as_ref().is_none_or(|a| a.input_armed())
278    }
279
280    /// Test/harness seam: treat the active prompt as opened long enough ago
281    /// that the arming delay has passed. Not compiled into production.
282    #[cfg(any(test, feature = "testing"))]
283    #[doc(hidden)]
284    pub fn age_input_guard_for_tests(&mut self, by: std::time::Duration) {
285        if let Some(a) = self.active.as_mut() {
286            if let Some(t) = a.activated_at.checked_sub(by) {
287                a.activated_at = t;
288            }
289        }
290    }
291
292    #[cfg(any(test, feature = "testing"))]
293    #[doc(hidden)]
294    pub fn arm_now_for_tests(&mut self) {
295        if let Some(a) = self.active.as_mut() {
296            if let Some(t) = std::time::Instant::now().checked_sub(CONFIRM_ARM_DELAY * 2) {
297                a.activated_at = t;
298            }
299        }
300    }
301
302    /// `Confirm` dialog: whether the "Allow all (session)" button is offered.
303    pub fn confirm_offers_allow_all(&self) -> bool {
304        self.active.as_ref().is_some_and(|a| a.offers_allow_all())
305    }
306
307    /// `Confirm` dialog: answer with `choice`. `AllowAll` on a prompt that
308    /// does not offer it does nothing (never a silent single Allow).
309    /// Returns the choice actually sent, or `None` when nothing was sent.
310    pub fn confirm_answer(&mut self, choice: ConfirmChoice) -> Option<ConfirmChoice> {
311        let active = self.active.as_mut()?;
312        if active.kind != PromptKind::Confirm {
313            return None;
314        }
315        if choice == ConfirmChoice::AllowAll && !active.offers_allow_all() {
316            return None;
317        }
318        match choice {
319            ConfirmChoice::Deny => self.cancel(),
320            ConfirmChoice::Allow | ConfirmChoice::AllowAll => {
321                active.buffer.clear();
322                active.buffer.push_str(if choice == ConfirmChoice::Allow {
323                    "y"
324                } else {
325                    CONFIRM_ANSWER_ALLOW_ALL
326                });
327                self.submit();
328            }
329        }
330        Some(choice)
331    }
332
333    /// `Confirm` dialog: activate the focused button (Enter). Deny is the
334    /// default focus, so an un-navigated Enter denies (fail-closed).
335    /// "Allow all" takes two Enters: the first only asks for confirmation
336    /// and restarts the input guard; the second (once armed) answers.
337    pub fn confirm_activate_focused(&mut self) -> Option<ConfirmChoice> {
338        let choice = self.confirm_focus();
339        if choice == ConfirmChoice::AllowAll {
340            let active = self.active.as_mut()?;
341            if !active.offers_allow_all() {
342                return None;
343            }
344            if !active.allow_all_pending {
345                active.allow_all_pending = true;
346                active.activated_at = std::time::Instant::now();
347                return None;
348            }
349        }
350        self.confirm_answer(choice)
351    }
352
353    pub fn submit(&mut self) {
354        if let Some(mut active) = self.active.take() {
355            let secret = std::mem::take(&mut active.buffer);
356            let _ = active.response_tx.send(Some(secret));
357        }
358        self.activate_next();
359    }
360
361    pub fn cancel(&mut self) {
362        if let Some(mut active) = self.active.take() {
363            active.buffer.clear();
364            let _ = active.response_tx.send(None);
365        }
366        self.activate_next();
367    }
368
369    /// Drop the active prompt WITHOUT answering (the oneshot is dropped, not
370    /// sent): another client resolved it (`PromptResolved` for a prompt this
371    /// client did not answer). The next pending prompt activates.
372    pub fn dismiss(&mut self) {
373        if let Some(mut active) = self.active.take() {
374            active.buffer.clear();
375            drop(active.response_tx);
376        }
377        self.activate_next();
378    }
379}
380
381#[cfg(test)]
382mod tests {
383    use super::*;
384
385    #[test]
386    fn dismiss_drops_without_answering_and_activates_next() {
387        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
388        let rx = Arc::new(Mutex::new(rx));
389        let (tx1, mut rx1) = tokio::sync::oneshot::channel();
390        let (tx2, mut rx2) = tokio::sync::oneshot::channel();
391        tx.send(SecretPromptRequest { kind: PromptKind::Secret, title: "a".into(), prompt: "p".into(), confirm_scope: None, response_tx: tx1 }).unwrap();
392        tx.send(SecretPromptRequest { kind: PromptKind::Secret, title: "b".into(), prompt: "p".into(), confirm_scope: None, response_tx: tx2 }).unwrap();
393        let mut q = SecretPromptQueue::new();
394        q.poll_requests(&rx);
395        assert_eq!(q.active().unwrap().title, "a");
396        q.push_char('z');
397        q.dismiss();
398        // Dropped, never sent: the waiter sees a closed channel (== cancelled).
399        assert!(matches!(rx1.try_recv(), Err(tokio::sync::oneshot::error::TryRecvError::Closed)));
400        // The next prompt is active with a fresh buffer.
401        assert_eq!(q.active().unwrap().title, "b");
402        assert_eq!(q.active().unwrap().buffer, "");
403        q.push_char('x');
404        q.submit();
405        assert_eq!(rx2.try_recv().unwrap(), Some("x".to_string()));
406        assert!(!q.is_active());
407    }
408
409    fn confirm(title: &str, scope: Option<ConfirmScope>) -> SecretPromptQueue {
410        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
411        let rx = Arc::new(Mutex::new(rx));
412        let (response_tx, _rx) = tokio::sync::oneshot::channel();
413        tx.send(SecretPromptRequest {
414            kind: PromptKind::Confirm,
415            title: title.into(),
416            prompt: "p".into(),
417            confirm_scope: scope,
418            response_tx,
419        })
420        .unwrap();
421        let mut q = SecretPromptQueue::new();
422        q.poll_requests(&rx);
423        q
424    }
425
426    #[test]
427    fn allow_all_requires_an_attributed_tool_call_confirm() {
428        let scope = ConfirmScope { extensions: vec!["jev".into()], tool: Some("bash".into()), rewritten_by: vec![] };
429        let title = crate::session::CONFIRM_PROMPT_TITLE;
430        assert!(confirm(title, Some(scope.clone())).confirm_offers_allow_all());
431        assert!(
432            !confirm(title, None).confirm_offers_allow_all(),
433            "an unattributed confirm cannot be latched, so no Allow all"
434        );
435        assert!(!confirm(
436            title,
437            Some(ConfirmScope { extensions: vec![], tool: None, rewritten_by: vec![] })
438        )
439        .confirm_offers_allow_all());
440        assert!(!confirm(crate::session::CONFIRM_ACTIVATION_PROMPT_TITLE, Some(scope))
441            .confirm_offers_allow_all());
442        // Without a scope, AllowAll does nothing (no silent single Allow).
443        let mut q = confirm(title, None);
444        assert_eq!(q.confirm_answer(ConfirmChoice::AllowAll), None);
445        assert!(q.is_active(), "prompt still waiting for a real answer");
446    }
447
448    #[test]
449    fn focus_does_not_wrap_and_allow_all_needs_two_enters() {
450        let scope = ConfirmScope { extensions: vec!["jev".into()], tool: None, rewritten_by: vec![] };
451        let mut q = confirm(crate::session::CONFIRM_PROMPT_TITLE, Some(scope));
452        q.arm_now_for_tests();
453        q.move_confirm_focus(true);
454        assert_eq!(q.confirm_focus(), ConfirmChoice::Deny, "no wrap past the end");
455        q.move_confirm_focus(false);
456        assert_eq!(q.confirm_focus(), ConfirmChoice::AllowAll);
457        assert_eq!(q.confirm_activate_focused(), None, "first Enter only asks");
458        assert!(q.allow_all_pending());
459        assert!(!q.input_armed(), "the confirmation gets its own fresh guard");
460        assert!(q.cancel_allow_all_pending());
461        assert!(!q.allow_all_pending());
462        assert_eq!(q.confirm_activate_focused(), None);
463        q.arm_now_for_tests();
464        assert_eq!(q.confirm_activate_focused(), Some(ConfirmChoice::AllowAll));
465        assert!(!q.is_active());
466    }
467
468    #[test]
469    fn dropping_allow_all_strips_queued_prompts_too() {
470        let scope = || Some(ConfirmScope { extensions: vec!["jev".into()], tool: None, rewritten_by: vec![] });
471        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
472        let rx = Arc::new(Mutex::new(rx));
473        let mut keep = Vec::new();
474        for _ in 0..2 {
475            let (response_tx, r) = tokio::sync::oneshot::channel();
476            keep.push(r);
477            tx.send(SecretPromptRequest {
478                kind: PromptKind::Confirm,
479                title: crate::session::CONFIRM_PROMPT_TITLE.into(),
480                prompt: "p".into(),
481                confirm_scope: scope(),
482                response_tx,
483            })
484            .unwrap();
485        }
486        let mut q = SecretPromptQueue::new();
487        q.poll_requests(&rx);
488        assert!(q.confirm_offers_allow_all());
489        q.drop_allow_all();
490        assert!(!q.confirm_offers_allow_all());
491        q.confirm_answer(ConfirmChoice::Deny);
492        assert!(q.is_active(), "the queued prompt is now active");
493        assert!(!q.confirm_offers_allow_all(), "queued prompt lost its button too");
494    }
495
496    #[test]
497    fn dropping_allow_all_moves_focus_off_it() {
498        let scope = ConfirmScope { extensions: vec!["jev".into()], tool: None, rewritten_by: vec![] };
499        let mut q = confirm(crate::session::CONFIRM_PROMPT_TITLE, Some(scope));
500        q.move_confirm_focus(false);
501        assert_eq!(q.confirm_focus(), ConfirmChoice::AllowAll);
502        q.drop_allow_all();
503        assert!(!q.confirm_offers_allow_all());
504        assert_eq!(q.confirm_focus(), ConfirmChoice::Deny);
505        q.move_confirm_focus(false);
506        assert_eq!(q.confirm_focus(), ConfirmChoice::Allow, "two buttons now");
507    }
508
509    #[test]
510    fn confirm_ignores_input_until_armed_secret_never_waits() {
511        let scope = ConfirmScope { extensions: vec!["jev".into()], tool: None, rewritten_by: vec![] };
512        let mut q = confirm(crate::session::CONFIRM_PROMPT_TITLE, Some(scope));
513        assert!(!q.input_armed(), "a fresh confirm is not armed");
514        q.arm_now_for_tests();
515        assert!(q.input_armed());
516        assert!(SecretPromptQueue::new().input_armed(), "no prompt → nothing to guard");
517    }
518}
agent_engine/tools/secret_prompt.rs

agent_engine/tools/
secret_prompt.rs
