Skip to main content

atman_runtime/tools/
form.rs

1use crate::error::RuntimeError;
2use crate::form::{CompositeForm, FormAnswer, FormKind, FormQuestion, PendingForm};
3use crate::tool::{ApprovalLevel, BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
4use crate::value::Value;
5
6pub struct FormAsk;
7
8impl Tool for FormAsk {
9    fn name(&self) -> &str {
10        "form.ask"
11    }
12
13    fn tier(&self) -> Tier {
14        Tier::Zero
15    }
16
17    fn approval_level(&self, _args: &ToolArgs, _ctx: &ToolCtx) -> ApprovalLevel {
18        ApprovalLevel::Auto
19    }
20
21    fn description(&self) -> Option<&str> {
22        Some(
23            "Ask the user a structured question through a form modal. Pass `kind`
24             plus fields required for that kind:
25             \
26             confirm       { kind:\"confirm\", prompt }
27             single_select { kind:\"single_select\", prompt, options[] }
28             multi_select  { kind:\"multi_select\", prompt, options[], min?, max? }
29             text          { kind:\"text\", prompt, placeholder?, multiline? }
30             \
31             For several independent answers, make one `form.ask` call with a `questions`
32             list. Each question has an `id`, `kind`, and the fields for that kind.
33             The UI keeps all answers as a draft and asks for one final Yes/No confirmation;
34             do not make multiple calls expecting the UI to merge them.
35             \
36             Returns a struct { kind, ... } where kind is one of \
37             confirmed | selected | multi_selected | text_entered | cancelled.",
38        )
39    }
40
41    fn input_schema(&self) -> serde_json::Value {
42        serde_json::json!({
43            "type": "object",
44            "properties": {
45                "kind": {"type": "string"},
46                "prompt": {"type": "string"},
47                "options": {"type": "array", "items": {"type": "string"}},
48                "min": {"type": "integer"},
49                "max": {"type": "integer"},
50                "placeholder": {"type": "string"},
51                "multiline": {"type": "boolean"},
52                "questions": {
53                    "type": "array",
54                    "minItems": 1,
55                    "items": {
56                        "type": "object",
57                        "properties": {
58                            "id": {"type": "string"},
59                            "kind": {"type": "string"},
60                            "prompt": {"type": "string"},
61                            "options": {"type": "array", "items": {"type": "string"}},
62                            "min": {"type": "integer"},
63                            "max": {"type": "integer"},
64                            "placeholder": {"type": "string"},
65                            "multiline": {"type": "boolean"}
66                        },
67                        "required": ["id", "kind", "prompt"]
68                    }
69                }
70            },
71            "oneOf": [
72                {"required": ["kind", "prompt"], "not": {"required": ["questions"]}},
73                {"required": ["questions"], "not": {"anyOf": [{"required": ["kind"]}, {"required": ["prompt"]}]}}
74            ]
75        })
76    }
77
78    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
79        Box::pin(async move {
80            let (form, kind, composite) = parse_form_request(&args)?;
81            // Daemon clients drive the modal over RPC via the prompt
82            // resolver; the in-process TUI subscribes to FormRegistry.
83            // Pick whichever the runtime host wired up, prefer the
84            // resolver so daemon overrides an accidental fallback.
85            if let Some(resolver) = ctx.prompt_resolver.clone() {
86                let id = crate::rendezvous::PromptId::now();
87                let payload = if composite {
88                    serde_json::to_value(&form).unwrap_or(serde_json::Value::Null)
89                } else {
90                    serde_json::to_value(&kind).unwrap_or(serde_json::Value::Null)
91                };
92                let timeout = std::time::Duration::from_secs(300);
93                let Some(answer_json) =
94                    crate::rendezvous::await_expirable_prompt_with_payload_cancel(
95                        &resolver,
96                        id,
97                        "form_ask",
98                        payload,
99                        timeout,
100                        &ctx.cancel,
101                    )
102                    .await?
103                else {
104                    return Ok(submission_to_value(
105                        &crate::form::FormSubmission::Rejected,
106                        composite,
107                    ));
108                };
109                let submission = serde_json::from_value::<crate::form::FormSubmission>(answer_json)
110                    .map_err(|error| {
111                        RuntimeError::ToolFailed(format!(
112                            "form.ask: invalid prompt submission: {error}"
113                        ))
114                    })?;
115                return Ok(submission_to_value(&submission, composite));
116            }
117            let forms = ctx.forms.as_ref().ok_or_else(|| {
118                RuntimeError::ToolFailed(
119                    "form.ask: no FormRegistry or PromptResolver attached".into(),
120                )
121            })?;
122            let run_id = ctx.flow_run_id.clone().ok_or_else(|| {
123                RuntimeError::ToolFailed("form.ask: no flow_run_id in ctx".into())
124            })?;
125            let form_id = uuid::Uuid::now_v7().to_string();
126            let pending = PendingForm {
127                form_id: form_id.clone(),
128                run_id,
129                tool_use_id: ctx.current_node_id.clone().unwrap_or_default(),
130                form,
131                kind,
132                emitted_at: chrono::Utc::now(),
133            };
134            let rx = forms.request(pending);
135            let submission = await_local_submission(
136                forms,
137                form_id,
138                rx,
139                std::time::Duration::from_secs(300),
140                &ctx.cancel,
141            )
142            .await;
143            Ok(submission_to_value(&submission, composite))
144        })
145    }
146}
147
148async fn await_local_submission(
149    forms: &crate::session::FormRegistry,
150    form_id: String,
151    mut rx: tokio::sync::oneshot::Receiver<crate::form::FormSubmission>,
152    timeout: std::time::Duration,
153    cancel: &tokio_util::sync::CancellationToken,
154) -> crate::form::FormSubmission {
155    tokio::select! {
156        result = tokio::time::timeout(timeout, &mut rx) => match result {
157            Ok(Ok(submission)) => submission,
158            Ok(Err(_)) => {
159                forms.cancel(&form_id);
160                crate::form::FormSubmission::Rejected
161            }
162            Err(_) => {
163                if forms.expire(&form_id) {
164                    crate::form::FormSubmission::Rejected
165                } else {
166                    tokio::select! {
167                        result = &mut rx => result.unwrap_or(crate::form::FormSubmission::Rejected),
168                        _ = cancel.cancelled() => {
169                            forms.cancel(&form_id);
170                            crate::form::FormSubmission::Rejected
171                        }
172                    }
173                }
174            }
175        },
176        _ = cancel.cancelled() => {
177            forms.cancel(&form_id);
178            crate::form::FormSubmission::Rejected
179        }
180    }
181}
182
183fn submission_to_value(submission: &crate::form::FormSubmission, composite: bool) -> Value {
184    match submission {
185        crate::form::FormSubmission::Submitted { answers } if composite => Value::Struct(vec![
186            ("kind".into(), Value::Str("submitted".into())),
187            (
188                "answers".into(),
189                Value::List(answers.iter().map(answer_to_value).collect()),
190            ),
191        ]),
192        crate::form::FormSubmission::Submitted { answers } => {
193            answers.first().map(answer_to_value).unwrap_or_else(|| {
194                Value::Struct(vec![("kind".into(), Value::Str("cancelled".into()))])
195            })
196        }
197        crate::form::FormSubmission::Rejected => {
198            Value::Struct(vec![("kind".into(), Value::Str("cancelled".into()))])
199        }
200    }
201}
202
203fn parse_form_request(args: &ToolArgs) -> Result<(CompositeForm, FormKind, bool), RuntimeError> {
204    match (args.named("questions"), args.named("kind")) {
205        (Some(Value::List(items)), None) => {
206            if items.is_empty() {
207                return Err(RuntimeError::ToolFailed(
208                    "form.ask: `questions` must be non-empty".into(),
209                ));
210            }
211            let mut questions = Vec::with_capacity(items.len());
212            for (index, item) in items.iter().enumerate() {
213                let Value::Struct(fields) = item else {
214                    return Err(RuntimeError::TypeMismatch {
215                        expected: "struct {id, kind, prompt, ...}".into(),
216                        actual: item.kind_name().into(),
217                    });
218                };
219                let get = |name: &str| {
220                    fields
221                        .iter()
222                        .find(|(key, _)| key == name)
223                        .map(|(_, value)| value)
224                };
225                let id = match get("id") {
226                    Some(Value::Str(value)) if !value.is_empty() => value.clone(),
227                    Some(value) => {
228                        return Err(RuntimeError::TypeMismatch {
229                            expected: "string".into(),
230                            actual: value.kind_name().into(),
231                        });
232                    }
233                    None => return Err(RuntimeError::MissingArg(format!("questions[{index}].id"))),
234                };
235                if questions
236                    .iter()
237                    .any(|question: &FormQuestion| question.id == id)
238                {
239                    return Err(RuntimeError::ToolFailed(format!(
240                        "form.ask: duplicate question id `{id}`"
241                    )));
242                }
243                let named = ToolArgs {
244                    positional: Vec::new(),
245                    named: fields.clone(),
246                };
247                let kind = parse_form_kind(&named)?;
248                questions.push(FormQuestion { id, kind });
249            }
250            let first = questions[0].kind.clone();
251            Ok((CompositeForm { questions }, first, true))
252        }
253        (Some(value), _) => Err(RuntimeError::TypeMismatch {
254            expected: "list<struct>".into(),
255            actual: value.kind_name().into(),
256        }),
257        (None, _) => {
258            let kind = parse_form_kind(args)?;
259            Ok((
260                CompositeForm {
261                    questions: vec![FormQuestion {
262                        id: "question".into(),
263                        kind: kind.clone(),
264                    }],
265                },
266                kind,
267                false,
268            ))
269        }
270    }
271}
272
273fn parse_form_kind(args: &ToolArgs) -> Result<FormKind, RuntimeError> {
274    let kind = named_str(args, "kind")?;
275    let prompt = named_str(args, "prompt")?;
276    match kind.as_str() {
277        "confirm" => Ok(FormKind::Confirm { prompt }),
278        "single_select" => {
279            let options = named_string_list(args, "options")?;
280            if options.is_empty() {
281                return Err(RuntimeError::ToolFailed(
282                    "form.ask(single_select): options must be non-empty".into(),
283                ));
284            }
285            Ok(FormKind::SingleSelect { prompt, options })
286        }
287        "multi_select" => {
288            let options = named_string_list(args, "options")?;
289            if options.is_empty() {
290                return Err(RuntimeError::ToolFailed(
291                    "form.ask(multi_select): options must be non-empty".into(),
292                ));
293            }
294            let min = named_usize(args, "min")?;
295            let max = named_usize(args, "max")?;
296            if let (Some(m), Some(mx)) = (min, max)
297                && mx < m
298            {
299                return Err(RuntimeError::ToolFailed(
300                    "form.ask(multi_select): max must be >= min".into(),
301                ));
302            }
303            Ok(FormKind::MultiSelect {
304                prompt,
305                options,
306                min,
307                max,
308            })
309        }
310        "text" => {
311            let placeholder = named_opt_str(args, "placeholder")?;
312            let multiline = matches!(args.named("multiline"), Some(Value::Bool(true)));
313            Ok(FormKind::Text {
314                prompt,
315                placeholder,
316                multiline,
317            })
318        }
319        other => Err(RuntimeError::ToolFailed(format!(
320            "form.ask: unknown kind `{other}` (expected confirm | single_select | multi_select | text)"
321        ))),
322    }
323}
324
325fn named_str(args: &ToolArgs, name: &str) -> Result<String, RuntimeError> {
326    match args.named(name) {
327        Some(Value::Str(s)) => Ok(s.clone()),
328        Some(v) => Err(RuntimeError::TypeMismatch {
329            expected: "string".into(),
330            actual: v.kind_name().into(),
331        }),
332        None => Err(RuntimeError::MissingArg(name.into())),
333    }
334}
335
336fn named_opt_str(args: &ToolArgs, name: &str) -> Result<Option<String>, RuntimeError> {
337    match args.named(name) {
338        Some(Value::Str(s)) => Ok(Some(s.clone())),
339        Some(Value::Unit) | None => Ok(None),
340        Some(v) => Err(RuntimeError::TypeMismatch {
341            expected: "string".into(),
342            actual: v.kind_name().into(),
343        }),
344    }
345}
346
347fn named_string_list(args: &ToolArgs, name: &str) -> Result<Vec<String>, RuntimeError> {
348    match args.named(name) {
349        Some(Value::List(items)) => items
350            .iter()
351            .map(|v| match v {
352                Value::Str(s) => Ok(s.clone()),
353                other => Err(RuntimeError::TypeMismatch {
354                    expected: "string".into(),
355                    actual: other.kind_name().into(),
356                }),
357            })
358            .collect(),
359        Some(v) => Err(RuntimeError::TypeMismatch {
360            expected: "list<string>".into(),
361            actual: v.kind_name().into(),
362        }),
363        None => Err(RuntimeError::MissingArg(name.into())),
364    }
365}
366
367fn named_usize(args: &ToolArgs, name: &str) -> Result<Option<usize>, RuntimeError> {
368    match args.named(name) {
369        Some(Value::Int(i)) if *i >= 0 => Ok(Some(*i as usize)),
370        Some(Value::Int(_)) => Err(RuntimeError::ToolFailed(format!(
371            "form.ask: `{name}` must be non-negative"
372        ))),
373        Some(Value::Unit) | None => Ok(None),
374        Some(v) => Err(RuntimeError::TypeMismatch {
375            expected: "int".into(),
376            actual: v.kind_name().into(),
377        }),
378    }
379}
380
381fn answer_to_value(answer: &FormAnswer) -> Value {
382    match answer {
383        FormAnswer::Confirmed { value } => Value::Struct(vec![
384            ("kind".into(), Value::Str("confirmed".into())),
385            ("value".into(), Value::Bool(*value)),
386        ]),
387        FormAnswer::Selected { index, label } => Value::Struct(vec![
388            ("kind".into(), Value::Str("selected".into())),
389            ("index".into(), Value::Int(*index as i64)),
390            ("label".into(), Value::Str(label.clone())),
391        ]),
392        FormAnswer::MultiSelected { indices, labels } => Value::Struct(vec![
393            ("kind".into(), Value::Str("multi_selected".into())),
394            (
395                "indices".into(),
396                Value::List(indices.iter().map(|i| Value::Int(*i as i64)).collect()),
397            ),
398            (
399                "labels".into(),
400                Value::List(labels.iter().map(|s| Value::Str(s.clone())).collect()),
401            ),
402        ]),
403        FormAnswer::TextEntered { text } => Value::Struct(vec![
404            ("kind".into(), Value::Str("text_entered".into())),
405            ("text".into(), Value::Str(text.clone())),
406        ]),
407        FormAnswer::Cancelled => {
408            Value::Struct(vec![("kind".into(), Value::Str("cancelled".into()))])
409        }
410    }
411}
412
413#[cfg(test)]
414mod tests {
415    use super::*;
416    use crate::form::FormKind;
417    use crate::tool::ToolArgs;
418
419    fn named(name: &str, v: Value) -> (String, Value) {
420        (name.into(), v)
421    }
422
423    #[tokio::test]
424    async fn local_form_timeout_keeps_pending_entry_answerable() {
425        let forms = crate::session::FormRegistry::new();
426        let _subscriber = forms.subscribe();
427        let form_id = "timed-out".to_string();
428        let form = crate::form::CompositeForm {
429            questions: vec![crate::form::FormQuestion {
430                id: "question".into(),
431                kind: FormKind::Confirm { prompt: "?".into() },
432            }],
433        };
434        let pending = PendingForm {
435            form_id: form_id.clone(),
436            run_id: crate::event::FlowRunId::now(),
437            tool_use_id: "tool".into(),
438            form,
439            kind: FormKind::Confirm { prompt: "?".into() },
440            emitted_at: chrono::Utc::now(),
441        };
442        let rx = forms.request(pending);
443        assert_eq!(
444            await_local_submission(
445                &forms,
446                form_id,
447                rx,
448                std::time::Duration::ZERO,
449                &tokio_util::sync::CancellationToken::new(),
450            )
451            .await,
452            crate::form::FormSubmission::Rejected
453        );
454        assert_eq!(forms.list_pending().len(), 1);
455        assert!(forms.submit(
456            "timed-out",
457            crate::form::FormSubmission::Submitted {
458                answers: vec![FormAnswer::Confirmed { value: true }],
459            }
460        ));
461        assert!(forms.list_pending().is_empty());
462    }
463
464    #[tokio::test]
465    async fn local_form_cancellation_removes_pending_entry() {
466        let forms = crate::session::FormRegistry::new();
467        let _subscriber = forms.subscribe();
468        let form_id = "cancelled".to_string();
469        let pending = PendingForm {
470            form_id: form_id.clone(),
471            run_id: crate::event::FlowRunId::now(),
472            tool_use_id: "tool".into(),
473            form: crate::form::CompositeForm {
474                questions: vec![crate::form::FormQuestion {
475                    id: "question".into(),
476                    kind: FormKind::Confirm { prompt: "?".into() },
477                }],
478            },
479            kind: FormKind::Confirm { prompt: "?".into() },
480            emitted_at: chrono::Utc::now(),
481        };
482        let rx = forms.request(pending);
483        let cancel = tokio_util::sync::CancellationToken::new();
484        cancel.cancel();
485
486        assert_eq!(
487            await_local_submission(
488                &forms,
489                form_id,
490                rx,
491                std::time::Duration::from_secs(300),
492                &cancel,
493            )
494            .await,
495            crate::form::FormSubmission::Rejected
496        );
497        assert!(forms.list_pending().is_empty());
498    }
499
500    #[test]
501    fn parse_composite_questions() {
502        let question = |id: &str, prompt: &str| {
503            Value::Struct(vec![
504                ("id".into(), Value::Str(id.into())),
505                ("kind".into(), Value::Str("text".into())),
506                ("prompt".into(), Value::Str(prompt.into())),
507            ])
508        };
509        let args = ToolArgs {
510            positional: vec![],
511            named: vec![(
512                "questions".into(),
513                Value::List(vec![question("name", "Name?"), question("team", "Team?")]),
514            )],
515        };
516        let (form, first, composite) = parse_form_request(&args).unwrap();
517        assert!(composite);
518        assert_eq!(form.questions.len(), 2);
519        assert_eq!(form.questions[0].id, "name");
520        assert_eq!(form.questions[1].id, "team");
521        assert!(matches!(first, FormKind::Text { .. }));
522    }
523
524    #[test]
525    fn parse_composite_questions_rejects_duplicate_ids() {
526        let question = |id: &str| {
527            Value::Struct(vec![
528                ("id".into(), Value::Str(id.into())),
529                ("kind".into(), Value::Str("confirm".into())),
530                ("prompt".into(), Value::Str("Continue?".into())),
531            ])
532        };
533        let args = ToolArgs {
534            positional: vec![],
535            named: vec![(
536                "questions".into(),
537                Value::List(vec![question("same"), question("same")]),
538            )],
539        };
540        let error = parse_form_request(&args).unwrap_err();
541        assert!(error.to_string().contains("duplicate question id"));
542    }
543
544    #[test]
545    fn parse_confirm_kind() {
546        let args = ToolArgs {
547            positional: vec![],
548            named: vec![
549                named("kind", Value::Str("confirm".into())),
550                named("prompt", Value::Str("sure?".into())),
551            ],
552        };
553        assert!(matches!(
554            parse_form_kind(&args).unwrap(),
555            FormKind::Confirm { .. }
556        ));
557    }
558
559    #[test]
560    fn parse_single_select_rejects_empty_options() {
561        let args = ToolArgs {
562            positional: vec![],
563            named: vec![
564                named("kind", Value::Str("single_select".into())),
565                named("prompt", Value::Str("pick".into())),
566                named("options", Value::List(vec![])),
567            ],
568        };
569        let err = parse_form_kind(&args).unwrap_err();
570        assert!(err.to_string().contains("non-empty"));
571    }
572
573    #[test]
574    fn parse_multi_select_validates_bounds() {
575        let args = ToolArgs {
576            positional: vec![],
577            named: vec![
578                named("kind", Value::Str("multi_select".into())),
579                named("prompt", Value::Str("tags".into())),
580                named(
581                    "options",
582                    Value::List(vec![Value::Str("a".into()), Value::Str("b".into())]),
583                ),
584                named("min", Value::Int(3)),
585                named("max", Value::Int(1)),
586            ],
587        };
588        let err = parse_form_kind(&args).unwrap_err();
589        assert!(err.to_string().contains("max must be >= min"));
590    }
591
592    #[test]
593    fn parse_text_defaults_multiline_to_false() {
594        let args = ToolArgs {
595            positional: vec![],
596            named: vec![
597                named("kind", Value::Str("text".into())),
598                named("prompt", Value::Str("name?".into())),
599            ],
600        };
601        match parse_form_kind(&args).unwrap() {
602            FormKind::Text { multiline, .. } => assert!(!multiline),
603            other => panic!("expected text, got {other:?}"),
604        }
605    }
606
607    #[test]
608    fn parse_unknown_kind_errors_with_hint() {
609        let args = ToolArgs {
610            positional: vec![],
611            named: vec![
612                named("kind", Value::Str("weird".into())),
613                named("prompt", Value::Str("?".into())),
614            ],
615        };
616        let err = parse_form_kind(&args).unwrap_err();
617        assert!(err.to_string().contains("weird"));
618        assert!(err.to_string().contains("confirm"));
619    }
620
621    #[test]
622    fn answer_confirmed_becomes_struct() {
623        let v = answer_to_value(&FormAnswer::Confirmed { value: true });
624        assert_eq!(v.field("kind").unwrap().kind_name(), "string");
625        assert!(matches!(v.field("value"), Some(Value::Bool(true))));
626    }
627
628    #[test]
629    fn answer_multi_selected_carries_indices_and_labels() {
630        let v = answer_to_value(&FormAnswer::MultiSelected {
631            indices: vec![0, 2],
632            labels: vec!["a".into(), "c".into()],
633        });
634        let indices = match v.field("indices").unwrap() {
635            Value::List(l) => l,
636            _ => panic!("expected list"),
637        };
638        assert_eq!(indices.len(), 2);
639    }
640
641    #[test]
642    fn answer_cancelled_is_kind_only_struct() {
643        let v = answer_to_value(&FormAnswer::Cancelled);
644        assert!(matches!(v.field("kind"), Some(Value::Str(s)) if s == "cancelled"));
645        assert!(v.field("value").is_none());
646    }
647}