agent_engine/tools/
secret_prompt.rs1use std::sync::{Arc, Mutex};
2
3pub use crate::session::PromptKind;
4
5#[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 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 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 assert!(matches!(rx1.try_recv(), Err(tokio::sync::oneshot::error::TryRecvError::Closed)));
164 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}