a3s_code_core/
dynamic_workflow.rs

1//! A3S Flow-backed dynamic workflow runtime.
2//!
3//! `DynamicWorkflowRuntime` lets hosts run a sandboxed PTC script as an A3S
4//! Flow runtime. Flow owns durable replay and step lifecycle; A3S Code's
5//! existing `program` tool remains the sandbox and tool-call boundary.
6
7use crate::tools::{Tool, ToolContext, ToolOutput, ToolRegistry, ToolResult};
8use crate::{
9    agent::AgentEvent,
10    planning::{Complexity, ExecutionPlan, Task, TaskStatus},
11};
12use a3s_flow::{
13    FlowEngine, FlowEvent, FlowEventEnvelope, FlowEventObserver, FlowEventStore, FlowRuntime,
14    InMemoryEventStore, LocalFileEventStore, RuntimeCommand, StepInvocation, WorkflowInvocation,
15    WorkflowRunStatus, WorkflowSpec,
16};
17use anyhow::Result;
18use async_trait::async_trait;
19use serde::{Deserialize, Serialize};
20use serde_json::{json, Map, Value};
21use std::collections::BTreeSet;
22use std::sync::Arc;
23use tokio::sync::{broadcast, Mutex};
24
25const DYNAMIC_WORKFLOW_TOOL: &str = "dynamic_workflow";
26const PROGRAM_TOOL: &str = "program";
27const PARALLEL_TASK_TOOL: &str = "parallel_task";
28
29/// Limits forwarded to the underlying PTC `program` tool.
30#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
31#[serde(rename_all = "camelCase")]
32pub struct DynamicWorkflowScriptLimits {
33    #[serde(skip_serializing_if = "Option::is_none")]
34    pub timeout_ms: Option<u64>,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub max_tool_calls: Option<usize>,
37    #[serde(skip_serializing_if = "Option::is_none")]
38    pub max_output_bytes: Option<usize>,
39}
40
41/// Runs A3S Flow workflow and step invocations through a sandboxed PTC script.
42#[derive(Clone)]
43pub struct DynamicWorkflowRuntime {
44    registry: Arc<ToolRegistry>,
45    context: ToolContext,
46    source: Arc<str>,
47    allowed_tools: Vec<String>,
48    limits: DynamicWorkflowScriptLimits,
49}
50
51impl DynamicWorkflowRuntime {
52    pub fn new(
53        registry: Arc<ToolRegistry>,
54        context: ToolContext,
55        source: impl Into<String>,
56    ) -> Self {
57        let allowed_tools = default_allowed_tools(&registry);
58        Self {
59            registry,
60            context,
61            source: Arc::from(source.into()),
62            allowed_tools,
63            limits: DynamicWorkflowScriptLimits::default(),
64        }
65    }
66
67    pub fn with_allowed_tools(mut self, allowed_tools: impl IntoIterator<Item = String>) -> Self {
68        self.allowed_tools = sanitize_allowed_tools(allowed_tools);
69        self
70    }
71
72    pub fn with_limits(mut self, limits: DynamicWorkflowScriptLimits) -> Self {
73        self.limits = limits;
74        self
75    }
76
77    async fn run_script(&self, payload: Value) -> a3s_flow::Result<ToolResult> {
78        let mut args = json!({
79            "type": "script",
80            "language": "javascript",
81            "source": self.source.as_ref(),
82            "inputs": payload,
83            "allowed_tools": self.allowed_tools,
84        });
85        if let Some(object) = args.as_object_mut() {
86            if let Ok(Value::Object(limits)) = serde_json::to_value(&self.limits) {
87                if !limits.is_empty() {
88                    object.insert("limits".to_string(), Value::Object(limits));
89                }
90            }
91        }
92
93        let result = self
94            .registry
95            .execute_with_context(PROGRAM_TOOL, &args, &self.context)
96            .await
97            .map_err(|err| a3s_flow::FlowError::Runtime(err.to_string()))?;
98        if result.exit_code != 0 {
99            return Err(a3s_flow::FlowError::Runtime(result.output));
100        }
101        Ok(result)
102    }
103
104    async fn run_tool_step(&self, tool_name: &str, args: Value) -> a3s_flow::Result<Value> {
105        let result = self
106            .registry
107            .execute_with_context(tool_name, &args, &self.context)
108            .await
109            .map_err(|err| a3s_flow::FlowError::Runtime(err.to_string()))?;
110        if result.exit_code != 0 {
111            return Err(a3s_flow::FlowError::Runtime(result.output));
112        }
113        Ok(json!({
114            "tool": result.name,
115            "output": result.output,
116            "exit_code": result.exit_code,
117            "metadata": result.metadata,
118        }))
119    }
120}
121
122#[async_trait]
123impl FlowRuntime for DynamicWorkflowRuntime {
124    async fn run_workflow(
125        &self,
126        invocation: WorkflowInvocation,
127    ) -> a3s_flow::Result<RuntimeCommand> {
128        let payload = invocation_payload("workflow", &invocation.run_id, &invocation.history)
129            .with("input", invocation.input);
130        let result = self.run_script(payload.into_value()).await?;
131        serde_json::from_value(script_result(&result)?).map_err(a3s_flow::FlowError::from)
132    }
133
134    async fn run_step(&self, invocation: StepInvocation) -> a3s_flow::Result<Value> {
135        if invocation.step_name == PARALLEL_TASK_TOOL {
136            return self
137                .run_tool_step(PARALLEL_TASK_TOOL, invocation.input)
138                .await;
139        }
140
141        let payload = invocation_payload("step", &invocation.run_id, &invocation.history)
142            .with("step_id", invocation.step_id)
143            .with("step_name", invocation.step_name)
144            .with("input", invocation.input);
145        let result = self.run_script(payload.into_value()).await?;
146        script_result(&result)
147    }
148}
149
150struct WorkflowProgressState {
151    tasks: Vec<Task>,
152}
153
154impl WorkflowProgressState {
155    fn new() -> Self {
156        Self { tasks: Vec::new() }
157    }
158
159    fn upsert_step(
160        &mut self,
161        step_id: &str,
162        step_name: &str,
163        input: Option<&Value>,
164        status: TaskStatus,
165    ) {
166        let content = workflow_step_description(step_id, step_name, input);
167        if let Some(task) = self.tasks.iter_mut().find(|task| task.id == step_id) {
168            task.content = content;
169            task.status = status;
170            task.tool = Some(step_name.to_string());
171        } else {
172            self.tasks
173                .push(Task::new(step_id.to_string(), content).with_tool(step_name));
174            if let Some(task) = self.tasks.last_mut() {
175                task.status = status;
176            }
177        }
178    }
179
180    fn mark_status(&mut self, step_id: &str, status: TaskStatus) {
181        if let Some(task) = self.tasks.iter_mut().find(|task| task.id == step_id) {
182            task.status = status;
183        }
184    }
185
186    fn step_position(&self, step_id: &str) -> (usize, usize) {
187        let total = self.tasks.len().max(1);
188        let number = self
189            .tasks
190            .iter()
191            .position(|task| task.id == step_id)
192            .map(|idx| idx + 1)
193            .unwrap_or(total);
194        (number, total)
195    }
196
197    fn step_description(&self, step_id: &str) -> String {
198        self.tasks
199            .iter()
200            .find(|task| task.id == step_id)
201            .map(|task| task.content.clone())
202            .unwrap_or_else(|| step_id.to_string())
203    }
204}
205
206struct AgentEventFlowObserver {
207    tx: broadcast::Sender<AgentEvent>,
208    session_id: String,
209    state: Mutex<WorkflowProgressState>,
210}
211
212impl AgentEventFlowObserver {
213    fn new(tx: broadcast::Sender<AgentEvent>, session_id: String) -> Self {
214        Self {
215            tx,
216            session_id,
217            state: Mutex::new(WorkflowProgressState::new()),
218        }
219    }
220
221    fn emit_task_update(&self, tasks: &[Task]) {
222        let _ = self.tx.send(AgentEvent::TaskUpdated {
223            session_id: self.session_id.clone(),
224            tasks: tasks.to_vec(),
225        });
226    }
227}
228
229#[async_trait]
230impl FlowEventObserver for AgentEventFlowObserver {
231    async fn observe(&self, envelope: FlowEventEnvelope) {
232        match envelope.event {
233            FlowEvent::RunStarted => {
234                let _ = self.tx.send(AgentEvent::PlanningStart {
235                    prompt: "dynamic_workflow".to_string(),
236                });
237            }
238            FlowEvent::StepCreated {
239                step_id,
240                step_name,
241                input,
242                ..
243            } => {
244                let mut state = self.state.lock().await;
245                state.upsert_step(&step_id, &step_name, Some(&input), TaskStatus::Pending);
246                self.emit_task_update(&state.tasks);
247                let mut plan = ExecutionPlan::new("dynamic workflow", Complexity::Medium);
248                for task in state.tasks.iter().cloned() {
249                    plan.add_step(task);
250                }
251                let _ = self.tx.send(AgentEvent::PlanningEnd {
252                    estimated_steps: plan.steps.len(),
253                    plan,
254                });
255            }
256            FlowEvent::StepStarted { step_id, .. } => {
257                let mut state = self.state.lock().await;
258                state.mark_status(&step_id, TaskStatus::InProgress);
259                self.emit_task_update(&state.tasks);
260                let (step_number, total_steps) = state.step_position(&step_id);
261                let _ = self.tx.send(AgentEvent::StepStart {
262                    description: state.step_description(&step_id),
263                    step_id,
264                    step_number,
265                    total_steps,
266                });
267            }
268            FlowEvent::StepCompleted { step_id, .. } => {
269                let mut state = self.state.lock().await;
270                state.mark_status(&step_id, TaskStatus::Completed);
271                self.emit_task_update(&state.tasks);
272                let (step_number, total_steps) = state.step_position(&step_id);
273                let _ = self.tx.send(AgentEvent::StepEnd {
274                    step_id,
275                    status: TaskStatus::Completed,
276                    step_number,
277                    total_steps,
278                });
279            }
280            FlowEvent::StepRetrying { step_id, .. } => {
281                let mut state = self.state.lock().await;
282                state.mark_status(&step_id, TaskStatus::InProgress);
283                self.emit_task_update(&state.tasks);
284            }
285            FlowEvent::StepFailed { step_id, .. } => {
286                let mut state = self.state.lock().await;
287                state.mark_status(&step_id, TaskStatus::Failed);
288                self.emit_task_update(&state.tasks);
289                let (step_number, total_steps) = state.step_position(&step_id);
290                let _ = self.tx.send(AgentEvent::StepEnd {
291                    step_id,
292                    status: TaskStatus::Failed,
293                    step_number,
294                    total_steps,
295                });
296            }
297            FlowEvent::RunFailed { .. } => {
298                let mut state = self.state.lock().await;
299                for task in &mut state.tasks {
300                    if task.status.is_active() {
301                        task.status = TaskStatus::Failed;
302                    }
303                }
304                self.emit_task_update(&state.tasks);
305            }
306            FlowEvent::RunCancelled { .. } => {
307                let mut state = self.state.lock().await;
308                for task in &mut state.tasks {
309                    if task.status.is_active() {
310                        task.status = TaskStatus::Cancelled;
311                    }
312                }
313                self.emit_task_update(&state.tasks);
314            }
315            _ => {}
316        }
317    }
318}
319
320fn workflow_step_description(step_id: &str, step_name: &str, input: Option<&Value>) -> String {
321    if step_name == PARALLEL_TASK_TOOL {
322        let count = input
323            .and_then(|value| value.get("tasks"))
324            .and_then(Value::as_array)
325            .map(Vec::len)
326            .unwrap_or(0);
327        if count > 0 {
328            return format!("Fan out {count} parallel subagent task(s)");
329        }
330    }
331
332    input
333        .and_then(|value| value.get("description").or_else(|| value.get("title")))
334        .and_then(Value::as_str)
335        .map(ToString::to_string)
336        .unwrap_or_else(|| {
337            if step_name == step_id {
338                step_id.to_string()
339            } else {
340                format!("{step_name}: {step_id}")
341            }
342        })
343}
344
345/// Model-visible tool that executes a dynamic workflow through A3S Flow.
346pub struct DynamicWorkflowTool {
347    registry: Arc<ToolRegistry>,
348}
349
350impl DynamicWorkflowTool {
351    pub fn new(registry: Arc<ToolRegistry>) -> Self {
352        Self { registry }
353    }
354}
355
356#[async_trait]
357impl Tool for DynamicWorkflowTool {
358    fn name(&self) -> &str {
359        DYNAMIC_WORKFLOW_TOOL
360    }
361
362    fn description(&self) -> &str {
363        "Run a local dynamic workflow with A3S Flow. The workflow source is a sandboxed JavaScript PTC script that may call allowed ctx tools; A3S Flow records workflow and step history."
364    }
365
366    fn parameters(&self) -> Value {
367        json!({
368            "type": "object",
369            "additionalProperties": false,
370            "properties": {
371                "source": {
372                    "type": "string",
373                    "description": "JavaScript PTC source defining async function run(ctx, inputs). For inputs.kind='workflow', return a Flow command: {type:'complete', output}, {type:'fail', error}, {type:'schedule_step', step_id, step_name, input, retry?}, or {type:'schedule_steps', steps:[...]}. For inputs.kind='step', return the step JSON output. A scheduled step with step_name='parallel_task' bypasses QuickJS and calls the host parallel_task tool directly with input as its arguments."
374                },
375                "input": {
376                    "type": "object",
377                    "description": "Initial workflow input."
378                },
379                "run_id": {
380                    "type": "string",
381                    "description": "Optional durable run id. Reusing it with the same source and input is idempotent."
382                },
383                "allowed_tools": {
384                    "type": "array",
385                    "description": "Tool names the workflow script may call through ctx. Defaults to all registered tools except program, dynamic_workflow, and parallel_task. Login-registered tools such as runtime are allowed when present.",
386                    "items": { "type": "string" }
387                },
388                "limits": {
389                    "type": "object",
390                    "additionalProperties": false,
391                    "properties": {
392                        "timeoutMs": { "type": "integer", "minimum": 1 },
393                        "maxToolCalls": { "type": "integer", "minimum": 1 },
394                        "maxOutputBytes": { "type": "integer", "minimum": 1 }
395                    }
396                }
397            },
398            "required": ["source"]
399        })
400    }
401
402    async fn execute(&self, args: &Value, ctx: &ToolContext) -> Result<ToolOutput> {
403        let Some(source) = args.get("source").and_then(Value::as_str) else {
404            return Ok(ToolOutput::error("dynamic_workflow requires source"));
405        };
406        let input = args.get("input").cloned().unwrap_or_else(|| json!({}));
407        let allowed_tools = args
408            .get("allowed_tools")
409            .and_then(Value::as_array)
410            .map(|items| {
411                items
412                    .iter()
413                    .filter_map(Value::as_str)
414                    .map(ToString::to_string)
415                    .collect::<Vec<_>>()
416            })
417            .unwrap_or_else(|| default_allowed_tools(&self.registry));
418        let limits = args
419            .get("limits")
420            .cloned()
421            .and_then(|value| serde_json::from_value(value).ok())
422            .unwrap_or_default();
423
424        let runtime = Arc::new(
425            DynamicWorkflowRuntime::new(Arc::clone(&self.registry), ctx.clone(), source)
426                .with_allowed_tools(allowed_tools)
427                .with_limits(limits),
428        );
429        let store = flow_store_for_context(ctx);
430        let engine = match ctx.agent_event_tx.clone() {
431            Some(tx) => FlowEngine::builder(runtime)
432                .with_store(store)
433                .with_observer(Arc::new(AgentEventFlowObserver::new(
434                    tx,
435                    ctx.session_id.clone().unwrap_or_default(),
436                )))
437                .build(),
438            None => FlowEngine::new(store, runtime),
439        };
440        let source_hash = source_hash(source);
441        let spec = WorkflowSpec::rust_embedded(
442            "a3s-code.dynamic-workflow",
443            source_hash.as_str(),
444            "ptc",
445            "run",
446        );
447
448        let run_id = match args.get("run_id").and_then(Value::as_str) {
449            Some(run_id) => match engine.start_with_id(run_id, spec, input).await {
450                Ok(run_id) => run_id,
451                Err(err) => return Ok(ToolOutput::error(err.to_string())),
452            },
453            None => match engine.start(spec, input).await {
454                Ok(run_id) => run_id,
455                Err(err) => return Ok(ToolOutput::error(err.to_string())),
456            },
457        };
458
459        let snapshot = match engine.snapshot(&run_id).await {
460            Ok(snapshot) => snapshot,
461            Err(err) => return Ok(ToolOutput::error(err.to_string())),
462        };
463        let history = match engine.history(&run_id).await {
464            Ok(history) => history,
465            Err(err) => return Ok(ToolOutput::error(err.to_string())),
466        };
467
468        let output = match &snapshot.output {
469            Some(output) => {
470                serde_json::to_string_pretty(output).unwrap_or_else(|_| output.to_string())
471            }
472            None => snapshot
473                .error
474                .clone()
475                .unwrap_or_else(|| format!("workflow status: {:?}", snapshot.status)),
476        };
477
478        let status = snapshot.status;
479        let metadata = json!({
480            "dynamic_workflow": {
481                "run_id": run_id,
482                "status": format!("{:?}", snapshot.status),
483                "last_sequence": snapshot.last_sequence,
484                "source_hash": source_hash,
485                "snapshot": snapshot,
486                "history": history,
487            }
488        });
489        let output = match status {
490            WorkflowRunStatus::Completed => ToolOutput::success(output),
491            WorkflowRunStatus::Failed | WorkflowRunStatus::Cancelled => ToolOutput::error(output),
492            _ => ToolOutput::error(format!(
493                "dynamic_workflow ended without a terminal result: {status:?}; {output}"
494            )),
495        };
496
497        Ok(output.with_metadata(metadata))
498    }
499}
500
501pub fn register_dynamic_workflow(registry: &Arc<ToolRegistry>) {
502    registry.register(Arc::new(DynamicWorkflowTool::new(Arc::clone(registry))));
503}
504
505fn flow_store_for_context(ctx: &ToolContext) -> Arc<dyn FlowEventStore> {
506    match ctx.workspace_services.local_root() {
507        Some(root) => Arc::new(LocalFileEventStore::new(
508            root.join(".a3s-flow").join("dynamic-workflows"),
509        )),
510        None => Arc::new(InMemoryEventStore::new()),
511    }
512}
513
514struct PayloadBuilder {
515    value: Map<String, Value>,
516}
517
518impl PayloadBuilder {
519    fn with(mut self, key: &str, value: impl Serialize) -> Self {
520        self.value.insert(
521            key.to_string(),
522            serde_json::to_value(value).unwrap_or(Value::Null),
523        );
524        self
525    }
526
527    fn into_value(self) -> Value {
528        Value::Object(self.value)
529    }
530}
531
532fn invocation_payload(kind: &str, run_id: &str, history: &[FlowEventEnvelope]) -> PayloadBuilder {
533    let mut value = Map::new();
534    value.insert("kind".to_string(), json!(kind));
535    value.insert("run_id".to_string(), json!(run_id));
536    value.insert("history".to_string(), json!(history));
537    value.insert("step_outputs".to_string(), completed_step_outputs(history));
538    value.insert("step_failures".to_string(), failed_step_outputs(history));
539    PayloadBuilder { value }
540}
541
542fn completed_step_outputs(history: &[FlowEventEnvelope]) -> Value {
543    let mut outputs = Map::new();
544    for envelope in history {
545        if let FlowEvent::StepCompleted { step_id, output } = &envelope.event {
546            outputs.insert(step_id.clone(), output.clone());
547        }
548    }
549    Value::Object(outputs)
550}
551
552fn failed_step_outputs(history: &[FlowEventEnvelope]) -> Value {
553    let mut outputs = Map::new();
554    for envelope in history {
555        if let FlowEvent::StepFailed {
556            step_id,
557            attempt,
558            error,
559        } = &envelope.event
560        {
561            outputs.insert(
562                step_id.clone(),
563                json!({
564                    "attempt": attempt,
565                    "error": error,
566                }),
567            );
568        }
569    }
570    Value::Object(outputs)
571}
572
573fn script_result(result: &ToolResult) -> a3s_flow::Result<Value> {
574    result
575        .metadata
576        .as_ref()
577        .and_then(|metadata| metadata.get("script_result"))
578        .cloned()
579        .ok_or_else(|| {
580            a3s_flow::FlowError::Runtime(
581                "PTC program result did not include script_result metadata".to_string(),
582            )
583        })
584}
585
586fn default_allowed_tools(registry: &ToolRegistry) -> Vec<String> {
587    sanitize_allowed_tools(registry.list())
588}
589
590fn sanitize_allowed_tools(items: impl IntoIterator<Item = String>) -> Vec<String> {
591    let mut tools = items.into_iter().collect::<BTreeSet<_>>();
592    tools.remove(PROGRAM_TOOL);
593    tools.remove(DYNAMIC_WORKFLOW_TOOL);
594    tools.remove(PARALLEL_TASK_TOOL);
595    tools.into_iter().collect()
596}
597
598fn source_hash(source: &str) -> String {
599    sha256::digest(source.as_bytes())
600}
601
602#[cfg(test)]
603mod tests {
604    use super::*;
605    use crate::tools::{ToolExecutor, ToolOutput};
606
607    struct FakeParallelTaskTool;
608
609    #[async_trait]
610    impl Tool for FakeParallelTaskTool {
611        fn name(&self) -> &str {
612            PARALLEL_TASK_TOOL
613        }
614
615        fn description(&self) -> &str {
616            "Fake parallel task tool for DynamicWorkflowRuntime tests."
617        }
618
619        fn parameters(&self) -> Value {
620            json!({ "type": "object" })
621        }
622
623        async fn execute(&self, args: &Value, _ctx: &ToolContext) -> Result<ToolOutput> {
624            let count = args
625                .get("tasks")
626                .and_then(Value::as_array)
627                .map(Vec::len)
628                .unwrap_or(0);
629            Ok(ToolOutput::success(format!("parallel:{count}"))
630                .with_metadata(json!({ "task_count": count })))
631        }
632    }
633
634    struct FakeRuntimeTool;
635
636    #[async_trait]
637    impl Tool for FakeRuntimeTool {
638        fn name(&self) -> &str {
639            "runtime"
640        }
641
642        fn description(&self) -> &str {
643            "Fake OS runtime tool for DynamicWorkflowRuntime tests."
644        }
645
646        fn parameters(&self) -> Value {
647            json!({ "type": "object" })
648        }
649
650        async fn execute(&self, args: &Value, _ctx: &ToolContext) -> Result<ToolOutput> {
651            let tasks = args
652                .get("tasks")
653                .and_then(Value::as_array)
654                .map(Vec::len)
655                .unwrap_or(0);
656            Ok(ToolOutput::success(format!("runtime:{tasks}"))
657                .with_metadata(json!({ "runtime_tasks": tasks })))
658        }
659    }
660
661    struct FailingRuntimeTool;
662
663    #[async_trait]
664    impl Tool for FailingRuntimeTool {
665        fn name(&self) -> &str {
666            "runtime"
667        }
668
669        fn description(&self) -> &str {
670            "Failing OS runtime tool for DynamicWorkflowRuntime tests."
671        }
672
673        fn parameters(&self) -> Value {
674            json!({ "type": "object" })
675        }
676
677        async fn execute(&self, _args: &Value, _ctx: &ToolContext) -> Result<ToolOutput> {
678            Ok(ToolOutput::error("runtime unavailable"))
679        }
680    }
681
682    #[tokio::test]
683    async fn dynamic_workflow_tool_runs_ptc_step_through_a3s_flow() {
684        let dir = tempfile::tempdir().unwrap();
685        tokio::fs::write(dir.path().join("fixture.txt"), "hello from fixture")
686            .await
687            .unwrap();
688        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
689        register_dynamic_workflow(executor.registry());
690
691        let source = r#"
692async function run(ctx, inputs) {
693  if (inputs.kind === "workflow") {
694    const read = inputs.step_outputs.read_fixture;
695    if (read) {
696      return { type: "complete", output: { text: read.output } };
697    }
698    return {
699      type: "schedule_step",
700      step_id: "read_fixture",
701      step_name: "read_fixture",
702      input: { path: inputs.input.path },
703      retry: { max_attempts: 1, delay_ms: 0 },
704    };
705  }
706
707  if (inputs.kind === "step" && inputs.step_name === "read_fixture") {
708    return await ctx.read(inputs.input.path);
709  }
710
711  return { error: "unknown invocation" };
712}
713"#;
714
715        let result = executor
716            .execute(
717                DYNAMIC_WORKFLOW_TOOL,
718                &json!({
719                    "source": source,
720                    "input": { "path": "fixture.txt" },
721                    "run_id": "test-dynamic-workflow",
722                    "allowed_tools": ["read"],
723                }),
724            )
725            .await
726            .unwrap();
727
728        assert_eq!(result.exit_code, 0, "{}", result.output);
729        assert!(
730            result.output.contains("hello from fixture"),
731            "{}",
732            result.output
733        );
734        let metadata = result.metadata.unwrap();
735        assert_eq!(
736            metadata["dynamic_workflow"]["run_id"],
737            "test-dynamic-workflow"
738        );
739        assert_eq!(metadata["dynamic_workflow"]["status"], "Completed");
740        assert_eq!(
741            metadata["dynamic_workflow"]["snapshot"]["steps"]["read_fixture"]["status"],
742            "completed"
743        );
744    }
745
746    #[tokio::test]
747    async fn dynamic_workflow_emits_agent_progress_events() {
748        let dir = tempfile::tempdir().unwrap();
749        tokio::fs::write(dir.path().join("fixture.txt"), "hello from fixture")
750            .await
751            .unwrap();
752        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
753        register_dynamic_workflow(executor.registry());
754        let (tx, mut rx) = broadcast::channel(64);
755        let ctx = ToolContext::new(dir.path().to_path_buf())
756            .with_session_id("progress-session")
757            .with_agent_event_tx(tx);
758
759        let source = r#"
760async function run(ctx, inputs) {
761  if (inputs.kind === "workflow") {
762    const read = inputs.step_outputs.read_fixture;
763    if (read) {
764      return { type: "complete", output: { text: read.output } };
765    }
766    return {
767      type: "schedule_step",
768      step_id: "read_fixture",
769      step_name: "read_fixture",
770      input: { path: inputs.input.path, description: "Read fixture" },
771      retry: { max_attempts: 1, delay_ms: 0 },
772    };
773  }
774
775  if (inputs.kind === "step" && inputs.step_name === "read_fixture") {
776    return await ctx.read(inputs.input.path);
777  }
778
779  return { error: "unknown invocation" };
780}
781"#;
782
783        let result = executor
784            .execute_with_context(
785                DYNAMIC_WORKFLOW_TOOL,
786                &json!({
787                    "source": source,
788                    "input": { "path": "fixture.txt" },
789                    "run_id": "test-dynamic-workflow-progress",
790                    "allowed_tools": ["read"],
791                }),
792                &ctx,
793            )
794            .await
795            .unwrap();
796
797        assert_eq!(result.exit_code, 0, "{}", result.output);
798        let mut events = Vec::new();
799        while let Ok(event) = rx.try_recv() {
800            events.push(event);
801        }
802
803        assert!(
804            events
805                .iter()
806                .any(|event| matches!(event, AgentEvent::PlanningStart { .. })),
807            "{events:?}"
808        );
809        assert!(
810            events.iter().any(|event| matches!(
811                event,
812                AgentEvent::TaskUpdated { tasks, .. }
813                    if tasks.iter().any(|task| task.id == "read_fixture")
814            )),
815            "{events:?}"
816        );
817        assert!(
818            events.iter().any(|event| matches!(
819                event,
820                AgentEvent::StepStart { step_id, .. } if step_id == "read_fixture"
821            )),
822            "{events:?}"
823        );
824        assert!(
825            events.iter().any(|event| matches!(
826                event,
827                AgentEvent::StepEnd { step_id, status, .. }
828                    if step_id == "read_fixture" && *status == TaskStatus::Completed
829            )),
830            "{events:?}"
831        );
832    }
833
834    #[tokio::test]
835    async fn dynamic_workflow_step_can_call_host_parallel_task_without_ptc_parallel_task() {
836        let dir = tempfile::tempdir().unwrap();
837        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
838        executor.register_dynamic_tool(Arc::new(FakeParallelTaskTool));
839        register_dynamic_workflow(executor.registry());
840
841        let source = r#"
842async function run(ctx, inputs) {
843  if (inputs.kind === "workflow") {
844    const fanout = inputs.step_outputs.fanout;
845    if (fanout) {
846      return { type: "complete", output: { fanout } };
847    }
848    return {
849      type: "schedule_step",
850      step_id: "fanout",
851      step_name: "parallel_task",
852      input: {
853        tasks: [
854          { agent: "explore", description: "alpha", prompt: "research alpha" },
855          { agent: "explore", description: "beta", prompt: "research beta" },
856        ],
857      },
858    };
859  }
860
861  return { error: "ptc step handler should not run for parallel_task" };
862}
863"#;
864
865        let result = executor
866            .execute(
867                DYNAMIC_WORKFLOW_TOOL,
868                &json!({
869                    "source": source,
870                    "run_id": "test-dynamic-workflow-parallel-step",
871                    "allowed_tools": [],
872                }),
873            )
874            .await
875            .unwrap();
876
877        assert_eq!(result.exit_code, 0, "{}", result.output);
878        assert!(result.output.contains("parallel:2"), "{}", result.output);
879        let metadata = result.metadata.unwrap();
880        assert_eq!(metadata["dynamic_workflow"]["status"], "Completed");
881        let step = &metadata["dynamic_workflow"]["snapshot"]["steps"]["fanout"];
882        assert_eq!(step["status"], "completed");
883        assert_eq!(step["output"]["tool"], PARALLEL_TASK_TOOL);
884        assert_eq!(step["output"]["metadata"]["task_count"], 2);
885    }
886
887    #[tokio::test]
888    async fn dynamic_workflow_ptc_step_can_call_login_registered_runtime_tool_by_default() {
889        let dir = tempfile::tempdir().unwrap();
890        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
891        executor.register_dynamic_tool(Arc::new(FakeRuntimeTool));
892        register_dynamic_workflow(executor.registry());
893
894        let source = r#"
895async function run(ctx, inputs) {
896  if (inputs.kind === "workflow") {
897    const runtime = inputs.step_outputs.runtime_fanout;
898    if (runtime) {
899      return { type: "complete", output: { runtime } };
900    }
901    return {
902      type: "schedule_step",
903      step_id: "runtime_fanout",
904      step_name: "runtime_fanout",
905      input: {
906        worker: "research-worker",
907        tasks: ["alpha", "beta", "gamma"],
908      },
909    };
910  }
911
912  if (inputs.kind === "step" && inputs.step_name === "runtime_fanout") {
913    return await ctx.tool("runtime", inputs.input);
914  }
915
916  return { error: "unknown invocation" };
917}
918"#;
919
920        let result = executor
921            .execute(
922                DYNAMIC_WORKFLOW_TOOL,
923                &json!({
924                    "source": source,
925                    "run_id": "test-dynamic-workflow-runtime-step",
926                }),
927            )
928            .await
929            .unwrap();
930
931        assert_eq!(result.exit_code, 0, "{}", result.output);
932        assert!(result.output.contains("runtime:3"), "{}", result.output);
933        let metadata = result.metadata.unwrap();
934        let step = &metadata["dynamic_workflow"]["snapshot"]["steps"]["runtime_fanout"];
935        assert_eq!(step["status"], "completed");
936        assert_eq!(step["output"]["name"], "runtime");
937        assert_eq!(step["output"]["metadata"]["runtime_tasks"], 3);
938    }
939
940    #[tokio::test]
941    async fn dynamic_workflow_ptc_step_can_call_legacy_ctx_tools_runtime_proxy() {
942        let dir = tempfile::tempdir().unwrap();
943        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
944        executor.register_dynamic_tool(Arc::new(FakeRuntimeTool));
945        register_dynamic_workflow(executor.registry());
946
947        let source = r#"
948async function run(ctx, inputs) {
949  if (inputs.kind === "workflow") {
950    const runtime = inputs.step_outputs.runtime_fanout;
951    if (runtime) {
952      return { type: "complete", output: { runtime } };
953    }
954    return {
955      type: "schedule_step",
956      step_id: "runtime_fanout",
957      step_name: "runtime_fanout",
958      input: {
959        worker: "research-worker",
960        tasks: ["alpha", "beta"],
961      },
962    };
963  }
964
965  if (inputs.kind === "step" && inputs.step_name === "runtime_fanout") {
966    return await ctx.tools.runtime(inputs.input);
967  }
968
969  return { error: "unknown invocation" };
970}
971"#;
972
973        let result = executor
974            .execute(
975                DYNAMIC_WORKFLOW_TOOL,
976                &json!({
977                    "source": source,
978                    "run_id": "test-dynamic-workflow-runtime-tools-proxy",
979                }),
980            )
981            .await
982            .unwrap();
983
984        assert_eq!(result.exit_code, 0, "{}", result.output);
985        assert!(result.output.contains("runtime:2"), "{}", result.output);
986        let metadata = result.metadata.unwrap();
987        let step = &metadata["dynamic_workflow"]["snapshot"]["steps"]["runtime_fanout"];
988        assert_eq!(step["status"], "completed");
989        assert_eq!(step["output"]["name"], "runtime");
990        assert_eq!(step["output"]["metadata"]["runtime_tasks"], 2);
991    }
992
993    #[tokio::test]
994    async fn dynamic_workflow_tool_returns_error_when_runtime_step_fails() {
995        let dir = tempfile::tempdir().unwrap();
996        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
997        executor.register_dynamic_tool(Arc::new(FailingRuntimeTool));
998        register_dynamic_workflow(executor.registry());
999
1000        let source = r#"
1001async function run(ctx, inputs) {
1002  if (inputs.kind === "workflow") {
1003    const runtime = inputs.step_outputs.runtime_fanout;
1004    if (runtime) {
1005      return { type: "complete", output: { runtime } };
1006    }
1007    return {
1008      type: "schedule_step",
1009      step_id: "runtime_fanout",
1010      step_name: "runtime_fanout",
1011      input: { worker: "research-worker", tasks: ["alpha"] },
1012    };
1013  }
1014
1015  if (inputs.kind === "step" && inputs.step_name === "runtime_fanout") {
1016    const result = await ctx.tool("runtime", inputs.input);
1017    if (result.exitCode !== 0) {
1018      throw new Error(result.output || "runtime failed");
1019    }
1020    return result;
1021  }
1022
1023  return { error: "unknown invocation" };
1024}
1025"#;
1026
1027        let result = executor
1028            .execute(
1029                DYNAMIC_WORKFLOW_TOOL,
1030                &json!({
1031                    "source": source,
1032                    "run_id": "test-dynamic-workflow-runtime-step-fails",
1033                }),
1034            )
1035            .await
1036            .unwrap();
1037
1038        assert_ne!(result.exit_code, 0, "{}", result.output);
1039        assert!(
1040            result.output.contains("runtime unavailable"),
1041            "{}",
1042            result.output
1043        );
1044        let metadata = result.metadata.unwrap();
1045        assert_eq!(metadata["dynamic_workflow"]["status"], "Failed");
1046        let step = &metadata["dynamic_workflow"]["snapshot"]["steps"]["runtime_fanout"];
1047        assert_eq!(step["status"], "failed");
1048    }
1049
1050    #[tokio::test]
1051    async fn dynamic_workflow_step_failure_can_continue_workflow_with_error_payload() {
1052        let dir = tempfile::tempdir().unwrap();
1053        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
1054        executor.register_dynamic_tool(Arc::new(FailingRuntimeTool));
1055        register_dynamic_workflow(executor.registry());
1056
1057        let source = r#"
1058async function run(ctx, inputs) {
1059  if (inputs.kind === "workflow") {
1060    const failure = inputs.step_failures.runtime_fanout;
1061    if (failure) {
1062      return { type: "complete", output: { recovered: true, error: failure.error } };
1063    }
1064    return {
1065      type: "schedule_step",
1066      step_id: "runtime_fanout",
1067      step_name: "runtime_fanout",
1068      input: { worker: "research-worker" },
1069      retry: { max_attempts: 1, delay_ms: 0, on_exhausted: "continue_workflow" },
1070    };
1071  }
1072
1073  if (inputs.kind === "step" && inputs.step_name === "runtime_fanout") {
1074    const result = await ctx.tool("runtime", inputs.input);
1075    if (result.exitCode !== 0) {
1076      throw new Error(result.output || "runtime failed");
1077    }
1078    return result;
1079  }
1080
1081  return { error: "unknown invocation" };
1082}
1083"#;
1084
1085        let result = executor
1086            .execute(
1087                DYNAMIC_WORKFLOW_TOOL,
1088                &json!({
1089                    "source": source,
1090                    "run_id": "test-dynamic-workflow-continue-after-step-failure",
1091                }),
1092            )
1093            .await
1094            .unwrap();
1095
1096        assert_eq!(result.exit_code, 0, "{}", result.output);
1097        assert!(
1098            result.output.contains("runtime unavailable"),
1099            "{}",
1100            result.output
1101        );
1102        let metadata = result.metadata.unwrap();
1103        assert_eq!(metadata["dynamic_workflow"]["status"], "Completed");
1104        let step = &metadata["dynamic_workflow"]["snapshot"]["steps"]["runtime_fanout"];
1105        assert_eq!(step["status"], "failed");
1106        assert!(step["error"]
1107            .as_str()
1108            .is_some_and(|error| error.contains("runtime unavailable")));
1109    }
1110
1111    #[tokio::test]
1112    async fn dynamic_workflow_tool_returns_error_when_run_is_suspended() {
1113        let dir = tempfile::tempdir().unwrap();
1114        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
1115        register_dynamic_workflow(executor.registry());
1116
1117        let source = r#"
1118async function run(ctx, inputs) {
1119  if (inputs.kind === "workflow") {
1120    return {
1121      type: "wait_until",
1122      wait_id: "external-research-still-running",
1123      resume_at: "2099-01-01T00:00:00Z",
1124    };
1125  }
1126
1127  return { error: "unknown invocation" };
1128}
1129"#;
1130
1131        let result = executor
1132            .execute(
1133                DYNAMIC_WORKFLOW_TOOL,
1134                &json!({
1135                    "source": source,
1136                    "run_id": "test-dynamic-workflow-suspended-is-error",
1137                }),
1138            )
1139            .await
1140            .unwrap();
1141
1142        assert_ne!(result.exit_code, 0, "{}", result.output);
1143        assert!(
1144            result
1145                .output
1146                .contains("dynamic_workflow ended without a terminal result: Suspended"),
1147            "{}",
1148            result.output
1149        );
1150        let metadata = result.metadata.unwrap();
1151        assert_eq!(metadata["dynamic_workflow"]["status"], "Suspended");
1152        assert_eq!(
1153            metadata["dynamic_workflow"]["snapshot"]["waits"]["external-research-still-running"]
1154                ["status"],
1155            "waiting"
1156        );
1157    }
1158
1159    #[test]
1160    fn default_allowed_tools_exclude_recursive_program_and_dynamic_workflow_tools() {
1161        let dir = tempfile::tempdir().unwrap();
1162        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
1163        register_dynamic_workflow(executor.registry());
1164
1165        let tools = default_allowed_tools(executor.registry());
1166
1167        assert!(!tools.contains(&PROGRAM_TOOL.to_string()));
1168        assert!(!tools.contains(&DYNAMIC_WORKFLOW_TOOL.to_string()));
1169        assert!(!tools.contains(&PARALLEL_TASK_TOOL.to_string()));
1170        assert!(tools.contains(&"read".to_string()));
1171    }
1172}
a3s_code_core/dynamic_workflow.rs

a3s_code_core/
dynamic_workflow.rs
