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 parsed = uses
896        .iter()
897        .enumerate()
898        .map(|(index, entry)| {
899            let Value::Struct(fields) = entry else {
900                return Err(RuntimeError::TypeMismatch {
901                    expected: "struct {id, name, input}".into(),
902                    actual: entry.kind_name().into(),
903                });
904            };
905            let get = |key: &str| {
906                fields
907                    .iter()
908                    .find(|(name, _)| name == key)
909                    .map(|(_, value)| value.clone())
910            };
911            let id = match get("id") {
912                Some(Value::Str(id)) => id,
913                _ => {
914                    return Err(RuntimeError::ToolFailed(
915                        "dispatch_all: tool_use missing `id` string".into(),
916                    ));
917                }
918            };
919            let name = match get("name") {
920                Some(Value::Str(name)) => name,
921                _ => {
922                    return Err(RuntimeError::ToolFailed(
923                        "dispatch_all: tool_use missing `name` string".into(),
924                    ));
925                }
926            };
927            Ok((index, id, name, get("input").unwrap_or(Value::Unit)))
928        })
929        .collect::<Result<Vec<_>, RuntimeError>>()?;
930
931    let mut prepared = Vec::with_capacity(parsed.len());
932    for (index, id, name, input) in parsed {
933        emit_tool_node(ctx, &id, &name, &input);
934        let Some(tool) = registry.get(&name) else {
935            prepared.push(PreparedEntry::Failed {
936                index,
937                msg: build_error_result(ctx, &id, &format!("dispatch_all: unknown tool `{name}`")),
938            });
939            continue;
940        };
941        let named = match &input {
942            Value::Struct(fields) => fields.clone(),
943            Value::Unit => Vec::new(),
944            other => {
945                prepared.push(PreparedEntry::Failed {
946                    index,
947                    msg: build_error_result(
948                        ctx,
949                        &id,
950                        &format!(
951                            "tool `{name}` expected struct or unit input, got {}",
952                            other.kind_name()
953                        ),
954                    ),
955                });
956                continue;
957            }
958        };
959        let missing = missing_required_fields(&tool.input_schema(), &named);
960        if !missing.is_empty() {
961            let content = format!(
962                "tool `{name}` received empty/incomplete input. Missing required fields: {}. Retry with a complete argument object like {{{}}} — do NOT reuse an empty {{}} input.",
963                missing.join(", "),
964                missing
965                    .iter()
966                    .map(|f| format!("\"{f}\":\"...\""))
967                    .collect::<Vec<_>>()
968                    .join(", ")
969            );
970            prepared.push(PreparedEntry::Failed {
971                index,
972                msg: build_error_result(ctx, &id, &content),
973            });
974            continue;
975        }
976        prepared.push(PreparedEntry::Ready {
977            index,
978            id,
979            name,
980            tool,
981            call_args: ToolArgs {
982                positional: Vec::new(),
983                named,
984            },
985        });
986    }
987    Ok(prepared)
988}
989
990struct Approved {
991    index: usize,
992    id: String,
993    name: String,
994    tool: std::sync::Arc<dyn Tool>,
995    call_args: ToolArgs,
996}
997
998async fn partition_and_gate(
999    prepared: Vec<PreparedEntry>,
1000    ctx: &ToolCtx,
1001) -> (Vec<Approved>, Vec<Approved>, Vec<Option<Value>>) {
1002    let total = prepared.len();
1003    let mut out_slots: Vec<Option<Value>> = vec![None; total];
1004    struct ReadyEntry {
1005        index: usize,
1006        id: String,
1007        name: String,
1008        tool: std::sync::Arc<dyn Tool>,
1009        call_args: ToolArgs,
1010    }
1011    let mut ready: Vec<ReadyEntry> = Vec::new();
1012    for entry in prepared {
1013        match entry {
1014            PreparedEntry::Failed { index, msg } => {
1015                out_slots[index] = Some(Value::Message(emit_tool_result(ctx, &msg)));
1016            }
1017            PreparedEntry::Ready {
1018                index,
1019                id,
1020                name,
1021                tool,
1022                call_args,
1023            } => {
1024                ready.push(ReadyEntry {
1025                    index,
1026                    id,
1027                    name,
1028                    tool,
1029                    call_args,
1030                });
1031            }
1032        }
1033    }
1034    // Parallel: serial awaits hid all but the first pending node from the UI.
1035    let gates = ready.iter().map(|r| {
1036        let level = r.tool.approval_level(&r.call_args, ctx);
1037        request_approval(
1038            ctx,
1039            &r.id,
1040            &r.name,
1041            &r.call_args,
1042            level,
1043            Some(r.tool.as_ref()),
1044        )
1045    });
1046    let outcomes = futures::future::join_all(gates).await;
1047    let mut auto_batch = Vec::new();
1048    let mut serial_batch = Vec::new();
1049    for (r, outcome) in ready.into_iter().zip(outcomes) {
1050        let level = r.tool.approval_level(&r.call_args, ctx);
1051        match outcome {
1052            ApprovalOutcome::Approve => {
1053                let a = Approved {
1054                    index: r.index,
1055                    id: r.id,
1056                    name: r.name.clone(),
1057                    tool: r.tool,
1058                    call_args: r.call_args,
1059                };
1060                if level == crate::tool::ApprovalLevel::Auto {
1061                    auto_batch.push(a);
1062                } else {
1063                    serial_batch.push(a);
1064                }
1065            }
1066            ApprovalOutcome::Deny { reason } => {
1067                let msg = build_error_result(
1068                    ctx,
1069                    &r.id,
1070                    &format!("tool `{}` denied by user: {reason}", r.name),
1071                );
1072                out_slots[r.index] = Some(Value::Message(emit_tool_result(ctx, &msg)));
1073            }
1074        }
1075    }
1076    (auto_batch, serial_batch, out_slots)
1077}
1078
1079async fn run_auto_parallel(batch: Vec<Approved>, ctx: &ToolCtx, out_slots: &mut [Option<Value>]) {
1080    use futures::StreamExt;
1081
1082    let mut pending = futures::stream::FuturesUnordered::new();
1083    for a in batch {
1084        pending.push(async move {
1085            let result = a.tool.call(a.call_args, ctx).await;
1086            (a.index, a.id, a.name, result)
1087        });
1088    }
1089    while let Some((index, id, name, result)) = pending.next().await {
1090        out_slots[index] = Some(finish_dispatch(ctx, &id, &name, result));
1091    }
1092}
1093
1094async fn run_serial(batch: Vec<Approved>, ctx: &ToolCtx, out_slots: &mut [Option<Value>]) {
1095    for a in batch {
1096        let result = a.tool.call(a.call_args, ctx).await;
1097        out_slots[a.index] = Some(finish_dispatch(ctx, &a.id, &a.name, result));
1098    }
1099}
1100
1101fn finish_dispatch(ctx: &ToolCtx, id: &str, name: &str, result: ToolResult) -> Value {
1102    let (content, is_error) = match &result {
1103        Ok(v) => (render_tool_result_text(v), false),
1104        Err(e) => (format!("{e}"), true),
1105    };
1106    if let Ok(v) = &result {
1107        emit_diff_preview_if_relevant(ctx, name, v);
1108    }
1109    let msg = crate::message::Message {
1110        role: crate::message::MessageRole::Tool,
1111        parts: vec![crate::message::MessagePart::ToolResult {
1112            tool_use_id: id.to_string(),
1113            content,
1114            is_error,
1115        }],
1116        turn_id: ctx
1117            .turn_id
1118            .clone()
1119            .unwrap_or_else(crate::event::TurnId::now),
1120        origin: crate::message::MessageOrigin::User,
1121    };
1122    Value::Message(emit_tool_result(ctx, &msg))
1123}
1124
1125type DiffPreviewData = (String, Option<String>, Option<String>, Option<String>);
1126
1127fn emit_diff_preview_if_relevant(ctx: &ToolCtx, tool_name: &str, value: &Value) {
1128    let Some(sink) = ctx.events.as_ref() else {
1129        return;
1130    };
1131    let data: Option<DiffPreviewData> = match tool_name {
1132        "fs.edit" => {
1133            let path = value_struct_string(value, "summary").and_then(|s| {
1134                s.strip_prefix("[fs.edit(")
1135                    .and_then(|s| s.split(':').next())
1136                    .map(|s| s.trim_end_matches(')').to_string())
1137            });
1138            let diff = value_struct_string(value, "diff");
1139            diff.map(|d| (path.unwrap_or_default(), None, None, Some(d)))
1140        }
1141        "fs.write" => {
1142            let path = value_struct_string(value, "path").unwrap_or_default();
1143            let diff = value_struct_string(value, "diff");
1144            diff.map(|d| (path, None, None, Some(d)))
1145        }
1146        "git.diff" => {
1147            let Some(diff) = value_struct_string(value, "diff") else {
1148                return;
1149            };
1150            Some(("git diff".into(), None, None, Some(diff)))
1151        }
1152        "git.show" => {
1153            let sha = value_struct_string(value, "sha").unwrap_or_default();
1154            let Some(diff) = value_struct_string(value, "diff") else {
1155                return;
1156            };
1157            Some((format!("git show {sha}"), None, None, Some(diff)))
1158        }
1159        "git.log" => {
1160            let Some(diff) = value_struct_string(value, "diff") else {
1161                return;
1162            };
1163            Some(("git log HEAD".into(), None, None, Some(diff)))
1164        }
1165        _ => None,
1166    };
1167    if let Some((title, old_content, new_content, unified_diff)) = data {
1168        sink.emit(crate::event::Event::DiffPreview {
1169            turn_id: ctx.turn_id.clone(),
1170            flow_run_id: ctx.flow_run_id.clone(),
1171            title,
1172            old_content,
1173            new_content,
1174            unified_diff,
1175        });
1176    }
1177}
1178
1179fn value_struct_string(value: &Value, field: &str) -> Option<String> {
1180    if let Value::Struct(fields) = value {
1181        fields
1182            .iter()
1183            .find(|(k, _)| k == field)
1184            .and_then(|(_, v)| match v {
1185                Value::Str(s) => Some(s.clone()),
1186                _ => None,
1187            })
1188    } else {
1189        None
1190    }
1191}
1192
1193fn emit_tool_node(ctx: &ToolCtx, id: &str, name: &str, input: &Value) {
1194    let (Some(run_id), Some(parent_node)) = (&ctx.flow_run_id, &ctx.current_node_id) else {
1195        return;
1196    };
1197    let args_preview = format!("{:?}", input)
1198        .chars()
1199        .take(4000)
1200        .collect::<String>();
1201    if let Some(sink) = &ctx.events {
1202        sink.emit(crate::event::Event::ToolNode {
1203            run_id: run_id.clone(),
1204            parent_node_id: parent_node.clone(),
1205            tool_use_id: id.to_string(),
1206            tool_name: name.to_string(),
1207            args_preview: args_preview.clone(),
1208        });
1209    }
1210    if let Some(tx) = &ctx.stream_tx {
1211        let _ = tx.send(crate::stream::StreamFrame::ToolNode {
1212            run_id: run_id.0.to_string(),
1213            parent_node_id: parent_node.clone(),
1214            tool_use_id: id.to_string(),
1215            tool: name.to_string(),
1216            args_preview,
1217        });
1218    }
1219}
1220
1221fn build_error_result(ctx: &ToolCtx, tool_use_id: &str, content: &str) -> crate::message::Message {
1222    crate::message::Message {
1223        role: crate::message::MessageRole::Tool,
1224        parts: vec![crate::message::MessagePart::ToolResult {
1225            tool_use_id: tool_use_id.to_string(),
1226            content: content.to_string(),
1227            is_error: true,
1228        }],
1229        turn_id: ctx
1230            .turn_id
1231            .clone()
1232            .unwrap_or_else(crate::event::TurnId::now),
1233        origin: crate::message::MessageOrigin::User,
1234    }
1235}
1236
1237fn missing_required_fields(schema: &serde_json::Value, named: &[(String, Value)]) -> Vec<String> {
1238    let Some(required) = schema.get("required").and_then(|v| v.as_array()) else {
1239        return Vec::new();
1240    };
1241    let have: std::collections::HashSet<&str> = named.iter().map(|(k, _)| k.as_str()).collect();
1242    required
1243        .iter()
1244        .filter_map(|v| v.as_str())
1245        .filter(|k| !have.contains(k))
1246        .map(String::from)
1247        .collect()
1248}
1249
1250fn emit_tool_result(ctx: &ToolCtx, msg: &crate::message::Message) -> crate::message::Message {
1251    let msg = crate::tools::tool_output::maybe_truncate_tool_message_with_budget(
1252        msg,
1253        ctx.output_store.as_deref(),
1254        ctx.tool_output_budget,
1255    );
1256    if let Some(tx) = &ctx.stream_tx {
1257        let _ = tx.send(crate::stream::StreamFrame::ToolResultMsg {
1258            flow_run_id: ctx.flow_run_id.as_ref().map(|r| r.0.to_string()),
1259            message: msg.clone(),
1260        });
1261    } else if let Some(sink) = &ctx.events {
1262        sink.emit(crate::event::Event::ToolResultMsg {
1263            turn_id: msg.turn_id.clone(),
1264            flow_run_id: ctx.flow_run_id.clone(),
1265            message: msg.clone(),
1266        });
1267    }
1268    msg
1269}
1270
1271fn render_tool_result_text(v: &Value) -> String {
1272    match v {
1273        Value::Str(s) => s.clone(),
1274        Value::Message(m) => m.text_concat(),
1275        other => other.to_json().to_string(),
1276    }
1277}
1278
1279fn extract_list(args: &ToolArgs, name: &str, pos: usize) -> Result<Vec<Value>, RuntimeError> {
1280    let value = match args.named(name) {
1281        Some(v) => v,
1282        None => args.positional(pos)?,
1283    };
1284    match value {
1285        Value::List(items) => Ok(items.clone()),
1286        other => Err(RuntimeError::TypeMismatch {
1287            expected: "list".into(),
1288            actual: other.kind_name().into(),
1289        }),
1290    }
1291}
1292
1293pub struct ComposeEmailPreview;
1294
1295impl Tool for ComposeEmailPreview {
1296    fn name(&self) -> &str {
1297        "compose_email_preview"
1298    }
1299
1300    fn tier(&self) -> Tier {
1301        Tier::Zero
1302    }
1303
1304    fn call<'a>(&'a self, args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
1305        Box::pin(async move {
1306            let subject = extract_string(&args, "subject", 0)?;
1307            let body = extract_string(&args, "body", 1)?;
1308            let to = extract_string_list(&args, "to", 2)?;
1309            Ok(Value::Str(compose_email_preview(&subject, &body, &to)))
1310        })
1311    }
1312}
1313
1314pub fn compose_email_preview(subject: &str, body: &str, to: &[String]) -> String {
1315    format!(
1316        "To: {}
1317Subject: {subject}
1318---
1319{body}",
1320        to.join(", ")
1321    )
1322}
1323
1324fn extract_string(args: &ToolArgs, name: &str, pos: usize) -> Result<String, RuntimeError> {
1325    let value = match args.named(name) {
1326        Some(v) => v,
1327        None => args.positional(pos)?,
1328    };
1329    match value {
1330        Value::Str(s) => Ok(s.clone()),
1331        other => Err(RuntimeError::TypeMismatch {
1332            expected: "string".into(),
1333            actual: other.kind_name().into(),
1334        }),
1335    }
1336}
1337
1338fn extract_string_list(
1339    args: &ToolArgs,
1340    name: &str,
1341    pos: usize,
1342) -> Result<Vec<String>, 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) => items
1349            .iter()
1350            .map(|v| match v {
1351                Value::Str(s) => Ok(s.clone()),
1352                other => Err(RuntimeError::TypeMismatch {
1353                    expected: "list of string".into(),
1354                    actual: other.kind_name().into(),
1355                }),
1356            })
1357            .collect(),
1358        other => Err(RuntimeError::TypeMismatch {
1359            expected: "list".into(),
1360            actual: other.kind_name().into(),
1361        }),
1362    }
1363}
1364
1365#[cfg(test)]
1366mod tests {
1367    use super::*;
1368
1369    #[test]
1370    fn shell_quote_wraps_and_escapes() {
1371        assert_eq!(shell_quote("hello"), "'hello'");
1372        assert_eq!(shell_quote("It's fine"), "'It'\\''s fine'");
1373        assert_eq!(shell_quote(""), "''");
1374        assert_eq!(shell_quote("a'b'c"), "'a'\\''b'\\''c'");
1375    }
1376
1377    struct ControlledTool {
1378        name: &'static str,
1379        release: std::sync::Arc<tokio::sync::Semaphore>,
1380    }
1381
1382    impl Tool for ControlledTool {
1383        fn name(&self) -> &str {
1384            self.name
1385        }
1386
1387        fn tier(&self) -> Tier {
1388            Tier::Zero
1389        }
1390
1391        fn call<'a>(&'a self, _args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
1392            Box::pin(async move {
1393                let _permit = self.release.acquire().await.unwrap();
1394                Ok(Value::Str(self.name.to_string()))
1395            })
1396        }
1397    }
1398
1399    #[tokio::test]
1400    async fn dispatch_all_emits_each_scoped_result_as_its_tool_finishes() {
1401        let fast_release = std::sync::Arc::new(tokio::sync::Semaphore::new(0));
1402        let slow_release = std::sync::Arc::new(tokio::sync::Semaphore::new(0));
1403        let registry = crate::tool::ToolRegistry::new();
1404        registry.register(std::sync::Arc::new(ControlledTool {
1405            name: "fast",
1406            release: fast_release.clone(),
1407        }));
1408        registry.register(std::sync::Arc::new(ControlledTool {
1409            name: "slow",
1410            release: slow_release.clone(),
1411        }));
1412        let run_id = crate::event::FlowRunId::now();
1413        let (stream_tx, mut stream_rx) = tokio::sync::broadcast::channel(32);
1414        let ctx = ToolCtx::new()
1415            .with_anchors(None, Some(run_id.clone()), None)
1416            .with_current_node(Some("dispatch_all".into()))
1417            .with_registry(std::sync::Arc::new(registry))
1418            .with_stream_tx(stream_tx);
1419        let uses = Value::List(vec![
1420            Value::Struct(vec![
1421                ("id".into(), Value::Str("slow_id".into())),
1422                ("name".into(), Value::Str("slow".into())),
1423                ("input".into(), Value::Struct(Vec::new())),
1424            ]),
1425            Value::Struct(vec![
1426                ("id".into(), Value::Str("fast_id".into())),
1427                ("name".into(), Value::Str("fast".into())),
1428                ("input".into(), Value::Struct(Vec::new())),
1429            ]),
1430        ]);
1431        let task = tokio::spawn(async move {
1432            DispatchAll
1433                .call(
1434                    ToolArgs {
1435                        positional: vec![uses],
1436                        named: Vec::new(),
1437                    },
1438                    &ctx,
1439                )
1440                .await
1441                .unwrap()
1442        });
1443
1444        fast_release.add_permits(1);
1445        let fast_result = tokio::time::timeout(std::time::Duration::from_secs(1), async {
1446            loop {
1447                if let crate::stream::StreamFrame::ToolResultMsg {
1448                    flow_run_id,
1449                    message,
1450                } = stream_rx.recv().await.unwrap()
1451                    && message.parts.iter().any(|part| {
1452                        matches!(
1453                            part,
1454                            crate::message::MessagePart::ToolResult { tool_use_id, .. }
1455                                if tool_use_id == "fast_id"
1456                        )
1457                    })
1458                {
1459                    break flow_run_id;
1460                }
1461            }
1462        })
1463        .await
1464        .expect("fast result before slow release");
1465        assert_eq!(fast_result.as_deref(), Some(run_id.0.to_string().as_str()));
1466        assert!(!task.is_finished());
1467
1468        slow_release.add_permits(1);
1469        let Value::List(results) = task.await.unwrap() else {
1470            panic!("dispatch result list");
1471        };
1472        let ids: Vec<&str> = results
1473            .iter()
1474            .map(|value| match value {
1475                Value::Message(message) => match &message.parts[0] {
1476                    crate::message::MessagePart::ToolResult { tool_use_id, .. } => {
1477                        tool_use_id.as_str()
1478                    }
1479                    _ => panic!("tool result part"),
1480                },
1481                _ => panic!("tool result message"),
1482            })
1483            .collect();
1484        assert_eq!(ids, vec!["slow_id", "fast_id"]);
1485    }
1486
1487    struct TextTool {
1488        name: &'static str,
1489        output: String,
1490    }
1491
1492    impl Tool for TextTool {
1493        fn name(&self) -> &str {
1494            self.name
1495        }
1496
1497        fn tier(&self) -> Tier {
1498            Tier::Zero
1499        }
1500
1501        fn call<'a>(&'a self, _args: ToolArgs, _ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
1502            let output = self.output.clone();
1503            Box::pin(async move { Ok(Value::Str(output)) })
1504        }
1505    }
1506
1507    #[tokio::test]
1508    async fn dispatch_all_preserves_fs_read_pagination_for_full_utf8_reassembly() {
1509        let dir = tempfile::tempdir().unwrap();
1510        let path = dir.path().join("large.txt");
1511        let expected = "界".repeat(349_525) + "a";
1512        assert_eq!(expected.len(), 1_048_576);
1513        tokio::fs::write(&path, &expected).await.unwrap();
1514
1515        let registry = crate::tool::ToolRegistry::new();
1516        crate::tools::register_tier_zero(&registry);
1517        let ctx = ToolCtx::new()
1518            .with_registry(std::sync::Arc::new(registry))
1519            .with_session_dir(dir.path().to_path_buf());
1520        let uses = Value::List(vec![Value::Struct(vec![
1521            ("id".into(), Value::Str("read_id".into())),
1522            ("name".into(), Value::Str("fs.read".into())),
1523            (
1524                "input".into(),
1525                Value::Struct(vec![("path".into(), Value::Path(path))]),
1526            ),
1527        ])]);
1528        let Value::List(results) = DispatchAll
1529            .call(
1530                ToolArgs {
1531                    positional: vec![uses],
1532                    named: Vec::new(),
1533                },
1534                &ctx,
1535            )
1536            .await
1537            .unwrap()
1538        else {
1539            panic!("dispatch result list");
1540        };
1541        let Value::Message(message) = &results[0] else {
1542            panic!("tool result message");
1543        };
1544        let crate::message::MessagePart::ToolResult { content, .. } = &message.parts[0] else {
1545            panic!("tool result part");
1546        };
1547        let envelope: serde_json::Value = serde_json::from_str(content).unwrap();
1548        let output_id = envelope["output_id"].as_str().unwrap();
1549        let mut reconstructed = envelope["content"].as_str().unwrap().to_string();
1550        let mut offset = envelope["next"]["offset"].as_u64().unwrap() as usize;
1551        while envelope["next"]["has_more"].as_bool().unwrap() || offset < expected.len() {
1552            let page = ctx
1553                .output_store
1554                .as_ref()
1555                .unwrap()
1556                .read_bytes(output_id, offset, usize::MAX, ctx.tool_output_budget)
1557                .unwrap();
1558            reconstructed.push_str(&page.content);
1559            if !page.has_more {
1560                break;
1561            }
1562            assert!(page.next_offset > offset);
1563            offset = page.next_offset;
1564        }
1565        assert_eq!(reconstructed.len(), 1_048_576);
1566        assert_eq!(reconstructed, expected);
1567    }
1568
1569    #[tokio::test]
1570    async fn dispatch_all_returns_the_same_truncated_message_it_emits() {
1571        let registry = crate::tool::ToolRegistry::new();
1572        registry.register(std::sync::Arc::new(TextTool {
1573            name: "text",
1574            output: "0123456789".repeat(20),
1575        }));
1576        let (stream_tx, mut stream_rx) = tokio::sync::broadcast::channel(8);
1577        let mut ctx = ToolCtx::new()
1578            .with_registry(std::sync::Arc::new(registry))
1579            .with_stream_tx(stream_tx);
1580        ctx.tool_output_budget = crate::tools::tool_output::ToolOutputBudget {
1581            max_lines: 32,
1582            max_bytes: 24,
1583            max_line_bytes: 24,
1584        };
1585        let uses = Value::List(vec![Value::Struct(vec![
1586            ("id".into(), Value::Str("text_id".into())),
1587            ("name".into(), Value::Str("text".into())),
1588            ("input".into(), Value::Struct(Vec::new())),
1589        ])]);
1590        let Value::List(results) = DispatchAll
1591            .call(
1592                ToolArgs {
1593                    positional: vec![uses],
1594                    named: Vec::new(),
1595                },
1596                &ctx,
1597            )
1598            .await
1599            .unwrap()
1600        else {
1601            panic!("dispatch result list");
1602        };
1603        let Value::Message(returned) = &results[0] else {
1604            panic!("tool result message");
1605        };
1606        let crate::stream::StreamFrame::ToolResultMsg {
1607            message: emitted, ..
1608        } = stream_rx.try_recv().unwrap()
1609        else {
1610            panic!("tool result stream frame");
1611        };
1612        assert_eq!(returned, &emitted);
1613        assert!(matches!(
1614            &returned.parts[0],
1615            crate::message::MessagePart::ToolResult { content, .. }
1616                if content.contains("Output truncated")
1617        ));
1618    }
1619
1620    #[test]
1621    fn finish_dispatch_returns_the_budgeted_message_it_emits() {
1622        let (stream_tx, mut stream_rx) = tokio::sync::broadcast::channel(8);
1623        let mut ctx = ToolCtx::new().with_stream_tx(stream_tx);
1624        ctx.tool_output_budget = crate::tools::tool_output::ToolOutputBudget {
1625            max_lines: 32,
1626            max_bytes: 24,
1627            max_line_bytes: 24,
1628        };
1629        let Value::Message(returned) = finish_dispatch(
1630            &ctx,
1631            "text_id",
1632            "text",
1633            Ok(Value::Str("0123456789".repeat(20))),
1634        ) else {
1635            panic!("tool result message");
1636        };
1637        let crate::stream::StreamFrame::ToolResultMsg {
1638            message: emitted, ..
1639        } = stream_rx.try_recv().unwrap()
1640        else {
1641            panic!("tool result stream frame");
1642        };
1643
1644        assert_eq!(returned, emitted);
1645        assert!(matches!(
1646            &returned.parts[0],
1647            crate::message::MessagePart::ToolResult { content, .. }
1648                if content.contains("Output truncated")
1649        ));
1650    }
1651
1652    #[test]
1653    fn finish_dispatch_preserves_error_flag_and_message_consistency() {
1654        let (stream_tx, mut stream_rx) = tokio::sync::broadcast::channel(8);
1655        let ctx = ToolCtx::new().with_stream_tx(stream_tx);
1656        let Value::Message(returned) = finish_dispatch(
1657            &ctx,
1658            "error_id",
1659            "failing",
1660            Err(RuntimeError::ToolFailed("expected failure".into())),
1661        ) else {
1662            panic!("tool result message");
1663        };
1664        let crate::stream::StreamFrame::ToolResultMsg {
1665            message: emitted, ..
1666        } = stream_rx.try_recv().unwrap()
1667        else {
1668            panic!("tool result stream frame");
1669        };
1670
1671        assert_eq!(returned, emitted);
1672        assert!(matches!(
1673            &returned.parts[0],
1674            crate::message::MessagePart::ToolResult {
1675                content,
1676                is_error: true,
1677                ..
1678            } if content.contains("expected failure")
1679        ));
1680    }
1681
1682    #[test]
1683    fn finish_dispatch_emits_full_diff_preview_before_result_truncation() {
1684        let events = crate::event::EventSink::new();
1685        let mut ctx = ToolCtx::new().with_events(events.clone());
1686        ctx.tool_output_budget = crate::tools::tool_output::ToolOutputBudget {
1687            max_lines: 1,
1688            max_bytes: 16,
1689            max_line_bytes: 16,
1690        };
1691        let diff = "-old\n+new\n".repeat(20);
1692        let Value::Message(returned) = finish_dispatch(
1693            &ctx,
1694            "edit_id",
1695            "fs.edit",
1696            Ok(Value::Struct(vec![
1697                (
1698                    "summary".into(),
1699                    Value::Str("[fs.edit(example.txt): updated]".into()),
1700                ),
1701                ("diff".into(), Value::Str(diff.clone())),
1702            ])),
1703        ) else {
1704            panic!("tool result message");
1705        };
1706
1707        assert!(matches!(
1708            &returned.parts[0],
1709            crate::message::MessagePart::ToolResult { content, .. }
1710                if content.contains("Output truncated")
1711        ));
1712        assert!(events.snapshot().iter().any(|event| matches!(
1713            event,
1714            crate::event::Event::DiffPreview {
1715                unified_diff: Some(preview),
1716                ..
1717            } if preview == &diff
1718        )));
1719    }
1720
1721    #[tokio::test]
1722    async fn dispatch_all_returns_the_same_truncated_preflight_error_it_emits() {
1723        let registry = crate::tool::ToolRegistry::new();
1724        let (stream_tx, mut stream_rx) = tokio::sync::broadcast::channel(8);
1725        let mut ctx = ToolCtx::new()
1726            .with_registry(std::sync::Arc::new(registry))
1727            .with_stream_tx(stream_tx);
1728        ctx.tool_output_budget = crate::tools::tool_output::ToolOutputBudget {
1729            max_lines: 32,
1730            max_bytes: 32,
1731            max_line_bytes: 32,
1732        };
1733        let uses = Value::List(vec![Value::Struct(vec![
1734            ("id".into(), Value::Str("unknown_id".into())),
1735            ("name".into(), Value::Str("missing".repeat(40))),
1736            ("input".into(), Value::Struct(Vec::new())),
1737        ])]);
1738
1739        let Value::List(results) = DispatchAll
1740            .call(
1741                ToolArgs {
1742                    positional: vec![uses],
1743                    named: Vec::new(),
1744                },
1745                &ctx,
1746            )
1747            .await
1748            .unwrap()
1749        else {
1750            panic!("dispatch result list");
1751        };
1752        let Value::Message(returned) = &results[0] else {
1753            panic!("tool result message");
1754        };
1755        let crate::stream::StreamFrame::ToolResultMsg {
1756            message: emitted, ..
1757        } = stream_rx.try_recv().unwrap()
1758        else {
1759            panic!("tool result stream frame");
1760        };
1761
1762        assert_eq!(returned, &emitted);
1763        assert!(matches!(
1764            &returned.parts[0],
1765            crate::message::MessagePart::ToolResult {
1766                content,
1767                is_error: true,
1768                ..
1769            } if content.contains("Output truncated")
1770        ));
1771    }
1772
1773    #[tokio::test]
1774    async fn dispatch_all_unknown_tool_finishes_its_workflow_node_with_error() {
1775        use crate::workflow::{NodeStatus, WorkflowGraph};
1776
1777        let registry = crate::tool::ToolRegistry::new();
1778        let run_id = crate::event::FlowRunId::now();
1779        let run = run_id.0.to_string();
1780        let (stream_tx, mut stream_rx) = tokio::sync::broadcast::channel(8);
1781        let ctx = ToolCtx::new()
1782            .with_anchors(None, Some(run_id), None)
1783            .with_current_node(Some("dispatch_all".into()))
1784            .with_registry(std::sync::Arc::new(registry))
1785            .with_stream_tx(stream_tx);
1786        let uses = Value::List(vec![Value::Struct(vec![
1787            ("id".into(), Value::Str("unknown_id".into())),
1788            ("name".into(), Value::Str("missing.tool".into())),
1789            ("input".into(), Value::Struct(Vec::new())),
1790        ])]);
1791
1792        DispatchAll
1793            .call(
1794                ToolArgs {
1795                    positional: vec![uses],
1796                    named: Vec::new(),
1797                },
1798                &ctx,
1799            )
1800            .await
1801            .unwrap();
1802
1803        let mut graph = WorkflowGraph::new(crate::event::TurnId::now());
1804        graph.apply_stream_frame(&crate::stream::StreamFrame::FlowStart {
1805            run_id: run.clone(),
1806            flow_name: "agent_loop".into(),
1807            parent_run_id: None,
1808            parent_node_id: None,
1809        });
1810        graph.apply_stream_frame(&crate::stream::StreamFrame::FlowNodeStart {
1811            run_id: run.clone(),
1812            node_id: "dispatch_all".into(),
1813            kind: crate::nodegraph::NodeKind::ToolCall {
1814                path: "dispatch_all".into(),
1815            },
1816            label: "dispatch_all".into(),
1817            parent_node_id: None,
1818        });
1819        while let Ok(frame) = stream_rx.try_recv() {
1820            graph.apply_stream_frame(&frame);
1821        }
1822
1823        let node = graph
1824            .find_node(&format!("tool:{run}:unknown_id"))
1825            .expect("unknown tool node");
1826        assert_eq!(node.status, NodeStatus::Err);
1827        assert!(matches!(
1828            &node.kind,
1829            crate::workflow::WorkflowNodeKind::ToolCall {
1830                result_preview: Some(preview),
1831                ..
1832            } if preview.contains("unknown tool")
1833        ));
1834        let result = node.output_preview.as_deref().unwrap();
1835        assert!(result.contains("unknown tool"));
1836        assert!(
1837            graph
1838                .root
1839                .iter()
1840                .flat_map(|root| &root.children)
1841                .all(|child| {
1842                    !matches!(
1843                        &child.kind,
1844                        crate::workflow::WorkflowNodeKind::ToolCall { tool_use_id, .. }
1845                            if tool_use_id == "unknown_id"
1846                    )
1847                })
1848        );
1849    }
1850
1851    #[test]
1852    fn prepare_dispatch_does_not_emit_partial_nodes_for_malformed_batch() {
1853        let registry = crate::tool::ToolRegistry::new();
1854        let run_id = crate::event::FlowRunId::now();
1855        let (stream_tx, mut stream_rx) = tokio::sync::broadcast::channel(8);
1856        let ctx = ToolCtx::new()
1857            .with_anchors(None, Some(run_id), None)
1858            .with_current_node(Some("dispatch_all".into()))
1859            .with_stream_tx(stream_tx);
1860        let uses = vec![
1861            Value::Struct(vec![
1862                ("id".into(), Value::Str("valid_id".into())),
1863                ("name".into(), Value::Str("missing.tool".into())),
1864            ]),
1865            Value::Struct(vec![("name".into(), Value::Str("missing.tool".into()))]),
1866        ];
1867
1868        assert!(prepare_dispatch(&uses, &registry, &ctx).is_err());
1869        assert!(matches!(
1870            stream_rx.try_recv(),
1871            Err(tokio::sync::broadcast::error::TryRecvError::Empty)
1872        ));
1873    }
1874
1875    #[test]
1876    fn compose_email_preview_formats_headers() {
1877        let preview = compose_email_preview(
1878            "Deploy status",
1879            "See attached",
1880            &["a@x.com".into(), "b@x.com".into()],
1881        );
1882        assert_eq!(
1883            preview,
1884            "To: a@x.com, b@x.com\nSubject: Deploy status\n---\nSee attached"
1885        );
1886    }
1887}