1use super::{AgentTool, AgentToolResult, ToolContext, ToolError};
22use crate::tools::typed::TypedTool;
23use async_trait::async_trait;
24use schemars::JsonSchema;
25use serde::{Deserialize, Serialize};
26use std::sync::Arc;
27use std::sync::atomic::{AtomicBool, Ordering};
28use std::time::Duration;
29use tokio::sync::oneshot;
30
31#[derive(Clone)]
34pub struct AskBridge {
35 inner: Arc<parking_lot::Mutex<Option<PendingAsk>>>,
36 ui_attached: Arc<AtomicBool>,
40 session_id: Arc<parking_lot::Mutex<Option<String>>>,
47 timeout: Option<Duration>,
50}
51
52impl AskBridge {
53 pub fn new() -> Self {
55 Self {
56 inner: Arc::new(parking_lot::Mutex::new(None)),
57 ui_attached: Arc::new(AtomicBool::new(false)),
58 session_id: Arc::new(parking_lot::Mutex::new(None)),
59 timeout: None,
60 }
61 }
62
63 pub fn with_timeout(timeout: Option<Duration>) -> Self {
65 Self {
66 timeout,
67 ..Self::new()
68 }
69 }
70
71 pub fn attach_with_session(&self, session_id: impl Into<String>) {
78 let id = session_id.into();
79 debug_assert!(
80 !id.is_empty(),
81 "AskBridge::attach_with_session called with empty session_id"
82 );
83 *self.session_id.lock() = Some(id);
84 self.ui_attached.store(true, Ordering::SeqCst);
85 }
86
87 pub fn is_ui_attached(&self) -> bool {
89 self.ui_attached.load(Ordering::SeqCst)
90 }
91
92 #[cfg(any(test, debug_assertions))]
96 pub fn attach(&self) {
97 self.ui_attached.store(true, Ordering::SeqCst);
98 }
99
100 pub fn session_id(&self) -> Option<String> {
102 self.session_id.lock().clone()
103 }
104 pub fn timeout(&self) -> Option<Duration> {
106 self.timeout
107 }
108
109 pub fn set(&self, pending: PendingAsk) -> bool {
113 let mut lock = self.inner.lock();
114 if lock.is_some() {
115 return false;
116 }
117 *lock = Some(pending);
118 true
119 }
120
121 pub fn try_take(&self) -> Option<PendingAsk> {
124 self.inner.lock().take()
125 }
126
127 pub fn has_pending(&self) -> bool {
129 self.inner.lock().is_some()
130 }
131}
132
133impl Default for AskBridge {
134 fn default() -> Self {
135 Self::new()
136 }
137}
138
139pub struct PendingAsk {
143 pub questions: Vec<Question>,
145 pub responder: oneshot::Sender<AskResponse>,
148 pub timeout: Option<Duration>,
150 pub session_id: Option<String>,
154}
155
156#[derive(Debug, Clone, Serialize, Deserialize)]
158pub struct Question {
159 pub id: String,
161 #[serde(default)]
164 pub label: String,
165 pub prompt: String,
167 #[serde(default)]
169 pub options: Vec<QuestionOption>,
170 #[serde(default = "default_true")]
174 pub allow_other: bool,
175 #[serde(default)]
177 pub multi_select: bool,
178 #[serde(default)]
182 pub recommended: Option<usize>,
183}
184
185fn default_true() -> bool {
186 true
187}
188
189#[derive(Debug, Clone, Serialize, Deserialize)]
191pub struct QuestionOption {
192 pub value: String,
194 pub label: String,
196 pub description: Option<String>,
198}
199
200#[derive(Debug, Clone, Serialize, Deserialize)]
202pub struct AskResponse {
203 pub answers: Vec<Answer>,
205 pub cancelled: bool,
207 #[serde(default)]
209 pub timed_out: bool,
210}
211
212#[derive(Debug, Clone, Serialize, Deserialize)]
214pub struct Answer {
215 pub id: String,
217 pub value: String,
219 pub label: String,
221 pub was_custom: bool,
223 pub index: Option<usize>,
225}
226
227#[derive(Deserialize, Serialize, JsonSchema)]
231pub struct AskArgs {
232 questions: Vec<AskArgsQuestion>,
233}
234#[derive(Deserialize, Serialize, JsonSchema)]
236pub struct AskArgsQuestion {
237 id: String,
238 label: Option<String>,
239 prompt: String,
240 #[serde(default)]
241 options: Vec<AskArgsOption>,
242 #[serde(rename = "allowOther")]
243 #[serde(default = "default_true")]
244 allow_other: bool,
245 #[serde(rename = "multiSelect")]
246 #[serde(default)]
247 multi_select: bool,
248 recommended: Option<usize>,
249}
250
251#[derive(Deserialize, Serialize, JsonSchema)]
253pub struct AskArgsOption {
254 value: String,
255 label: String,
256 description: Option<String>,
257}
258
259pub struct AskTool {
263 bridge: Arc<AskBridge>,
264}
265
266impl AskTool {
267 pub fn new(bridge: Arc<AskBridge>) -> Self {
269 Self { bridge }
270 }
271}
272
273impl Clone for AskTool {
276 fn clone(&self) -> Self {
277 Self {
278 bridge: self.bridge.clone(),
279 }
280 }
281}
282
283#[async_trait]
284impl AgentTool for AskTool {
285 fn name(&self) -> &str {
286 "ask"
287 }
288
289 fn label(&self) -> &str {
290 "Ask"
291 }
292
293 fn description(&self) -> &str {
294 "Ask the user a clarifying question when choices have materially \
295 different tradeoffs the user must decide. Default to action — pick \
296 the conservative/standard option and proceed when a reasonable \
297 default exists; only ask when the user must weigh the tradeoff. Do \
298 NOT include an 'Other' option — the UI appends 'Other (type your \
299 own)' automatically. Use 'recommended' (0-indexed) to mark the \
300 default; a '(Recommended)' suffix is added automatically. Set \
301 'multiSelect' true to allow multiple selections. Provide 2-5 \
302 concise options with short labels; put explanatory tradeoffs in \
303 'description'. Batch related questions in one call via 'questions'."
304 }
305
306 fn parameters_schema(&self) -> serde_json::Value {
307 serde_json::json!({
308 "type": "object",
309 "properties": {
310 "questions": {
311 "type": "array",
312 "description": "Questions to ask the user",
313 "items": {
314 "type": "object",
315 "properties": {
316 "id": {
317 "type": "string",
318 "description": "Unique identifier for this question"
319 },
320 "label": {
321 "type": "string",
322 "description": "Short contextual label (defaults to the id)"
323 },
324 "prompt": {
325 "type": "string",
326 "description": "The full question text to display"
327 },
328 "options": {
329 "type": "array",
330 "description": "Available options (2-5). Do NOT include 'Other' — the UI adds it automatically.",
331 "default": [],
332 "items": {
333 "type": "object",
334 "properties": {
335 "value": {
336 "type": "string",
337 "description": "The value returned when selected"
338 },
339 "label": {
340 "type": "string",
341 "description": "Short display label for the option"
342 },
343 "description": {
344 "type": "string",
345 "description": "Optional explanatory tradeoff shown below the label"
346 }
347 },
348 "required": ["value", "label"]
349 }
350 },
351 "allowOther": {
352 "type": "boolean",
353 "description": "Show 'Other (type your own)' (default: true)",
354 "default": true
355 },
356 "multiSelect": {
357 "type": "boolean",
358 "description": "Allow multiple selections (default: false)",
359 "default": false
360 },
361 "recommended": {
362 "type": "number",
363 "description": "Recommended option index (0-based). Marks the default and is used for timeout auto-selection.",
364 "minimum": 0
365 }
366 },
367 "required": ["id", "prompt"]
368 }
369 },
370 },
371 "required": ["questions"]
372 })
373 }
374
375 async fn execute(
376 &self,
377 _tool_call_id: &str,
378 params: serde_json::Value,
379 signal: Option<oneshot::Receiver<()>>,
380 _ctx: &ToolContext,
381 ) -> Result<AgentToolResult, ToolError> {
382 let args: AskArgs =
383 serde_json::from_value(params).map_err(|e| format!("invalid params: {e}"))?;
384 self.execute_typed(_tool_call_id, args, signal, _ctx).await
385 }
386}
387
388#[async_trait]
389impl TypedTool for AskTool {
390 type Args = AskArgs;
391
392 async fn execute_typed(
393 &self,
394 _tool_call_id: &str,
395 args: Self::Args,
396 signal: Option<oneshot::Receiver<()>>,
397 _ctx: &ToolContext,
398 ) -> Result<AgentToolResult, ToolError> {
399 if !self.bridge.is_ui_attached() {
401 return Ok(AgentToolResult::error(
402 "Ask requires interactive TUI mode. \
403 Not available in --print or RPC mode.",
404 ));
405 }
406
407 let session_id = self.bridge.session_id();
409 debug_assert!(
410 session_id.as_deref().is_some_and(|s| !s.is_empty()),
411 "AskBridge was attached without a non-empty session_id; refusing to run"
412 );
413
414 let params = serde_json::to_value(&args).map_err(|e| format!("serialize: {e}"))?;
416 let questions = parse_questions(¶ms)?;
417 let timeout = self.bridge.timeout();
418
419 let (tx, rx) = oneshot::channel();
421
422 if !self.bridge.set(PendingAsk {
424 questions,
425 responder: tx,
426 timeout,
427 session_id,
428 }) {
429 return Ok(AgentToolResult::error("Another ask is already pending"));
430 }
431
432 select_with_abort(rx, signal, &self.bridge).await
434 }
435}
436
437async fn select_with_abort(
439 rx: oneshot::Receiver<AskResponse>,
440 signal: Option<oneshot::Receiver<()>>,
441 bridge: &AskBridge,
442) -> Result<AgentToolResult, ToolError> {
443 let abort = async {
445 if let Some(sig) = signal {
446 let _ = sig.await;
447 } else {
448 std::future::pending::<()>().await;
449 }
450 };
451
452 tokio::select! {
453 response = rx => {
454 match response {
455 Ok(resp) => {
456 if resp.cancelled {
457 Ok(AgentToolResult::success("User cancelled the question"))
458 } else {
459 Ok(AgentToolResult::success(format_answers(
460 &resp.answers,
461 resp.timed_out,
462 )))
463 }
464 }
465 Err(_) => {
466 Ok(AgentToolResult::success("Question dismissed"))
468 }
469 }
470 }
471 () = abort => {
472 bridge.try_take();
474 Ok(AgentToolResult::success("Question cancelled by user interrupt"))
475 }
476 }
477}
478
479fn parse_questions(params: &serde_json::Value) -> Result<Vec<Question>, ToolError> {
481 let questions = params
482 .get("questions")
483 .and_then(|v| v.as_array())
484 .cloned()
485 .ok_or_else(|| "Missing or invalid 'questions' field".to_string())?;
486
487 let questions: Vec<Question> = questions
488 .into_iter()
489 .map(|v| serde_json::from_value(v).map_err(|e| e.to_string()))
490 .collect::<Result<Vec<_>, _>>()
491 .map_err(|e| format!("Invalid question: {}", e))?;
492
493 if questions.is_empty() {
494 return Err("At least one question is required".to_string());
495 }
496
497 let questions: Vec<Question> = questions
499 .into_iter()
500 .map(|mut q| {
501 if q.label.is_empty() {
502 q.label = q.id.clone();
503 }
504 q
505 })
506 .collect();
507
508 let mut ids = std::collections::HashSet::new();
510 for q in &questions {
511 if !ids.insert(&q.id) {
512 return Err(format!("Duplicate question id: {}", q.id));
513 }
514 }
515
516 Ok(questions)
517}
518
519pub fn format_answers(answers: &[Answer], timed_out: bool) -> String {
530 let suffix = if timed_out {
531 " (auto-selected after timeout)"
532 } else {
533 ""
534 };
535 answers
536 .iter()
537 .map(|a| {
538 let base = if a.was_custom {
539 format!("{}: \"{}\"", a.id, a.label)
540 } else if a.value.contains(',') {
541 let labels: Vec<&str> = a.label.split(", ").collect();
543 format!("{}: [{}]", a.id, labels.join(", "))
544 } else {
545 format!("{}: {}", a.id, a.label)
546 };
547 format!("{base}{suffix}")
548 })
549 .collect::<Vec<_>>()
550 .join("\n")
551}
552
553#[cfg(test)]
554mod tests {
555 use super::*;
556
557 #[test]
558 fn test_parse_questions_valid() {
559 let json = serde_json::json!({
560 "questions": [
561 {
562 "id": "lang",
563 "prompt": "Pick a language",
564 "options": [
565 { "value": "rust", "label": "Rust" },
566 { "value": "ts", "label": "TypeScript" }
567 ]
568 }
569 ]
570 });
571 let questions = parse_questions(&json).unwrap();
572 assert_eq!(questions.len(), 1);
573 assert_eq!(questions[0].id, "lang");
574 assert_eq!(questions[0].label, "lang"); assert_eq!(questions[0].options.len(), 2);
576 assert!(questions[0].allow_other); assert!(!questions[0].multi_select); }
579
580 #[test]
581 fn test_parse_questions_with_label() {
582 let json = serde_json::json!({
583 "questions": [
584 {
585 "id": "lang",
586 "label": "Language",
587 "prompt": "Pick a language"
588 }
589 ]
590 });
591 let questions = parse_questions(&json).unwrap();
592 assert_eq!(questions[0].label, "Language");
593 }
594
595 #[test]
596 fn test_parse_questions_empty_options() {
597 let json = serde_json::json!({
599 "questions": [
600 {
601 "id": "name",
602 "prompt": "What's your project name?",
603 "allowOther": true
604 }
605 ]
606 });
607 let questions = parse_questions(&json).unwrap();
608 assert_eq!(questions[0].options.len(), 0);
609 assert!(questions[0].allow_other);
610 }
611
612 #[test]
613 fn test_parse_questions_missing_questions() {
614 let json = serde_json::json!({});
615 let err = parse_questions(&json).unwrap_err();
616 assert!(err.contains("questions"));
617 }
618
619 #[test]
620 fn test_parse_questions_empty_array() {
621 let json = serde_json::json!({ "questions": [] });
622 let err = parse_questions(&json).unwrap_err();
623 assert!(err.contains("one question"));
624 }
625
626 #[test]
627 fn test_parse_questions_duplicate_ids() {
628 let json = serde_json::json!({
629 "questions": [
630 { "id": "a", "prompt": "Q1" },
631 { "id": "a", "prompt": "Q2" }
632 ]
633 });
634 let err = parse_questions(&json).unwrap_err();
635 assert!(err.contains("Duplicate"));
636 }
637
638 #[test]
639 fn test_format_answers_single() {
640 let answers = vec![Answer {
641 id: "lang".into(),
642 value: "rust".into(),
643 label: "Rust".into(),
644 was_custom: false,
645 index: Some(1),
646 }];
647 let text = format_answers(&answers, false);
648 assert_eq!(text, "lang: Rust");
649 }
650
651 #[test]
652 fn test_format_answers_custom() {
653 let answers = vec![Answer {
654 id: "name".into(),
655 value: "myproj".into(),
656 label: "myproj".into(),
657 was_custom: true,
658 index: None,
659 }];
660 let text = format_answers(&answers, false);
661 assert_eq!(text, "name: \"myproj\"");
662 }
663
664 #[test]
665 fn test_format_answers_multi() {
666 let answers = vec![Answer {
667 id: "lang".into(),
668 value: "rust, go".into(), label: "Rust, Go".into(),
670 was_custom: false,
671 index: None,
672 }];
673 let text = format_answers(&answers, false);
674 assert_eq!(text, "lang: [Rust, Go]");
675 }
676
677 #[test]
678 fn test_format_answers_timed_out() {
679 let answers = vec![Answer {
680 id: "auth".into(),
681 value: "oauth".into(),
682 label: "OAuth2".into(),
683 was_custom: false,
684 index: Some(2),
685 }];
686 let text = format_answers(&answers, true);
687 assert_eq!(text, "auth: OAuth2 (auto-selected after timeout)");
688 }
689
690 #[test]
691 fn test_bridge_set_take() {
692 let bridge = AskBridge::new();
693 assert!(!bridge.has_pending());
694
695 let (tx, _rx) = oneshot::channel();
696 let pending = PendingAsk {
697 questions: vec![],
698 responder: tx,
699 timeout: None,
700 session_id: None,
701 };
702 assert!(bridge.set(pending));
703 assert!(bridge.has_pending());
704
705 let taken = bridge.try_take();
706 assert!(taken.is_some());
707 assert!(!bridge.has_pending());
708
709 assert!(bridge.try_take().is_none());
711 }
712
713 #[test]
714 fn test_bridge_set_idempotent() {
715 let bridge = AskBridge::new();
716 let (tx1, _rx1) = oneshot::channel();
717 let (tx2, _rx2) = oneshot::channel();
718
719 bridge.set(PendingAsk {
720 questions: vec![],
721 responder: tx1,
722 timeout: None,
723 session_id: None,
724 });
725 assert!(!bridge.set(PendingAsk {
726 questions: vec![],
727 responder: tx2,
728 timeout: None,
729 session_id: None,
730 }));
731 }
732
733 #[test]
734 fn test_ui_attached_flag() {
735 let bridge = AskBridge::new();
736 assert!(!bridge.is_ui_attached());
737 bridge.attach();
738 assert!(bridge.is_ui_attached());
739 }
740
741 #[test]
742 fn test_bridge_with_timeout() {
743 let bridge = AskBridge::with_timeout(Some(Duration::from_secs(30)));
744 assert_eq!(bridge.timeout(), Some(Duration::from_secs(30)));
745 assert!(!bridge.is_ui_attached()); let no_timeout = AskBridge::new();
748 assert_eq!(no_timeout.timeout(), None);
749 }
750
751 #[test]
752 fn test_question_deserializes_without_recommended() {
753 let json = serde_json::json!({
755 "id": "test",
756 "prompt": "Test question?",
757 "options": [{"value": "a", "label": "A"}]
758 });
759 let q: Question = serde_json::from_value(json).unwrap();
760 assert_eq!(q.recommended, None);
761 }
762
763 #[test]
764 fn test_question_deserializes_with_recommended() {
765 let json = serde_json::json!({
766 "id": "test",
767 "prompt": "Test question?",
768 "options": [{"value": "a", "label": "A"}, {"value": "b", "label": "B"}],
769 "recommended": 1
770 });
771 let q: Question = serde_json::from_value(json).unwrap();
772 assert_eq!(q.recommended, Some(1));
773 }
774
775 #[test]
776 fn test_tool_name_is_ask() {
777 let bridge = Arc::new(AskBridge::new());
778 let tool = AskTool::new(bridge);
779 assert_eq!(tool.name(), "ask");
780 assert_eq!(tool.label(), "Ask");
781 }
782
783 #[test]
784 fn test_attach_with_session_stores_id() {
785 let bridge = AskBridge::new();
786 assert!(!bridge.is_ui_attached());
787 assert_eq!(bridge.session_id(), None);
788 bridge.attach_with_session("tui");
789 assert!(bridge.is_ui_attached());
790 assert_eq!(bridge.session_id().as_deref(), Some("tui"));
791 }
792
793 #[test]
794 fn test_format_answers_multi_with_comma_label() {
795 let answers = vec![Answer {
799 id: "tags".into(),
800 value: "a,b".into(),
801 label: "A, B".into(),
802 was_custom: false,
803 index: None,
804 }];
805 let text = format_answers(&answers, false);
806 assert_eq!(text, "tags: [A, B]");
807 }
808
809 #[test]
810 fn test_format_answers_cancelled_marker() {
811 let answers = vec![Answer {
812 id: "q1".into(),
813 value: String::new(),
814 label: String::new(),
815 was_custom: false,
816 index: None,
817 }];
818 let text = format_answers(&answers, false);
822 assert_eq!(text, "q1: ");
823 }
824}