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