atman_runtime/
executor.rs

1use std::collections::HashMap;
2
3use atman_dsl::ast::{File, FlowDecl};
4
5use crate::error::RuntimeError;
6use crate::event::{Event, EventSink, FlowRunId, FlowStatus, TurnId};
7use crate::exec::exec_flow_with_siblings;
8use crate::provider::ProviderRegistry;
9use crate::session::Session;
10use crate::tool::{ToolCtx, ToolRegistry};
11use crate::value::Value;
12
13#[derive(Clone)]
14pub struct Executor {
15    pub tools: ToolRegistry,
16    pub providers: ProviderRegistry,
17    pub events: EventSink,
18    pub tool_ctx: ToolCtx,
19    pub safety: Option<crate::safety::SafetyConfig>,
20    /// When set, relative `@` paths are resolved against this directory.
21    pub source_dir: Option<std::path::PathBuf>,
22}
23
24impl Executor {
25    pub fn new() -> Self {
26        Self {
27            tools: ToolRegistry::new(),
28            providers: ProviderRegistry::new(),
29            events: EventSink::new(),
30            tool_ctx: ToolCtx::new(),
31            safety: None,
32            source_dir: None,
33        }
34    }
35
36    pub fn with_events(events: EventSink) -> Self {
37        Self {
38            tools: ToolRegistry::new(),
39            providers: ProviderRegistry::new(),
40            events,
41            tool_ctx: ToolCtx::new(),
42            safety: None,
43            source_dir: None,
44        }
45    }
46
47    pub fn with_safety(mut self, safety: crate::safety::SafetyConfig) -> Self {
48        self.safety = Some(safety);
49        self
50    }
51
52    pub async fn run(
53        &self,
54        file: &File,
55        flow_name: &str,
56        args: Vec<(String, Value)>,
57    ) -> Result<Value, RuntimeError> {
58        self.run_in_turn(file, flow_name, args, None, None).await
59    }
60
61    pub async fn run_in_turn(
62        &self,
63        file: &File,
64        flow_name: &str,
65        args: Vec<(String, Value)>,
66        turn_id: Option<TurnId>,
67        session: Option<std::sync::Arc<Session>>,
68    ) -> Result<Value, RuntimeError> {
69        self.run_in_turn_with_run_id(file, flow_name, args, turn_id, session, None)
70            .await
71    }
72
73    pub async fn run_in_turn_with_run_id(
74        &self,
75        file: &File,
76        flow_name: &str,
77        args: Vec<(String, Value)>,
78        turn_id: Option<TurnId>,
79        session: Option<std::sync::Arc<Session>>,
80        first_run_id: Option<FlowRunId>,
81    ) -> Result<Value, RuntimeError> {
82        let flows: HashMap<_, _> = file
83            .flows
84            .iter()
85            .map(|f| (f.name.name.clone(), f.clone()))
86            .collect();
87        let mut current = flow_name.to_string();
88        let mut current_args = args;
89        let mut next_run_id = first_run_id;
90        for _ in 0..5 {
91            let flow = flows
92                .get(&current)
93                .ok_or_else(|| RuntimeError::UndefinedTool(format!("flow `{current}`")))?;
94            match self
95                .run_flow(
96                    flow,
97                    current_args,
98                    &flows,
99                    turn_id.clone(),
100                    session.clone(),
101                    next_run_id.take(),
102                )
103                .await
104            {
105                Err(RuntimeError::Redirect(target)) => {
106                    current = target;
107                    current_args = Vec::new();
108                    continue;
109                }
110                other => return other,
111            }
112        }
113        Err(RuntimeError::ToolFailed(
114            "redirect chain exceeded max depth (5)".into(),
115        ))
116    }
117
118    async fn run_flow(
119        &self,
120        flow: &FlowDecl,
121        args: Vec<(String, Value)>,
122        flows: &HashMap<String, FlowDecl>,
123        turn_id: Option<TurnId>,
124        session: Option<std::sync::Arc<Session>>,
125        run_id: Option<FlowRunId>,
126    ) -> Result<Value, RuntimeError> {
127        let run_id = run_id.unwrap_or_else(FlowRunId::now);
128        self.events.emit(Event::FlowStart {
129            run_id: run_id.clone(),
130            flow_name: flow.name.name.clone(),
131            parent_run_id: None,
132            parent_node_id: None,
133        });
134        if let Some(sess) = session.as_ref() {
135            let _ = sess
136                .stream_tx()
137                .send(crate::stream::StreamFrame::FlowStart {
138                    run_id: run_id.0.to_string(),
139                    flow_name: flow.name.name.clone(),
140                    parent_run_id: None,
141                    parent_node_id: None,
142                });
143        }
144        let graph = crate::nodegraph::extract_graph(flow);
145        self.events.emit(Event::FlowGraph {
146            run_id: run_id.clone(),
147            graph: graph.clone(),
148        });
149        if let Some(sess) = session.as_ref() {
150            let _ = sess
151                .stream_tx()
152                .send(crate::stream::StreamFrame::FlowGraph {
153                    run_id: run_id.0.to_string(),
154                    graph,
155                });
156        }
157        let flow_cancel = session
158            .as_ref()
159            .map(|s| s.flow_cancel_token())
160            .unwrap_or_default();
161        let exec_fut = exec_flow_with_siblings(
162            flow,
163            args,
164            &self.tools,
165            &self.tool_ctx,
166            &self.providers,
167            flows,
168            Some(&self.events),
169            turn_id,
170            Some(run_id.clone()),
171            session.clone(),
172            flow_cancel.clone(),
173            self.safety.as_ref(),
174            self.source_dir.clone(),
175        );
176        let result = tokio::select! {
177            biased;
178            _ = flow_cancel.cancelled() => Err(RuntimeError::Cancelled("flow cancelled by user".into())),
179            r = exec_fut => r,
180        };
181        let status = match &result {
182            Ok(v) => {
183                if let Value::Err(e) = v
184                    && matches!(e, RuntimeError::Cancelled(_))
185                {
186                    FlowStatus::Cancelled
187                } else {
188                    FlowStatus::Ok
189                }
190            }
191            Err(e) => {
192                if matches!(e, RuntimeError::Cancelled(_)) {
193                    FlowStatus::Cancelled
194                } else {
195                    FlowStatus::Errored {
196                        message: e.to_string(),
197                    }
198                }
199            }
200        };
201        let cancelled = matches!(status, FlowStatus::Cancelled);
202        self.events.emit(Event::FlowEnd {
203            run_id: run_id.clone(),
204            flow_name: flow.name.name.clone(),
205            status: status.clone(),
206        });
207        if let Some(sess) = session.as_ref() {
208            let _ = sess.stream_tx().send(crate::stream::StreamFrame::FlowDone {
209                run_id: run_id.0.to_string(),
210                flow_name: flow.name.name.clone(),
211                ok: matches!(status, FlowStatus::Ok),
212                cancelled,
213            });
214        }
215        result
216    }
217}
218
219impl Default for Executor {
220    fn default() -> Self {
221        Self::new()
222    }
223}
atman_runtime/executor.rs

atman_runtime/
executor.rs
