Skip to main content

atman_runtime/tools/
agent_ctrl.rs

1use crate::approval::{ApprovalOutcome, request_approval};
2use crate::error::RuntimeError;
3use crate::event::{Event, FlowRunId, FlowStatus, NodeEvent};
4use crate::message::{Message, MessagePart, MessageRole};
5use crate::provider::LlmRequest;
6use crate::tool::{ApprovalLevel, BoxFut, Tier, Tool, ToolArgs, ToolCtx, ToolResult, ToolSpec};
7use crate::value::Value;
8use std::path::PathBuf;
9
10pub struct AgentSpawn;
11
12const DEFAULT_MAX_ITER: u64 = 20;
13const MAX_ITER_HARD_CAP: u64 = 200;
14
15impl Tool for AgentSpawn {
16    fn name(&self) -> &str {
17        "agent.spawn"
18    }
19
20    fn tier(&self) -> Tier {
21        Tier::Two
22    }
23
24    fn approval_level(&self, _args: &ToolArgs, _ctx: &ToolCtx) -> ApprovalLevel {
25        ApprovalLevel::Approve
26    }
27
28    fn description(&self) -> Option<&str> {
29        Some(
30            "Spawn an independent sub-agent to handle a focused sub-goal. The sub-agent runs its \
31             own message history and iteration counter, uses the same tool registry (or a subset \
32             you pick), and returns its final assistant text as this tool's result. Prefer this \
33             over doing large exploratory work directly when it would otherwise flood the main \
34             conversation with search output or scratch reasoning. Parameters: \
35             `goal` (required string), `tools` (optional list of tool-name strings — defaults \
36             to all tools available to you), `max_iterations` (optional int, default 20, capped \
37             at 200), `model` (optional model name — defaults to the last model this session \
38             used, then configured models, then claude-opus-4.7), `flow` (optional .at file path \
39             or command name; goal is passed as the first flow argument).",
40        )
41    }
42
43    fn input_schema(&self) -> serde_json::Value {
44        serde_json::json!({
45            "type": "object",
46            "properties": {
47                "goal": {"type": "string"},
48                "tools": {"type": "array", "items": {"type": "string"}},
49                "max_iterations": {"type": "integer"},
50                "model": {"type": "string"},
51                "flow": {"type": "string"}
52            },
53            "required": ["goal"]
54        })
55    }
56
57    fn call<'a>(&'a self, args: ToolArgs, ctx: &'a ToolCtx) -> BoxFut<'a, ToolResult> {
58        Box::pin(async move { run_sub_agent(args, ctx).await })
59    }
60}
61
62async fn run_sub_agent(args: ToolArgs, ctx: &ToolCtx) -> ToolResult {
63    let goal = extract_goal(&args)?;
64    if let Some(flow) = extract_flow(&args)? {
65        return run_flow_agent(&flow, goal, ctx).await;
66    }
67    let max_iter = extract_max_iter(&args);
68    let tool_filter = extract_tool_filter(&args)?;
69    let model = pick_model(&args, ctx);
70    let Some(providers) = ctx.providers.as_ref() else {
71        return Err(RuntimeError::ToolFailed(
72            "agent.spawn: no provider registry available on ctx".into(),
73        ));
74    };
75    let Some(provider) = providers.resolve(&model) else {
76        return Ok(Value::Str(format!(
77            "[sub-agent failed: no provider for model `{model}`]"
78        )));
79    };
80    let Some(registry) = ctx.registry.as_ref() else {
81        return Err(RuntimeError::ToolFailed(
82            "agent.spawn: no tool registry available on ctx".into(),
83        ));
84    };
85    let tool_specs = build_tool_specs(registry.as_ref(), tool_filter.as_deref());
86    let child_run_id = FlowRunId::now();
87    emit_child_flow_start(ctx, &child_run_id, &goal);
88    let turn = ctx
89        .turn_id
90        .clone()
91        .unwrap_or_else(crate::event::TurnId::now);
92    let mut messages: Vec<Message> = vec![Message::user_text(turn.clone(), goal.clone())];
93    let mut final_text: Option<String> = None;
94    let mut failure_reason: Option<String> = None;
95    for iter in 0..max_iter {
96        if ctx.cancel.is_cancelled() {
97            failure_reason = Some("cancelled by parent".into());
98            break;
99        }
100        let req = LlmRequest {
101            model: model.clone(),
102            messages: messages.clone(),
103            system: None,
104            input: Value::Unit,
105            schema: None,
106            cache_prompt: false,
107            tools: tool_specs.clone(),
108            thinking_enabled: false,
109            stall_timeout_secs: 120,
110        };
111        let outcome = call_streaming_sub_agent(provider.as_ref(), req, ctx).await;
112        match outcome {
113            Ok(am) => {
114                emit_child_llm_call(ctx, &child_run_id, &model, &am);
115                let uses = extract_tool_uses(&am.message);
116                messages.push(am.message.clone());
117                emit_assistant_msg(ctx, &child_run_id, &am.message);
118                if uses.is_empty() {
119                    final_text = Some(am.text_concat());
120                    break;
121                }
122                let tool_results = dispatch_child_tools(&uses, registry.as_ref(), ctx).await;
123                let turn_for_results = am.message.turn_id.clone();
124                let combined = Message {
125                    turn_id: turn_for_results,
126                    role: MessageRole::Tool,
127                    parts: tool_results,
128                };
129                emit_tool_result_msg(ctx, &child_run_id, &combined);
130                messages.push(combined);
131            }
132            Err(e) => {
133                failure_reason = Some(format!("provider error at iter {iter}: {e}"));
134                break;
135            }
136        }
137    }
138    let status = if final_text.is_some() {
139        FlowStatus::Ok
140    } else {
141        FlowStatus::Errored {
142            message: failure_reason
143                .clone()
144                .unwrap_or_else(|| format!("hit max iterations {max_iter} without a final answer")),
145        }
146    };
147    emit_child_flow_end(ctx, &child_run_id, &status);
148    if let Some(text) = final_text {
149        Ok(Value::Str(text))
150    } else {
151        let reason = failure_reason
152            .unwrap_or_else(|| format!("hit max iterations {max_iter} without a final answer"));
153        let last = messages
154            .iter()
155            .rev()
156            .find(|m| matches!(m.role, MessageRole::Assistant))
157            .map(|m| m.text_concat())
158            .unwrap_or_default();
159        let partial = if last.is_empty() {
160            String::new()
161        } else {
162            format!("\n[partial output: {}]", truncate(&last, 400))
163        };
164        Ok(Value::Str(format!("[sub-agent failed: {reason}]{partial}")))
165    }
166}
167
168async fn run_flow_agent(flow_ref: &str, goal: String, ctx: &ToolCtx) -> ToolResult {
169    let Some(registry) = ctx.registry.as_ref() else {
170        return Err(RuntimeError::ToolFailed(
171            "agent.spawn: no tool registry available on ctx".into(),
172        ));
173    };
174    let Some(providers) = ctx.providers.as_ref() else {
175        return Err(RuntimeError::ToolFailed(
176            "agent.spawn: no provider registry available on ctx".into(),
177        ));
178    };
179    let (path, src) = read_flow_source(flow_ref).await?;
180    let file = atman_dsl::parse::parse_file(&src).map_err(|e| {
181        RuntimeError::ToolFailed(format!("agent.spawn: parse {}: {e}", path.display()))
182    })?;
183    let Some(flow) = file.flows.first() else {
184        return Err(RuntimeError::ToolFailed(format!(
185            "agent.spawn: no flow in {}",
186            path.display()
187        )));
188    };
189    let args = flow
190        .params
191        .first()
192        .map(|(ident, _)| vec![(ident.name.clone(), Value::Str(goal))])
193        .unwrap_or_default();
194    let flows = file
195        .flows
196        .iter()
197        .map(|flow| (flow.name.name.clone(), flow.clone()))
198        .collect();
199    let run_id = FlowRunId::now();
200    emit_flow_agent_start(ctx, &run_id, &flow.name.name);
201    let out = crate::exec::exec_flow_with_siblings(
202        flow,
203        args,
204        registry.as_ref(),
205        ctx,
206        providers.as_ref(),
207        &flows,
208        ctx.events.as_ref(),
209        ctx.turn_id.clone(),
210        Some(run_id.clone()),
211        None,
212        ctx.cancel.clone(),
213        None,
214    )
215    .await;
216    let status = match &out {
217        Ok(_) => FlowStatus::Ok,
218        Err(e) => FlowStatus::Errored {
219            message: e.to_string(),
220        },
221    };
222    emit_child_flow_end(ctx, &run_id, &status);
223    out
224}
225
226fn extract_goal(args: &ToolArgs) -> Result<String, RuntimeError> {
227    match args.named("goal").or_else(|| args.positional.first()) {
228        Some(Value::Str(s)) if !s.trim().is_empty() => Ok(s.clone()),
229        Some(other) => Err(RuntimeError::TypeMismatch {
230            expected: "non-empty goal string".into(),
231            actual: other.kind_name().into(),
232        }),
233        None => Err(RuntimeError::MissingArg("agent.spawn.goal".into())),
234    }
235}
236
237fn extract_max_iter(args: &ToolArgs) -> u64 {
238    match args.named("max_iterations") {
239        Some(Value::Int(n)) if *n > 0 => (*n as u64).min(MAX_ITER_HARD_CAP),
240        _ => DEFAULT_MAX_ITER,
241    }
242}
243
244fn extract_flow(args: &ToolArgs) -> Result<Option<String>, RuntimeError> {
245    match args.named("flow") {
246        Some(Value::Str(s)) if !s.trim().is_empty() => Ok(Some(s.clone())),
247        Some(Value::Unit) | None => Ok(None),
248        Some(other) => Err(RuntimeError::TypeMismatch {
249            expected: "flow string".into(),
250            actual: other.kind_name().into(),
251        }),
252    }
253}
254
255fn extract_tool_filter(args: &ToolArgs) -> Result<Option<Vec<String>>, RuntimeError> {
256    match args.named("tools") {
257        Some(Value::List(items)) => {
258            let mut out = Vec::with_capacity(items.len());
259            for it in items {
260                match it {
261                    Value::Str(s) => out.push(s.clone()),
262                    other => {
263                        return Err(RuntimeError::TypeMismatch {
264                            expected: "string tool name".into(),
265                            actual: other.kind_name().into(),
266                        });
267                    }
268                }
269            }
270            Ok(Some(out))
271        }
272        Some(Value::Unit) | None => Ok(None),
273        Some(other) => Err(RuntimeError::TypeMismatch {
274            expected: "list of tool names".into(),
275            actual: other.kind_name().into(),
276        }),
277    }
278}
279
280fn pick_model(args: &ToolArgs, ctx: &ToolCtx) -> String {
281    if let Some(Value::Str(s)) = args.named("model")
282        && !s.is_empty()
283    {
284        return crate::model_registry::resolve_alias(s);
285    }
286    if let Some(model) = &ctx.current_model {
287        return model.clone();
288    }
289    if let Some((name, _)) = crate::model_registry::all_model_entries().first() {
290        return name.clone();
291    }
292    "claude-opus-4.7".into()
293}
294
295async fn read_flow_source(flow_ref: &str) -> Result<(PathBuf, String), RuntimeError> {
296    for path in flow_candidates(flow_ref) {
297        match tokio::fs::read_to_string(&path).await {
298            Ok(src) => return Ok((path, src)),
299            Err(e) if e.kind() == std::io::ErrorKind::NotFound => {}
300            Err(e) => {
301                return Err(RuntimeError::ToolFailed(format!(
302                    "agent.spawn: read {}: {e}",
303                    path.display()
304                )));
305            }
306        }
307    }
308    Err(RuntimeError::ToolFailed(format!(
309        "agent.spawn: flow `{flow_ref}` not found"
310    )))
311}
312
313fn flow_candidates(flow_ref: &str) -> Vec<PathBuf> {
314    let path = PathBuf::from(flow_ref);
315    if path.is_absolute() {
316        return vec![path];
317    }
318    let file_name = if flow_ref.ends_with(".at") {
319        flow_ref.to_string()
320    } else {
321        format!("{flow_ref}.at")
322    };
323    let mut out = Vec::new();
324    if let Some(home) = std::env::var_os("HOME") {
325        out.push(
326            PathBuf::from(home)
327                .join(".config")
328                .join("atman")
329                .join("commands")
330                .join(&file_name),
331        );
332    }
333    out.push(PathBuf::from(file_name));
334    out
335}
336
337async fn call_streaming_sub_agent(
338    provider: &dyn crate::provider::Provider,
339    req: LlmRequest,
340    ctx: &ToolCtx,
341) -> Result<crate::provider::AssistantMessage, RuntimeError> {
342    let model_name = req.model.clone();
343    let obs = provider.call_streaming(req);
344    let mut events = obs.events;
345    let output = obs.output;
346    tokio::pin!(output);
347    let result = loop {
348        tokio::select! {
349            biased;
350            ev = events.recv() => forward_stream_event(ev, ctx, &model_name),
351            result = &mut output => break result,
352        }
353    };
354    while let Ok(ev) = events.try_recv() {
355        forward_stream_event(Ok(ev), ctx, &model_name);
356    }
357    result
358}
359
360fn forward_stream_event(
361    ev: Result<NodeEvent, tokio::sync::broadcast::error::RecvError>,
362    ctx: &ToolCtx,
363    model: &str,
364) {
365    let Some(tx) = &ctx.stream_tx else {
366        return;
367    };
368    match ev {
369        Ok(NodeEvent::LlmChunk { text, .. }) => {
370            let _ = tx.send(crate::stream::StreamFrame::LlmChunk {
371                text,
372                model: model.to_string(),
373            });
374        }
375        Ok(NodeEvent::ThinkingChunk { text }) => {
376            let _ = tx.send(crate::stream::StreamFrame::ThinkingChunk { text });
377        }
378        Ok(NodeEvent::LlmDone { total_tokens }) => {
379            let _ = tx.send(crate::stream::StreamFrame::LlmDone { total_tokens });
380        }
381        _ => {}
382    }
383}
384
385fn build_tool_specs(
386    registry: &crate::tool::ToolRegistry,
387    filter: Option<&[String]>,
388) -> Vec<ToolSpec> {
389    let mut specs = Vec::new();
390    for (name, tool) in registry.iter() {
391        if let Some(allow) = filter
392            && !allow.iter().any(|n| n == name)
393        {
394            continue;
395        }
396        specs.push(crate::tool::tool_spec(tool.as_ref()));
397    }
398    specs
399}
400
401fn extract_tool_uses(msg: &Message) -> Vec<(String, String, Value)> {
402    let mut out = Vec::new();
403    for part in &msg.parts {
404        if let MessagePart::ToolUse { id, name, input } = part {
405            let value = Value::from_json(input.clone());
406            out.push((id.clone(), name.clone(), value));
407        }
408    }
409    out
410}
411
412async fn dispatch_child_tools(
413    uses: &[(String, String, Value)],
414    registry: &crate::tool::ToolRegistry,
415    ctx: &ToolCtx,
416) -> Vec<MessagePart> {
417    struct Ready {
418        idx: usize,
419        id: String,
420        name: String,
421        tool: std::sync::Arc<dyn crate::tool::Tool>,
422        call_args: ToolArgs,
423    }
424    let mut out: Vec<Option<MessagePart>> = vec![None; uses.len()];
425    let mut ready: Vec<Ready> = Vec::new();
426    for (idx, (id, name, input)) in uses.iter().enumerate() {
427        let Some(tool) = registry.get(name) else {
428            out[idx] = Some(MessagePart::ToolResult {
429                tool_use_id: id.clone(),
430                content: format!("sub-agent: unknown tool `{name}`"),
431                is_error: true,
432            });
433            continue;
434        };
435        let named = match input {
436            Value::Struct(fields) => fields.clone(),
437            Value::Unit => Vec::new(),
438            _ => Vec::new(),
439        };
440        ready.push(Ready {
441            idx,
442            id: id.clone(),
443            name: name.clone(),
444            tool,
445            call_args: ToolArgs {
446                positional: Vec::new(),
447                named,
448            },
449        });
450    }
451    // Parallel: serial awaits hid all but the first pending node from the UI.
452    for r in &ready {
453        emit_tool_use_start(ctx, &r.name, &r.id, &r.call_args);
454    }
455    let gates = ready.iter().map(|r| {
456        let level = r.tool.approval_level(&r.call_args, ctx);
457        request_approval(
458            ctx,
459            &r.id,
460            &r.name,
461            &r.call_args,
462            level,
463            Some(r.tool.as_ref()),
464        )
465    });
466    let outcomes = futures::future::join_all(gates).await;
467    for (r, gate) in ready.into_iter().zip(outcomes) {
468        let part = match gate {
469            ApprovalOutcome::Deny { reason } => MessagePart::ToolResult {
470                tool_use_id: r.id.clone(),
471                content: format!("sub-agent: tool `{}` denied — {reason}", r.name),
472                is_error: true,
473            },
474            ApprovalOutcome::Approve => match r.tool.call(r.call_args, ctx).await {
475                Ok(v) => MessagePart::ToolResult {
476                    tool_use_id: r.id.clone(),
477                    content: format_value(&v),
478                    is_error: false,
479                },
480                Err(e) => MessagePart::ToolResult {
481                    tool_use_id: r.id.clone(),
482                    content: format!("{e}"),
483                    is_error: true,
484                },
485            },
486        };
487        emit_tool_use_done(ctx, &r.name, &r.id, &part);
488        out[r.idx] = Some(part);
489    }
490    out.into_iter().flatten().collect()
491}
492
493fn emit_tool_use_start(ctx: &ToolCtx, name: &str, id: &str, args: &ToolArgs) {
494    if let Some(tx) = &ctx.stream_tx {
495        let _ = tx.send(crate::stream::StreamFrame::ToolUseStart {
496            tool: name.to_string(),
497            args_preview: preview_tool_args(args),
498            id: id.to_string(),
499        });
500    }
501}
502
503fn emit_tool_use_done(ctx: &ToolCtx, name: &str, id: &str, part: &MessagePart) {
504    if let Some(tx) = &ctx.stream_tx {
505        let (ok, preview) = match part {
506            MessagePart::ToolResult {
507                content, is_error, ..
508            } => (!is_error, truncate(content, 400)),
509            _ => (false, String::new()),
510        };
511        let _ = tx.send(crate::stream::StreamFrame::ToolUseDone {
512            tool: name.to_string(),
513            ok,
514            preview,
515            id: id.to_string(),
516        });
517    }
518}
519
520fn preview_tool_args(args: &ToolArgs) -> String {
521    let mut parts: Vec<String> = args.positional.iter().map(preview_value).collect();
522    for (k, v) in &args.named {
523        parts.push(format!("{k}={}", preview_value(v)));
524    }
525    truncate(&parts.join(", "), 4000)
526}
527
528fn preview_value(v: &Value) -> String {
529    match v {
530        Value::Str(s) => format!("{s:?}"),
531        Value::Int(n) => n.to_string(),
532        Value::Bool(b) => b.to_string(),
533        Value::Float(f) => f.to_string(),
534        Value::Unit => "()".into(),
535        Value::List(items) => format!("list[{}]", items.len()),
536        Value::Struct(items) => format!("struct[{}]", items.len()),
537        Value::Message(_) => "<message>".into(),
538        Value::Path(p) => p.display().to_string(),
539        Value::EditProposal(_) => "<edit proposal>".into(),
540        Value::Err(e) => format!("error({e})"),
541    }
542}
543
544fn format_value(v: &Value) -> String {
545    match v {
546        Value::Str(s) => s.clone(),
547        other => format!("{other:?}"),
548    }
549}
550
551fn emit_child_flow_start(ctx: &ToolCtx, run_id: &FlowRunId, goal: &str) {
552    let parent_run_id = ctx.flow_run_id.clone();
553    let parent_node_id = ctx.current_node_id.clone();
554    if let Some(sink) = &ctx.events {
555        sink.emit(Event::FlowStart {
556            seq: 0,
557            run_id: run_id.clone(),
558            flow_name: "agent.sub".into(),
559            parent_run_id: parent_run_id.clone(),
560            parent_node_id: parent_node_id.clone(),
561            ts: chrono::Utc::now(),
562        });
563    }
564    if let Some(tx) = &ctx.stream_tx {
565        let _ = tx.send(crate::stream::StreamFrame::FlowStart {
566            run_id: run_id.0.to_string(),
567            flow_name: format!("agent.sub · {}", truncate(goal, 60)),
568            parent_run_id: parent_run_id.as_ref().map(|r| r.0.to_string()),
569            parent_node_id,
570        });
571    }
572}
573
574fn emit_flow_agent_start(ctx: &ToolCtx, run_id: &FlowRunId, flow_name: &str) {
575    let parent_run_id = ctx.flow_run_id.clone();
576    let parent_node_id = ctx.current_node_id.clone();
577    if let Some(sink) = &ctx.events {
578        sink.emit(Event::FlowStart {
579            seq: 0,
580            run_id: run_id.clone(),
581            flow_name: flow_name.into(),
582            parent_run_id: parent_run_id.clone(),
583            parent_node_id: parent_node_id.clone(),
584            ts: chrono::Utc::now(),
585        });
586    }
587    if let Some(tx) = &ctx.stream_tx {
588        let _ = tx.send(crate::stream::StreamFrame::FlowStart {
589            run_id: run_id.0.to_string(),
590            flow_name: flow_name.into(),
591            parent_run_id: parent_run_id.as_ref().map(|r| r.0.to_string()),
592            parent_node_id,
593        });
594    }
595}
596
597fn emit_child_flow_end(ctx: &ToolCtx, run_id: &FlowRunId, status: &FlowStatus) {
598    if let Some(sink) = &ctx.events {
599        sink.emit(Event::FlowEnd {
600            seq: 0,
601            run_id: run_id.clone(),
602            flow_name: "agent.sub".into(),
603            status: status.clone(),
604            ts: chrono::Utc::now(),
605        });
606    }
607    if let Some(tx) = &ctx.stream_tx {
608        let _ = tx.send(crate::stream::StreamFrame::FlowDone {
609            run_id: run_id.0.to_string(),
610            flow_name: "agent.sub".into(),
611            ok: matches!(status, FlowStatus::Ok),
612            cancelled: false,
613        });
614    }
615}
616
617fn emit_child_llm_call(
618    ctx: &ToolCtx,
619    _run_id: &FlowRunId,
620    model: &str,
621    am: &crate::provider::AssistantMessage,
622) {
623    if let Some(sink) = &ctx.events {
624        sink.emit(Event::LlmCall {
625            seq: 0,
626            model: model.into(),
627            provider: "sub".into(),
628            usage: am.token_usage.clone(),
629            wallclock_ms: 0,
630            ttft_ms: am.timing.ttft_ms,
631            tokens_per_second: am.timing.tokens_per_second(am.token_usage.output),
632            status: crate::event::LlmCallStatus::Ok,
633            run_id: None,
634            node_id: None,
635            ts: chrono::Utc::now(),
636        });
637    }
638}
639
640fn emit_assistant_msg(ctx: &ToolCtx, run_id: &FlowRunId, message: &Message) {
641    let turn_id = message.turn_id.clone();
642    if let Some(sink) = &ctx.events {
643        sink.emit(Event::AssistantMsg {
644            seq: 0,
645            turn_id: turn_id.clone(),
646            flow_run_id: Some(run_id.clone()),
647            message: message.clone(),
648            ts: chrono::Utc::now(),
649        });
650    }
651    if let Some(tx) = &ctx.stream_tx {
652        let _ = tx.send(crate::stream::StreamFrame::AssistantMsg {
653            flow_run_id: Some(run_id.0.to_string()),
654            message: message.clone(),
655        });
656    }
657}
658
659fn emit_tool_result_msg(ctx: &ToolCtx, run_id: &FlowRunId, message: &Message) {
660    let turn_id = message.turn_id.clone();
661    if let Some(sink) = &ctx.events {
662        sink.emit(Event::ToolResultMsg {
663            seq: 0,
664            turn_id: turn_id.clone(),
665            flow_run_id: Some(run_id.clone()),
666            message: message.clone(),
667            ts: chrono::Utc::now(),
668        });
669    }
670    if let Some(tx) = &ctx.stream_tx {
671        let _ = tx.send(crate::stream::StreamFrame::ToolResultMsg {
672            flow_run_id: Some(run_id.0.to_string()),
673            message: message.clone(),
674        });
675    }
676}
677
678fn truncate(s: &str, n: usize) -> String {
679    let chars: Vec<char> = s.chars().collect();
680    if chars.len() <= n {
681        s.to_string()
682    } else {
683        chars.iter().take(n).collect::<String>() + "…"
684    }
685}