Skip to main content

atman_runtime/tools/
stdlib.rs

1use crate::approval::{ApprovalOutcome, request_approval};
2use crate::error::RuntimeError;
3use crate::tool::{BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult};
4use crate::value::Value;
5
6pub struct ShellQuote;
7
8impl Tool for ShellQuote {
9    fn name(&self) -> &str {
10        "shell_quote"
11    }
12
13    fn tier(&self) -> Tier {
14        Tier::Zero
15    }
16
17    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
18        Box::pin(async move {
19            let s = extract_string(&args, "s", 0)?;
20            Ok(Value::Str(shell_quote(&s)))
21        })
22    }
23}
24
25pub fn shell_quote(s: &str) -> String {
26    // POSIX-safe: wrap in single quotes, escape any internal ' as '\''.
27    let mut out = String::with_capacity(s.len() + 2);
28    out.push('\'');
29    for c in s.chars() {
30        if c == '\'' {
31            out.push_str("'\\''");
32        } else {
33            out.push(c);
34        }
35    }
36    out.push('\'');
37    out
38}
39
40pub struct Len;
41
42impl Tool for Len {
43    fn name(&self) -> &str {
44        "len"
45    }
46
47    fn tier(&self) -> Tier {
48        Tier::Zero
49    }
50
51    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
52        Box::pin(async move {
53            let v = args.positional(0)?;
54            match v {
55                Value::List(items) => Ok(Value::Int(items.len() as i64)),
56                Value::Str(s) => Ok(Value::Int(s.chars().count() as i64)),
57                other => Err(RuntimeError::TypeMismatch {
58                    expected: "list or string".into(),
59                    actual: other.kind_name().into(),
60                }),
61            }
62        })
63    }
64}
65
66pub struct Head;
67
68impl Tool for Head {
69    fn name(&self) -> &str {
70        "head"
71    }
72
73    fn tier(&self) -> Tier {
74        Tier::Zero
75    }
76
77    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
78        Box::pin(async move {
79            match args.positional(0)? {
80                Value::List(items) => items
81                    .first()
82                    .cloned()
83                    .ok_or_else(|| RuntimeError::ToolFailed("head: empty list".into())),
84                other => Err(RuntimeError::TypeMismatch {
85                    expected: "list".into(),
86                    actual: other.kind_name().into(),
87                }),
88            }
89        })
90    }
91}
92
93pub struct Tail;
94
95impl Tool for Tail {
96    fn name(&self) -> &str {
97        "tail"
98    }
99
100    fn tier(&self) -> Tier {
101        Tier::Zero
102    }
103
104    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
105        Box::pin(async move {
106            match args.positional(0)? {
107                Value::List(items) if !items.is_empty() => Ok(Value::List(items[1..].to_vec())),
108                Value::List(_) => Err(RuntimeError::ToolFailed("tail: empty list".into())),
109                other => Err(RuntimeError::TypeMismatch {
110                    expected: "list".into(),
111                    actual: other.kind_name().into(),
112                }),
113            }
114        })
115    }
116}
117
118pub struct IsEmpty;
119
120impl Tool for IsEmpty {
121    fn name(&self) -> &str {
122        "is_empty"
123    }
124
125    fn tier(&self) -> Tier {
126        Tier::Zero
127    }
128
129    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
130        Box::pin(async move {
131            let v = args.positional(0)?;
132            match v {
133                Value::List(items) => Ok(Value::Bool(items.is_empty())),
134                Value::Str(s) => Ok(Value::Bool(s.is_empty())),
135                other => Err(RuntimeError::TypeMismatch {
136                    expected: "list or string".into(),
137                    actual: other.kind_name().into(),
138                }),
139            }
140        })
141    }
142}
143
144pub struct EstimateTokens;
145
146impl Tool for EstimateTokens {
147    fn name(&self) -> &str {
148        "estimate_tokens"
149    }
150    fn tier(&self) -> Tier {
151        Tier::Zero
152    }
153    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
154        Box::pin(async move {
155            let v = args.positional(0)?;
156            match v {
157                Value::List(items) => {
158                    let mut msgs = Vec::with_capacity(items.len());
159                    for it in items {
160                        match it {
161                            Value::Message(m) => msgs.push(m.clone()),
162                            other => {
163                                return Err(RuntimeError::TypeMismatch {
164                                    expected: "list of message".into(),
165                                    actual: other.kind_name().into(),
166                                });
167                            }
168                        }
169                    }
170                    let n = crate::compaction::estimate_tokens_for_messages(&msgs);
171                    Ok(Value::Int(n as i64))
172                }
173                Value::Message(m) => Ok(Value::Int(
174                    crate::compaction::estimate_tokens_for_message(m) as i64,
175                )),
176                Value::Str(s) => {
177                    let approx = ((s.len() as f64) / 3.5).ceil() as i64;
178                    Ok(Value::Int(approx))
179                }
180                other => Err(RuntimeError::TypeMismatch {
181                    expected: "message | list of message | string".into(),
182                    actual: other.kind_name().into(),
183                }),
184            }
185        })
186    }
187}
188
189pub struct FindCompactRange;
190
191impl Tool for FindCompactRange {
192    fn name(&self) -> &str {
193        "find_compact_range"
194    }
195    fn tier(&self) -> Tier {
196        Tier::Zero
197    }
198    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
199        Box::pin(async move {
200            let messages = extract_message_list(&args, "messages", 0)?;
201            let budget = extract_int(&args, "budget", 1)? as u64;
202            match crate::compaction::find_compact_range(&messages, budget) {
203                Some(range) => Ok(Value::Struct(vec![
204                    ("start".into(), Value::Int(range.start as i64)),
205                    ("end".into(), Value::Int(range.end as i64)),
206                    (
207                        "tokens_saved".into(),
208                        Value::Int(range.tokens_saved_estimate as i64),
209                    ),
210                    ("found".into(), Value::Bool(true)),
211                ])),
212                None => Ok(Value::Struct(vec![
213                    ("start".into(), Value::Int(0)),
214                    ("end".into(), Value::Int(0)),
215                    ("tokens_saved".into(), Value::Int(0)),
216                    ("found".into(), Value::Bool(false)),
217                ])),
218            }
219        })
220    }
221}
222
223pub struct ReplaceMessagesRange;
224
225impl Tool for ReplaceMessagesRange {
226    fn name(&self) -> &str {
227        "replace_messages_range"
228    }
229    fn tier(&self) -> Tier {
230        Tier::Zero
231    }
232    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
233        Box::pin(async move {
234            let messages = extract_message_list(&args, "messages", 0)?;
235            let start = extract_int(&args, "start", 1)? as usize;
236            let end = extract_int(&args, "end", 2)? as usize;
237            let summary = extract_string_arg(&args, "summary", 3)?;
238            if start > end || end > messages.len() {
239                return Err(RuntimeError::ToolFailed(format!(
240                    "replace_messages_range: invalid range start={start} end={end} len={}",
241                    messages.len()
242                )));
243            }
244            let before_tokens = crate::compaction::estimate_tokens_for_messages(&messages);
245            let seq_span = messages
246                .get(start..end.min(messages.len()))
247                .and_then(|slice| {
248                    Some((
249                        slice.first().map(|_| start as u64)?,
250                        slice.last().map(|_| end.saturating_sub(1) as u64)?,
251                    ))
252                })
253                .unwrap_or((start as u64, end.saturating_sub(1) as u64));
254            let range = crate::compaction::CompactRange {
255                start,
256                end,
257                tokens_saved_estimate: 0,
258            };
259            let turn_id = messages
260                .first()
261                .map(|m| m.turn_id.clone())
262                .unwrap_or_else(crate::event::TurnId::now);
263            let out =
264                crate::compaction::replace_range_with_summary(&messages, &range, summary, turn_id);
265            let after_tokens = crate::compaction::estimate_tokens_for_messages(&out);
266            if let Some(sink) = &ctx.events {
267                sink.mark_compacted();
268                sink.emit(crate::event::Event::ContextCompact {
269                    session_id: ctx
270                        .turn_id
271                        .as_ref()
272                        .map(|t| t.0.to_string())
273                        .unwrap_or_default(),
274                    before_tokens,
275                    after_tokens,
276                    compacted_range_start: seq_span.0,
277                    compacted_range_end: seq_span.1,
278                    summary_text: None,
279                    replacement_msg_seq: None,
280                });
281            }
282            if let Some(tx) = &ctx.lifecycle_fire_tx {
283                let _ = tx.send(atman_dsl::ast::LifecycleEvent::ContextCompact);
284            }
285            let list: Vec<Value> = out.into_iter().map(Value::Message).collect();
286            Ok(Value::List(list))
287        })
288    }
289}
290
291fn extract_message_list(
292    args: &ToolArgs,
293    name: &str,
294    pos: usize,
295) -> Result<Vec<crate::message::Message>, RuntimeError> {
296    let value = match args.named(name) {
297        Some(v) => v,
298        None => args.positional(pos)?,
299    };
300    match value {
301        Value::List(items) => {
302            let mut out = Vec::with_capacity(items.len());
303            for it in items {
304                match it {
305                    Value::Message(m) => out.push(m.clone()),
306                    other => {
307                        return Err(RuntimeError::TypeMismatch {
308                            expected: "list of message".into(),
309                            actual: other.kind_name().into(),
310                        });
311                    }
312                }
313            }
314            Ok(out)
315        }
316        other => Err(RuntimeError::TypeMismatch {
317            expected: "list of message".into(),
318            actual: other.kind_name().into(),
319        }),
320    }
321}
322
323fn extract_int(args: &ToolArgs, name: &str, pos: usize) -> Result<i64, RuntimeError> {
324    let value = match args.named(name) {
325        Some(v) => v,
326        None => args.positional(pos)?,
327    };
328    match value {
329        Value::Int(n) => Ok(*n),
330        other => Err(RuntimeError::TypeMismatch {
331            expected: "int".into(),
332            actual: other.kind_name().into(),
333        }),
334    }
335}
336
337fn extract_string_arg(args: &ToolArgs, name: &str, pos: usize) -> Result<String, RuntimeError> {
338    let value = match args.named(name) {
339        Some(v) => v,
340        None => args.positional(pos)?,
341    };
342    match value {
343        Value::Str(s) => Ok(s.clone()),
344        other => Err(RuntimeError::TypeMismatch {
345            expected: "string".into(),
346            actual: other.kind_name().into(),
347        }),
348    }
349}
350
351pub struct RenderPromptXml;
352pub struct RenderPromptMarkdown;
353pub struct RenderPromptTerse;
354
355fn extract_prompt_spec(v: &Value) -> Result<PromptSpec<'_>, RuntimeError> {
356    let Value::Struct(fields) = v else {
357        return Err(RuntimeError::TypeMismatch {
358            expected: "struct { role?, context?, task, examples?, schema? }".into(),
359            actual: v.kind_name().into(),
360        });
361    };
362    let get = |k: &str| fields.iter().find(|(n, _)| n == k).map(|(_, v)| v);
363    let task = match get("task") {
364        Some(Value::Str(s)) => s.clone(),
365        Some(other) => {
366            return Err(RuntimeError::TypeMismatch {
367                expected: "string (task)".into(),
368                actual: other.kind_name().into(),
369            });
370        }
371        None => return Err(RuntimeError::MissingArg("prompt.task".into())),
372    };
373    let role = match get("role") {
374        Some(Value::Str(s)) => Some(s.clone()),
375        Some(Value::Unit) | None => None,
376        Some(other) => {
377            return Err(RuntimeError::TypeMismatch {
378                expected: "string (role)".into(),
379                actual: other.kind_name().into(),
380            });
381        }
382    };
383    let context = get("context");
384    let schema = match get("schema") {
385        Some(Value::Str(s)) => Some(s.clone()),
386        _ => None,
387    };
388    let examples = match get("examples") {
389        Some(Value::List(items)) => items.iter().collect(),
390        _ => Vec::new(),
391    };
392    Ok(PromptSpec {
393        role,
394        context,
395        task,
396        examples,
397        schema,
398    })
399}
400
401struct PromptSpec<'a> {
402    role: Option<String>,
403    context: Option<&'a Value>,
404    task: String,
405    examples: Vec<&'a Value>,
406    schema: Option<String>,
407}
408
409fn json_str(v: &Value) -> String {
410    serde_json::to_string_pretty(&v.to_json()).unwrap_or_default()
411}
412
413fn render_xml(spec: &PromptSpec<'_>) -> String {
414    let mut out = String::new();
415    if let Some(role) = &spec.role {
416        out.push_str(&format!("<role>{}</role>\n", role));
417    }
418    if let Some(ctx) = spec.context {
419        out.push_str(&format!("<context>\n{}\n</context>\n", json_str(ctx)));
420    }
421    if !spec.examples.is_empty() {
422        out.push_str("<examples>\n");
423        for (i, ex) in spec.examples.iter().enumerate() {
424            out.push_str(&format!(
425                "  <example n=\"{}\">\n{}\n  </example>\n",
426                i + 1,
427                json_str(ex)
428            ));
429        }
430        out.push_str("</examples>\n");
431    }
432    out.push_str(&format!("<task>{}</task>\n", spec.task));
433    if let Some(schema) = &spec.schema {
434        out.push_str(&format!("<schema>{}</schema>\n", schema));
435    }
436    out
437}
438
439fn render_markdown(spec: &PromptSpec<'_>) -> String {
440    let mut out = String::new();
441    if let Some(role) = &spec.role {
442        out.push_str(&format!("# Role\n{}\n\n", role));
443    }
444    if let Some(ctx) = spec.context {
445        out.push_str(&format!("# Context\n```json\n{}\n```\n\n", json_str(ctx)));
446    }
447    if !spec.examples.is_empty() {
448        out.push_str("# Examples\n");
449        for (i, ex) in spec.examples.iter().enumerate() {
450            out.push_str(&format!(
451                "{}. `{}`\n",
452                i + 1,
453                json_str(ex).replace('\n', " ")
454            ));
455        }
456        out.push('\n');
457    }
458    out.push_str(&format!("# Task\n{}\n", spec.task));
459    if let Some(schema) = &spec.schema {
460        out.push_str(&format!("\n# Schema\n{}\n", schema));
461    }
462    out
463}
464
465fn render_terse(spec: &PromptSpec<'_>) -> String {
466    let mut out = String::new();
467    if let Some(role) = &spec.role {
468        out.push_str(&format!("Role: {}\n", role));
469    }
470    if let Some(ctx) = spec.context {
471        out.push_str(&format!("Context: {}\n", json_str(ctx).replace('\n', " ")));
472    }
473    out.push_str(&format!("Task: {}\n", spec.task));
474    if let Some(schema) = &spec.schema {
475        out.push_str(&format!("Schema: {}\n", schema));
476    }
477    for (i, ex) in spec.examples.iter().enumerate() {
478        out.push_str(&format!(
479            "Example {}: {}\n",
480            i + 1,
481            json_str(ex).replace('\n', " ")
482        ));
483    }
484    out
485}
486
487impl Tool for RenderPromptXml {
488    fn name(&self) -> &str {
489        "render_prompt_xml"
490    }
491    fn tier(&self) -> Tier {
492        Tier::Zero
493    }
494    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
495        Box::pin(async move {
496            let v = args.positional(0)?;
497            let spec = extract_prompt_spec(v)?;
498            Ok(Value::Str(render_xml(&spec)))
499        })
500    }
501}
502
503impl Tool for RenderPromptMarkdown {
504    fn name(&self) -> &str {
505        "render_prompt_markdown"
506    }
507    fn tier(&self) -> Tier {
508        Tier::Zero
509    }
510    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
511        Box::pin(async move {
512            let v = args.positional(0)?;
513            let spec = extract_prompt_spec(v)?;
514            Ok(Value::Str(render_markdown(&spec)))
515        })
516    }
517}
518
519impl Tool for RenderPromptTerse {
520    fn name(&self) -> &str {
521        "render_prompt_terse"
522    }
523    fn tier(&self) -> Tier {
524        Tier::Zero
525    }
526    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
527        Box::pin(async move {
528            let v = args.positional(0)?;
529            let spec = extract_prompt_spec(v)?;
530            Ok(Value::Str(render_terse(&spec)))
531        })
532    }
533}
534
535pub struct ToJsonString;
536
537impl Tool for ToJsonString {
538    fn name(&self) -> &str {
539        "to_json_string"
540    }
541
542    fn tier(&self) -> Tier {
543        Tier::Zero
544    }
545
546    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
547        Box::pin(async move {
548            let v = args.positional(0)?.clone();
549            let json = v.to_json();
550            let s = serde_json::to_string_pretty(&json)
551                .map_err(|e| RuntimeError::ToolFailed(format!("to_json_string: {e}")))?;
552            Ok(Value::Str(s))
553        })
554    }
555}
556
557pub struct TextConcat;
558
559impl Tool for TextConcat {
560    fn name(&self) -> &str {
561        "text_concat"
562    }
563
564    fn tier(&self) -> Tier {
565        Tier::Zero
566    }
567
568    fn description(&self) -> Option<&str> {
569        Some("Flatten the text parts of a Message into a single string.")
570    }
571
572    fn input_schema(&self) -> serde_json::Value {
573        serde_json::json!({
574            "type": "object",
575            "properties": {"message": {"description": "A Message value from an llm call."}},
576            "required": ["message"]
577        })
578    }
579
580    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
581        Box::pin(async move {
582            let v = match args.named("message") {
583                Some(v) => v,
584                None => args.positional(0)?,
585            };
586            match v {
587                Value::Message(m) => Ok(Value::Str(m.text_concat())),
588                Value::Str(s) => Ok(Value::Str(s.clone())),
589                other => Err(RuntimeError::TypeMismatch {
590                    expected: "message or string".into(),
591                    actual: other.kind_name().into(),
592                }),
593            }
594        })
595    }
596}
597
598pub struct Concat;
599
600impl Tool for Concat {
601    fn name(&self) -> &str {
602        "concat"
603    }
604
605    fn tier(&self) -> Tier {
606        Tier::Zero
607    }
608
609    fn description(&self) -> Option<&str> {
610        Some("Concatenate two lists into a single new list.")
611    }
612
613    fn input_schema(&self) -> serde_json::Value {
614        serde_json::json!({
615            "type": "object",
616            "properties": {
617                "left": {"type": "array"},
618                "right": {"type": "array"}
619            },
620            "required": ["left", "right"]
621        })
622    }
623
624    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
625        Box::pin(async move {
626            let left = extract_list(&args, "left", 0)?;
627            let right = extract_list(&args, "right", 1)?;
628            let mut out = Vec::with_capacity(left.len() + right.len());
629            out.extend(left);
630            out.extend(right);
631            Ok(Value::List(out))
632        })
633    }
634}
635
636pub struct MessageUser;
637
638impl Tool for MessageUser {
639    fn name(&self) -> &str {
640        "message.user"
641    }
642    fn tier(&self) -> Tier {
643        Tier::Zero
644    }
645    fn description(&self) -> Option<&str> {
646        Some(
647            "Construct a user-role Message from a text string. Use with session.push to inject user instructions into the session history before an llm.call(context: session) call.",
648        )
649    }
650    fn input_schema(&self) -> serde_json::Value {
651        serde_json::json!({
652            "type": "object",
653            "properties": {"text": {"type": "string"}},
654            "required": ["text"]
655        })
656    }
657    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
658        Box::pin(async move {
659            let text = extract_string(&args, "text", 0)?;
660            let turn_id = ctx
661                .turn_id
662                .clone()
663                .unwrap_or_else(crate::event::TurnId::now);
664            Ok(Value::Message(crate::message::Message::user_text(
665                turn_id, text,
666            )))
667        })
668    }
669}
670
671pub struct MessageAssistant;
672
673impl Tool for MessageAssistant {
674    fn name(&self) -> &str {
675        "message.assistant"
676    }
677    fn tier(&self) -> Tier {
678        Tier::Zero
679    }
680    fn description(&self) -> Option<&str> {
681        Some("Construct an assistant-role Message from a text string.")
682    }
683    fn input_schema(&self) -> serde_json::Value {
684        serde_json::json!({
685            "type": "object",
686            "properties": {"text": {"type": "string"}},
687            "required": ["text"]
688        })
689    }
690    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
691        Box::pin(async move {
692            let text = extract_string(&args, "text", 0)?;
693            let turn_id = ctx
694                .turn_id
695                .clone()
696                .unwrap_or_else(crate::event::TurnId::now);
697            Ok(Value::Message(crate::message::Message::assistant_text(
698                turn_id, text,
699            )))
700        })
701    }
702}
703
704pub struct MessageSystem;
705
706impl Tool for MessageSystem {
707    fn name(&self) -> &str {
708        "message.system"
709    }
710    fn tier(&self) -> Tier {
711        Tier::Zero
712    }
713    fn description(&self) -> Option<&str> {
714        Some("Construct a system-role Message from a text string.")
715    }
716    fn input_schema(&self) -> serde_json::Value {
717        serde_json::json!({
718            "type": "object",
719            "properties": {"text": {"type": "string"}},
720            "required": ["text"]
721        })
722    }
723    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
724        Box::pin(async move {
725            let text = extract_string(&args, "text", 0)?;
726            let turn_id = ctx
727                .turn_id
728                .clone()
729                .unwrap_or_else(crate::event::TurnId::now);
730            Ok(Value::Message(crate::message::Message::system_text(
731                turn_id, text,
732            )))
733        })
734    }
735}
736
737pub struct MessageTool;
738
739impl Tool for MessageTool {
740    fn name(&self) -> &str {
741        "message.tool"
742    }
743    fn tier(&self) -> Tier {
744        Tier::Zero
745    }
746    fn description(&self) -> Option<&str> {
747        Some(
748            "Construct a tool-role Message from a text string. Rarely needed directly — dispatch_all already returns tool-role Messages.",
749        )
750    }
751    fn input_schema(&self) -> serde_json::Value {
752        serde_json::json!({
753            "type": "object",
754            "properties": {"text": {"type": "string"}},
755            "required": ["text"]
756        })
757    }
758    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
759        Box::pin(async move {
760            let text = extract_string(&args, "text", 0)?;
761            let turn_id = ctx
762                .turn_id
763                .clone()
764                .unwrap_or_else(crate::event::TurnId::now);
765            Ok(Value::Message(crate::message::Message {
766                turn_id,
767                role: crate::message::MessageRole::Tool,
768                parts: vec![crate::message::MessagePart::Text { text }],
769                origin: crate::message::MessageOrigin::User,
770            }))
771        })
772    }
773}
774
775pub struct ExtractToolUses;
776
777impl Tool for ExtractToolUses {
778    fn name(&self) -> &str {
779        "extract_tool_uses"
780    }
781
782    fn tier(&self) -> Tier {
783        Tier::Zero
784    }
785
786    fn description(&self) -> Option<&str> {
787        Some(
788            "Pull the tool_use parts out of an assistant Message. Returns a list of \
789             {id, name, input} structs suitable for dispatch_all.",
790        )
791    }
792
793    fn input_schema(&self) -> serde_json::Value {
794        serde_json::json!({
795            "type": "object",
796            "properties": {"message": {"description": "Assistant Message value."}},
797            "required": ["message"]
798        })
799    }
800
801    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
802        Box::pin(async move {
803            let v = match args.named("message") {
804                Some(v) => v,
805                None => args.positional(0)?,
806            };
807            let m = match v {
808                Value::Message(m) => m,
809                Value::Str(_) => return Ok(Value::List(Vec::new())),
810                other => {
811                    return Err(RuntimeError::TypeMismatch {
812                        expected: "message or string".into(),
813                        actual: other.kind_name().into(),
814                    });
815                }
816            };
817            let mut out = Vec::new();
818            for part in &m.parts {
819                if let crate::message::MessagePart::ToolUse { id, name, input } = part {
820                    out.push(Value::Struct(vec![
821                        ("id".into(), Value::Str(id.clone())),
822                        ("name".into(), Value::Str(name.clone())),
823                        ("input".into(), Value::from_json(input.clone())),
824                    ]));
825                }
826            }
827            Ok(Value::List(out))
828        })
829    }
830}
831
832pub struct DispatchAll;
833
834impl Tool for DispatchAll {
835    fn name(&self) -> &str {
836        "dispatch_all"
837    }
838
839    fn tier(&self) -> Tier {
840        Tier::Zero
841    }
842
843    fn description(&self) -> Option<&str> {
844        Some(
845            "Dispatch each tool_use in the list against the current tool registry and \
846             return a list of tool_result Message values.",
847        )
848    }
849
850    fn input_schema(&self) -> serde_json::Value {
851        serde_json::json!({
852            "type": "object",
853            "properties": {"tool_uses": {"type": "array"}},
854            "required": ["tool_uses"]
855        })
856    }
857
858    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
859        Box::pin(async move {
860            let uses = extract_list(&args, "tool_uses", 0)?;
861            let Some(registry) = ctx.registry.as_ref() else {
862                return Err(RuntimeError::ToolFailed(
863                    "dispatch_all: no tool registry available on ctx".into(),
864                ));
865            };
866            let prepared = prepare_dispatch(&uses, registry.as_ref(), ctx)?;
867            let (auto_batch, serial_batch, mut out_slots) = partition_and_gate(prepared, ctx).await;
868            run_auto_parallel(auto_batch, ctx, &mut out_slots).await;
869            run_serial(serial_batch, ctx, &mut out_slots).await;
870            let out: Vec<Value> = out_slots.into_iter().flatten().collect();
871            Ok(Value::List(out))
872        })
873    }
874}
875
876enum PreparedEntry {
877    Ready {
878        index: usize,
879        id: String,
880        name: String,
881        tool: std::sync::Arc<dyn Tool>,
882        call_args: ToolArgs,
883    },
884    Failed {
885        index: usize,
886        msg: crate::message::Message,
887    },
888}
889
890fn prepare_dispatch(
891    uses: &[Value],
892    registry: &crate::tool::ToolRegistry,
893    ctx: &ToolCtx,
894) -> Result<Vec<PreparedEntry>, RuntimeError> {
895    let mut prepared = Vec::with_capacity(uses.len());
896    for (index, entry) in uses.iter().enumerate() {
897        let Value::Struct(fields) = entry else {
898            return Err(RuntimeError::TypeMismatch {
899                expected: "struct {id, name, input}".into(),
900                actual: entry.kind_name().into(),
901            });
902        };
903        let get = |k: &str| fields.iter().find(|(n, _)| n == k).map(|(_, v)| v.clone());
904        let id = match get("id") {
905            Some(Value::Str(s)) => s,
906            _ => {
907                return Err(RuntimeError::ToolFailed(
908                    "dispatch_all: tool_use missing `id` string".into(),
909                ));
910            }
911        };
912        let name = match get("name") {
913            Some(Value::Str(s)) => s,
914            _ => {
915                return Err(RuntimeError::ToolFailed(
916                    "dispatch_all: tool_use missing `name` string".into(),
917                ));
918            }
919        };
920        let input = get("input").unwrap_or(Value::Unit);
921        let Some(tool) = registry.get(&name) else {
922            prepared.push(PreparedEntry::Failed {
923                index,
924                msg: build_error_result(ctx, &id, &format!("dispatch_all: unknown tool `{name}`")),
925            });
926            continue;
927        };
928        let named = match &input {
929            Value::Struct(fields) => fields.clone(),
930            Value::Unit => Vec::new(),
931            other => {
932                return Err(RuntimeError::TypeMismatch {
933                    expected: "struct or unit for tool input".into(),
934                    actual: other.kind_name().into(),
935                });
936            }
937        };
938        let missing = missing_required_fields(&tool.input_schema(), &named);
939        if !missing.is_empty() {
940            let content = format!(
941                "tool `{name}` received empty/incomplete input. Missing required fields: {}. Retry with a complete argument object like {{{}}} — do NOT reuse an empty {{}} input.",
942                missing.join(", "),
943                missing
944                    .iter()
945                    .map(|f| format!("\"{f}\":\"...\""))
946                    .collect::<Vec<_>>()
947                    .join(", ")
948            );
949            prepared.push(PreparedEntry::Failed {
950                index,
951                msg: build_error_result(ctx, &id, &content),
952            });
953            continue;
954        }
955        emit_tool_node(ctx, &id, &name, &input);
956        prepared.push(PreparedEntry::Ready {
957            index,
958            id,
959            name,
960            tool,
961            call_args: ToolArgs {
962                positional: Vec::new(),
963                named,
964            },
965        });
966    }
967    Ok(prepared)
968}
969
970struct Approved {
971    index: usize,
972    id: String,
973    name: String,
974    tool: std::sync::Arc<dyn Tool>,
975    call_args: ToolArgs,
976}
977
978async fn partition_and_gate(
979    prepared: Vec<PreparedEntry>,
980    ctx: &ToolCtx,
981) -> (Vec<Approved>, Vec<Approved>, Vec<Option<Value>>) {
982    let total = prepared.len();
983    let mut out_slots: Vec<Option<Value>> = vec![None; total];
984    struct ReadyEntry {
985        index: usize,
986        id: String,
987        name: String,
988        tool: std::sync::Arc<dyn Tool>,
989        call_args: ToolArgs,
990    }
991    let mut ready: Vec<ReadyEntry> = Vec::new();
992    for entry in prepared {
993        match entry {
994            PreparedEntry::Failed { index, msg } => {
995                emit_tool_result(ctx, &msg);
996                out_slots[index] = Some(Value::Message(msg));
997            }
998            PreparedEntry::Ready {
999                index,
1000                id,
1001                name,
1002                tool,
1003                call_args,
1004            } => {
1005                ready.push(ReadyEntry {
1006                    index,
1007                    id,
1008                    name,
1009                    tool,
1010                    call_args,
1011                });
1012            }
1013        }
1014    }
1015    // Parallel: serial awaits hid all but the first pending node from the UI.
1016    let gates = ready.iter().map(|r| {
1017        let level = r.tool.approval_level(&r.call_args, ctx);
1018        request_approval(
1019            ctx,
1020            &r.id,
1021            &r.name,
1022            &r.call_args,
1023            level,
1024            Some(r.tool.as_ref()),
1025        )
1026    });
1027    let outcomes = futures::future::join_all(gates).await;
1028    let mut auto_batch = Vec::new();
1029    let mut serial_batch = Vec::new();
1030    for (r, outcome) in ready.into_iter().zip(outcomes) {
1031        let level = r.tool.approval_level(&r.call_args, ctx);
1032        match outcome {
1033            ApprovalOutcome::Approve => {
1034                let a = Approved {
1035                    index: r.index,
1036                    id: r.id,
1037                    name: r.name.clone(),
1038                    tool: r.tool,
1039                    call_args: r.call_args,
1040                };
1041                if level == crate::tool::ApprovalLevel::Auto {
1042                    auto_batch.push(a);
1043                } else {
1044                    serial_batch.push(a);
1045                }
1046            }
1047            ApprovalOutcome::Deny { reason } => {
1048                let msg = build_error_result(
1049                    ctx,
1050                    &r.id,
1051                    &format!("tool `{}` denied by user: {reason}", r.name),
1052                );
1053                emit_tool_result(ctx, &msg);
1054                out_slots[r.index] = Some(Value::Message(msg));
1055            }
1056        }
1057    }
1058    (auto_batch, serial_batch, out_slots)
1059}
1060
1061async fn run_auto_parallel(batch: Vec<Approved>, ctx: &ToolCtx, out_slots: &mut [Option<Value>]) {
1062    if batch.is_empty() {
1063        return;
1064    }
1065    let futs = batch.iter().map(|a| a.tool.call(a.call_args.clone(), ctx));
1066    let results = futures::future::join_all(futs).await;
1067    for (a, r) in batch.into_iter().zip(results) {
1068        emit_dispatch_node_start(ctx, &a.id, &a.name);
1069        let (content, is_error) = match &r {
1070            Ok(v) => (render_tool_result_text(v), false),
1071            Err(e) => (format!("{e}"), true),
1072        };
1073        emit_dispatch_node_end(ctx, &a.id, &a.name, is_error);
1074        if let Ok(v) = &r {
1075            emit_diff_preview_if_relevant(ctx, &a.name, v);
1076        }
1077        let msg = crate::message::Message {
1078            role: crate::message::MessageRole::Tool,
1079            parts: vec![crate::message::MessagePart::ToolResult {
1080                tool_use_id: a.id.clone(),
1081                content,
1082                is_error,
1083            }],
1084            turn_id: ctx
1085                .turn_id
1086                .clone()
1087                .unwrap_or_else(crate::event::TurnId::now),
1088            origin: crate::message::MessageOrigin::User,
1089        };
1090        emit_tool_result(ctx, &msg);
1091        out_slots[a.index] = Some(Value::Message(msg));
1092    }
1093}
1094
1095async fn run_serial(batch: Vec<Approved>, ctx: &ToolCtx, out_slots: &mut [Option<Value>]) {
1096    for a in batch {
1097        emit_dispatch_node_start(ctx, &a.id, &a.name);
1098        let r = a.tool.call(a.call_args, ctx).await;
1099        let (content, is_error) = match &r {
1100            Ok(v) => (render_tool_result_text(v), false),
1101            Err(e) => (format!("{e}"), true),
1102        };
1103        emit_dispatch_node_end(ctx, &a.id, &a.name, is_error);
1104        if let Ok(v) = &r {
1105            emit_diff_preview_if_relevant(ctx, &a.name, v);
1106        }
1107        let msg = crate::message::Message {
1108            role: crate::message::MessageRole::Tool,
1109            parts: vec![crate::message::MessagePart::ToolResult {
1110                tool_use_id: a.id.clone(),
1111                content,
1112                is_error,
1113            }],
1114            turn_id: ctx
1115                .turn_id
1116                .clone()
1117                .unwrap_or_else(crate::event::TurnId::now),
1118            origin: crate::message::MessageOrigin::User,
1119        };
1120        emit_tool_result(ctx, &msg);
1121        out_slots[a.index] = Some(Value::Message(msg));
1122    }
1123}
1124
1125fn emit_dispatch_node_start(ctx: &ToolCtx, id: &str, name: &str) {
1126    use crate::nodegraph::NodeKind;
1127    let kind = NodeKind::ToolCall {
1128        path: name.to_string(),
1129    };
1130    let label = format!("⟶ {name}");
1131    let node_id = format!("dispatch:{id}");
1132    if let Some(sink) = ctx.events.as_ref()
1133        && let Some(run_id) = ctx.flow_run_id.as_ref()
1134    {
1135        sink.emit(crate::event::Event::FlowNodeStart {
1136            run_id: run_id.clone(),
1137            node_id: node_id.clone(),
1138            kind: kind.clone(),
1139            label: label.clone(),
1140            parent_node_id: ctx.current_node_id.clone(),
1141        });
1142    }
1143    if let Some(tx) = &ctx.stream_tx
1144        && let Some(run_id) = ctx.flow_run_id.as_ref()
1145    {
1146        let _ = tx.send(crate::stream::StreamFrame::FlowNodeStart {
1147            run_id: run_id.0.to_string(),
1148            node_id,
1149            kind,
1150            label,
1151            parent_node_id: ctx.current_node_id.clone(),
1152        });
1153    }
1154}
1155
1156fn emit_dispatch_node_end(ctx: &ToolCtx, id: &str, name: &str, is_error: bool) {
1157    let node_id = format!("dispatch:{id}");
1158    let status = if is_error {
1159        crate::event::FlowNodeStatus::Err
1160    } else {
1161        crate::event::FlowNodeStatus::Ok
1162    };
1163    let preview = name.to_string();
1164    if let Some(sink) = ctx.events.as_ref()
1165        && let Some(run_id) = ctx.flow_run_id.as_ref()
1166    {
1167        sink.emit(crate::event::Event::FlowNodeEnd {
1168            run_id: run_id.clone(),
1169            node_id: node_id.clone(),
1170            status: status.clone(),
1171            output_preview: Some(preview.clone()),
1172        });
1173    }
1174    if let Some(tx) = &ctx.stream_tx
1175        && let Some(run_id) = ctx.flow_run_id.as_ref()
1176    {
1177        let _ = tx.send(crate::stream::StreamFrame::FlowNodeEnd {
1178            run_id: run_id.0.to_string(),
1179            node_id,
1180            status,
1181            output_preview: Some(preview),
1182            parent_node_id: ctx.current_node_id.clone(),
1183        });
1184    }
1185}
1186
1187type DiffPreviewData = (String, Option<String>, Option<String>, Option<String>);
1188
1189fn emit_diff_preview_if_relevant(ctx: &ToolCtx, tool_name: &str, value: &Value) {
1190    let Some(sink) = ctx.events.as_ref() else {
1191        return;
1192    };
1193    let data: Option<DiffPreviewData> = match tool_name {
1194        "fs.edit" => {
1195            let path = value_struct_string(value, "summary").and_then(|s| {
1196                s.strip_prefix("[fs.edit(")
1197                    .and_then(|s| s.split(':').next())
1198                    .map(|s| s.trim_end_matches(')').to_string())
1199            });
1200            let diff = value_struct_string(value, "diff");
1201            diff.map(|d| (path.unwrap_or_default(), None, None, Some(d)))
1202        }
1203        "fs.write" => {
1204            let path = value_struct_string(value, "path").unwrap_or_default();
1205            let diff = value_struct_string(value, "diff");
1206            diff.map(|d| (path, None, None, Some(d)))
1207        }
1208        "git.diff" => {
1209            let Some(diff) = value_struct_string(value, "diff") else {
1210                return;
1211            };
1212            Some(("git diff".into(), None, None, Some(diff)))
1213        }
1214        "git.show" => {
1215            let sha = value_struct_string(value, "sha").unwrap_or_default();
1216            let Some(diff) = value_struct_string(value, "diff") else {
1217                return;
1218            };
1219            Some((format!("git show {sha}"), None, None, Some(diff)))
1220        }
1221        "git.log" => {
1222            let Some(diff) = value_struct_string(value, "diff") else {
1223                return;
1224            };
1225            Some(("git log HEAD".into(), None, None, Some(diff)))
1226        }
1227        _ => None,
1228    };
1229    if let Some((title, old_content, new_content, unified_diff)) = data {
1230        sink.emit(crate::event::Event::DiffPreview {
1231            turn_id: ctx.turn_id.clone(),
1232            flow_run_id: ctx.flow_run_id.clone(),
1233            title,
1234            old_content,
1235            new_content,
1236            unified_diff,
1237        });
1238    }
1239}
1240
1241fn value_struct_string(value: &Value, field: &str) -> Option<String> {
1242    if let Value::Struct(fields) = value {
1243        fields
1244            .iter()
1245            .find(|(k, _)| k == field)
1246            .and_then(|(_, v)| match v {
1247                Value::Str(s) => Some(s.clone()),
1248                _ => None,
1249            })
1250    } else {
1251        None
1252    }
1253}
1254
1255fn emit_tool_node(ctx: &ToolCtx, id: &str, name: &str, input: &Value) {
1256    if let (Some(sink), Some(run_id), Some(parent_node)) = (
1257        ctx.events.as_ref(),
1258        ctx.flow_run_id.clone(),
1259        &ctx.current_node_id,
1260    ) {
1261        let args_preview = format!("{:?}", input)
1262            .chars()
1263            .take(4000)
1264            .collect::<String>();
1265        sink.emit(crate::event::Event::ToolNode {
1266            run_id: run_id.clone(),
1267            parent_node_id: parent_node.clone(),
1268            tool_use_id: id.to_string(),
1269            tool_name: name.to_string(),
1270            args_preview: args_preview.clone(),
1271        });
1272        if let Some(tx) = &ctx.stream_tx {
1273            let _ = tx.send(crate::stream::StreamFrame::ToolNode {
1274                run_id: run_id.0.to_string(),
1275                parent_node_id: parent_node.clone(),
1276                tool_use_id: id.to_string(),
1277                tool: name.to_string(),
1278                args_preview,
1279            });
1280        }
1281    }
1282}
1283
1284fn build_error_result(ctx: &ToolCtx, tool_use_id: &str, content: &str) -> crate::message::Message {
1285    crate::message::Message {
1286        role: crate::message::MessageRole::Tool,
1287        parts: vec![crate::message::MessagePart::ToolResult {
1288            tool_use_id: tool_use_id.to_string(),
1289            content: content.to_string(),
1290            is_error: true,
1291        }],
1292        turn_id: ctx
1293            .turn_id
1294            .clone()
1295            .unwrap_or_else(crate::event::TurnId::now),
1296        origin: crate::message::MessageOrigin::User,
1297    }
1298}
1299
1300fn missing_required_fields(schema: &serde_json::Value, named: &[(String, Value)]) -> Vec<String> {
1301    let Some(required) = schema.get("required").and_then(|v| v.as_array()) else {
1302        return Vec::new();
1303    };
1304    let have: std::collections::HashSet<&str> = named.iter().map(|(k, _)| k.as_str()).collect();
1305    required
1306        .iter()
1307        .filter_map(|v| v.as_str())
1308        .filter(|k| !have.contains(k))
1309        .map(String::from)
1310        .collect()
1311}
1312
1313fn emit_tool_result(ctx: &ToolCtx, msg: &crate::message::Message) {
1314    let msg =
1315        crate::tools::tool_output::maybe_truncate_tool_message(msg, ctx.session_dir.as_deref());
1316    if let Some(tx) = &ctx.stream_tx {
1317        let _ = tx.send(crate::stream::StreamFrame::ToolResultMsg {
1318            flow_run_id: if ctx.session_runtime.is_some() {
1319                None
1320            } else {
1321                ctx.flow_run_id.as_ref().map(|r| r.0.to_string())
1322            },
1323            message: msg,
1324        });
1325    } else if let Some(sink) = &ctx.events {
1326        sink.emit(crate::event::Event::ToolResultMsg {
1327            turn_id: msg.turn_id.clone(),
1328            flow_run_id: ctx.flow_run_id.clone(),
1329            message: msg,
1330        });
1331    }
1332}
1333
1334fn render_tool_result_text(v: &Value) -> String {
1335    match v {
1336        Value::Str(s) => s.clone(),
1337        Value::Message(m) => m.text_concat(),
1338        other => other.to_json().to_string(),
1339    }
1340}
1341
1342fn extract_list(args: &ToolArgs, name: &str, pos: usize) -> Result<Vec<Value>, RuntimeError> {
1343    let value = match args.named(name) {
1344        Some(v) => v,
1345        None => args.positional(pos)?,
1346    };
1347    match value {
1348        Value::List(items) => Ok(items.clone()),
1349        other => Err(RuntimeError::TypeMismatch {
1350            expected: "list".into(),
1351            actual: other.kind_name().into(),
1352        }),
1353    }
1354}
1355
1356pub struct ComposeEmailPreview;
1357
1358impl Tool for ComposeEmailPreview {
1359    fn name(&self) -> &str {
1360        "compose_email_preview"
1361    }
1362
1363    fn tier(&self) -> Tier {
1364        Tier::Zero
1365    }
1366
1367    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
1368        Box::pin(async move {
1369            let subject = extract_string(&args, "subject", 0)?;
1370            let body = extract_string(&args, "body", 1)?;
1371            let to = extract_string_list(&args, "to", 2)?;
1372            Ok(Value::Str(compose_email_preview(&subject, &body, &to)))
1373        })
1374    }
1375}
1376
1377pub fn compose_email_preview(subject: &str, body: &str, to: &[String]) -> String {
1378    format!(
1379        "To: {}
1380Subject: {subject}
1381---
1382{body}",
1383        to.join(", ")
1384    )
1385}
1386
1387fn extract_string(args: &ToolArgs, name: &str, pos: usize) -> Result<String, RuntimeError> {
1388    let value = match args.named(name) {
1389        Some(v) => v,
1390        None => args.positional(pos)?,
1391    };
1392    match value {
1393        Value::Str(s) => Ok(s.clone()),
1394        other => Err(RuntimeError::TypeMismatch {
1395            expected: "string".into(),
1396            actual: other.kind_name().into(),
1397        }),
1398    }
1399}
1400
1401fn extract_string_list(
1402    args: &ToolArgs,
1403    name: &str,
1404    pos: usize,
1405) -> Result<Vec<String>, RuntimeError> {
1406    let value = match args.named(name) {
1407        Some(v) => v,
1408        None => args.positional(pos)?,
1409    };
1410    match value {
1411        Value::List(items) => items
1412            .iter()
1413            .map(|v| match v {
1414                Value::Str(s) => Ok(s.clone()),
1415                other => Err(RuntimeError::TypeMismatch {
1416                    expected: "list of string".into(),
1417                    actual: other.kind_name().into(),
1418                }),
1419            })
1420            .collect(),
1421        other => Err(RuntimeError::TypeMismatch {
1422            expected: "list".into(),
1423            actual: other.kind_name().into(),
1424        }),
1425    }
1426}
1427
1428#[cfg(test)]
1429mod tests {
1430    use super::*;
1431
1432    #[test]
1433    fn shell_quote_wraps_and_escapes() {
1434        assert_eq!(shell_quote("hello"), "'hello'");
1435        assert_eq!(shell_quote("It's fine"), "'It'\\''s fine'");
1436        assert_eq!(shell_quote(""), "''");
1437        assert_eq!(shell_quote("a'b'c"), "'a'\\''b'\\''c'");
1438    }
1439
1440    #[test]
1441    fn compose_email_preview_formats_headers() {
1442        let preview = compose_email_preview(
1443            "Deploy status",
1444            "See attached",
1445            &["a@x.com".into(), "b@x.com".into()],
1446        );
1447        assert_eq!(
1448            preview,
1449            "To: a@x.com, b@x.com\nSubject: Deploy status\n---\nSee attached"
1450        );
1451    }
1452}