Skip to main content

agent_engine/tools/
secret_prompt.rs

1use std::sync::{Arc, Mutex};
2
3pub use crate::session::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        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
22        let request = SecretPromptRequest {
23            kind: PromptKind::from_title(&title),
24            title,
25            prompt,
26            response_tx,
27        };
28        self.tx.send(request).ok()?;
29        response_rx.await.ok().flatten()
30    }
31}
32
33pub struct SecretPromptRequest {
34    /// `Confirm` renders as a y/n dialog (body visible), `Secret` as a masked
35    /// field. Derived from the title by `SecretPromptHandle::prompt`; carried
36    /// verbatim from the wire `PromptRequest` by the daemon PromptBridge.
37    pub kind: PromptKind,
38    pub title: String,
39    pub prompt: String,
40    pub response_tx: tokio::sync::oneshot::Sender<Option<String>>,
41}
42
43pub struct PendingSecretPrompt {
44    pub kind: PromptKind,
45    pub title: String,
46    pub prompt: String,
47    pub buffer: String,
48    pub response_tx: tokio::sync::oneshot::Sender<Option<String>>,
49}
50
51pub struct SecretPromptQueue {
52    active: Option<PendingSecretPrompt>,
53    pending: std::collections::VecDeque<SecretPromptRequest>,
54}
55
56impl Default for SecretPromptQueue {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62impl SecretPromptQueue {
63    pub fn new() -> Self {
64        Self {
65            active: None,
66            pending: std::collections::VecDeque::new(),
67        }
68    }
69
70    pub fn poll_requests(
71        &mut self,
72        rx: &Arc<Mutex<tokio::sync::mpsc::UnboundedReceiver<SecretPromptRequest>>>,
73    ) {
74        if let Ok(mut rx) = rx.lock() {
75            while let Ok(req) = rx.try_recv() {
76                self.pending.push_back(req);
77            }
78        }
79        self.activate_next();
80    }
81
82    fn activate_next(&mut self) {
83        if self.active.is_some() {
84            return;
85        }
86        if let Some(req) = self.pending.pop_front() {
87            self.active = Some(PendingSecretPrompt {
88                kind: req.kind,
89                title: req.title,
90                prompt: req.prompt,
91                buffer: String::new(),
92                response_tx: req.response_tx,
93            });
94        }
95    }
96
97    pub fn is_active(&self) -> bool {
98        self.active.is_some()
99    }
100
101    pub fn active(&self) -> Option<&PendingSecretPrompt> {
102        self.active.as_ref()
103    }
104
105    pub fn push_char(&mut self, ch: char) {
106        if let Some(active) = self.active.as_mut() {
107            active.buffer.push(ch);
108        }
109    }
110
111    pub fn backspace(&mut self) {
112        if let Some(active) = self.active.as_mut() {
113            active.buffer.pop();
114        }
115    }
116
117    pub fn submit(&mut self) {
118        if let Some(mut active) = self.active.take() {
119            let secret = std::mem::take(&mut active.buffer);
120            let _ = active.response_tx.send(Some(secret));
121        }
122        self.activate_next();
123    }
124
125    pub fn cancel(&mut self) {
126        if let Some(mut active) = self.active.take() {
127            active.buffer.clear();
128            let _ = active.response_tx.send(None);
129        }
130        self.activate_next();
131    }
132
133    /// Drop the active prompt WITHOUT answering (the oneshot is dropped, not
134    /// sent): another client resolved it (`PromptResolved` for a prompt this
135    /// client did not answer). The next pending prompt activates.
136    pub fn dismiss(&mut self) {
137        if let Some(mut active) = self.active.take() {
138            active.buffer.clear();
139            drop(active.response_tx);
140        }
141        self.activate_next();
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    #[test]
150    fn dismiss_drops_without_answering_and_activates_next() {
151        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
152        let rx = Arc::new(Mutex::new(rx));
153        let (tx1, mut rx1) = tokio::sync::oneshot::channel();
154        let (tx2, mut rx2) = tokio::sync::oneshot::channel();
155        tx.send(SecretPromptRequest { kind: PromptKind::Secret, title: "a".into(), prompt: "p".into(), response_tx: tx1 }).unwrap();
156        tx.send(SecretPromptRequest { kind: PromptKind::Secret, title: "b".into(), prompt: "p".into(), response_tx: tx2 }).unwrap();
157        let mut q = SecretPromptQueue::new();
158        q.poll_requests(&rx);
159        assert_eq!(q.active().unwrap().title, "a");
160        q.push_char('z');
161        q.dismiss();
162        // Dropped, never sent: the waiter sees a closed channel (== cancelled).
163        assert!(matches!(rx1.try_recv(), Err(tokio::sync::oneshot::error::TryRecvError::Closed)));
164        // The next prompt is active with a fresh buffer.
165        assert_eq!(q.active().unwrap().title, "b");
166        assert_eq!(q.active().unwrap().buffer, "");
167        q.push_char('x');
168        q.submit();
169        assert_eq!(rx2.try_recv().unwrap(), Some("x".to_string()));
170        assert!(!q.is_active());
171    }
172}