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::success(output),
493        };
494
495        Ok(output.with_metadata(metadata))
496    }
497}
498
499pub fn register_dynamic_workflow(registry: &Arc<ToolRegistry>) {
500    registry.register(Arc::new(DynamicWorkflowTool::new(Arc::clone(registry))));
501}
502
503fn flow_store_for_context(ctx: &ToolContext) -> Arc<dyn FlowEventStore> {
504    match ctx.workspace_services.local_root() {
505        Some(root) => Arc::new(LocalFileEventStore::new(
506            root.join(".a3s-flow").join("dynamic-workflows"),
507        )),
508        None => Arc::new(InMemoryEventStore::new()),
509    }
510}
511
512struct PayloadBuilder {
513    value: Map<String, Value>,
514}
515
516impl PayloadBuilder {
517    fn with(mut self, key: &str, value: impl Serialize) -> Self {
518        self.value.insert(
519            key.to_string(),
520            serde_json::to_value(value).unwrap_or(Value::Null),
521        );
522        self
523    }
524
525    fn into_value(self) -> Value {
526        Value::Object(self.value)
527    }
528}
529
530fn invocation_payload(kind: &str, run_id: &str, history: &[FlowEventEnvelope]) -> PayloadBuilder {
531    let mut value = Map::new();
532    value.insert("kind".to_string(), json!(kind));
533    value.insert("run_id".to_string(), json!(run_id));
534    value.insert("history".to_string(), json!(history));
535    value.insert("step_outputs".to_string(), completed_step_outputs(history));
536    PayloadBuilder { value }
537}
538
539fn completed_step_outputs(history: &[FlowEventEnvelope]) -> Value {
540    let mut outputs = Map::new();
541    for envelope in history {
542        if let FlowEvent::StepCompleted { step_id, output } = &envelope.event {
543            outputs.insert(step_id.clone(), output.clone());
544        }
545    }
546    Value::Object(outputs)
547}
548
549fn script_result(result: &ToolResult) -> a3s_flow::Result<Value> {
550    result
551        .metadata
552        .as_ref()
553        .and_then(|metadata| metadata.get("script_result"))
554        .cloned()
555        .ok_or_else(|| {
556            a3s_flow::FlowError::Runtime(
557                "PTC program result did not include script_result metadata".to_string(),
558            )
559        })
560}
561
562fn default_allowed_tools(registry: &ToolRegistry) -> Vec<String> {
563    sanitize_allowed_tools(registry.list())
564}
565
566fn sanitize_allowed_tools(items: impl IntoIterator<Item = String>) -> Vec<String> {
567    let mut tools = items.into_iter().collect::<BTreeSet<_>>();
568    tools.remove(PROGRAM_TOOL);
569    tools.remove(DYNAMIC_WORKFLOW_TOOL);
570    tools.remove(PARALLEL_TASK_TOOL);
571    tools.into_iter().collect()
572}
573
574fn source_hash(source: &str) -> String {
575    sha256::digest(source.as_bytes())
576}
577
578#[cfg(test)]
579mod tests {
580    use super::*;
581    use crate::tools::{ToolExecutor, ToolOutput};
582
583    struct FakeParallelTaskTool;
584
585    #[async_trait]
586    impl Tool for FakeParallelTaskTool {
587        fn name(&self) -> &str {
588            PARALLEL_TASK_TOOL
589        }
590
591        fn description(&self) -> &str {
592            "Fake parallel task tool for DynamicWorkflowRuntime tests."
593        }
594
595        fn parameters(&self) -> Value {
596            json!({ "type": "object" })
597        }
598
599        async fn execute(&self, args: &Value, _ctx: &ToolContext) -> Result<ToolOutput> {
600            let count = args
601                .get("tasks")
602                .and_then(Value::as_array)
603                .map(Vec::len)
604                .unwrap_or(0);
605            Ok(ToolOutput::success(format!("parallel:{count}"))
606                .with_metadata(json!({ "task_count": count })))
607        }
608    }
609
610    struct FakeRuntimeTool;
611
612    #[async_trait]
613    impl Tool for FakeRuntimeTool {
614        fn name(&self) -> &str {
615            "runtime"
616        }
617
618        fn description(&self) -> &str {
619            "Fake OS runtime tool for DynamicWorkflowRuntime tests."
620        }
621
622        fn parameters(&self) -> Value {
623            json!({ "type": "object" })
624        }
625
626        async fn execute(&self, args: &Value, _ctx: &ToolContext) -> Result<ToolOutput> {
627            let tasks = args
628                .get("tasks")
629                .and_then(Value::as_array)
630                .map(Vec::len)
631                .unwrap_or(0);
632            Ok(ToolOutput::success(format!("runtime:{tasks}"))
633                .with_metadata(json!({ "runtime_tasks": tasks })))
634        }
635    }
636
637    struct FailingRuntimeTool;
638
639    #[async_trait]
640    impl Tool for FailingRuntimeTool {
641        fn name(&self) -> &str {
642            "runtime"
643        }
644
645        fn description(&self) -> &str {
646            "Failing OS runtime tool for DynamicWorkflowRuntime tests."
647        }
648
649        fn parameters(&self) -> Value {
650            json!({ "type": "object" })
651        }
652
653        async fn execute(&self, _args: &Value, _ctx: &ToolContext) -> Result<ToolOutput> {
654            Ok(ToolOutput::error("runtime unavailable"))
655        }
656    }
657
658    #[tokio::test]
659    async fn dynamic_workflow_tool_runs_ptc_step_through_a3s_flow() {
660        let dir = tempfile::tempdir().unwrap();
661        tokio::fs::write(dir.path().join("fixture.txt"), "hello from fixture")
662            .await
663            .unwrap();
664        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
665        register_dynamic_workflow(executor.registry());
666
667        let source = r#"
668async function run(ctx, inputs) {
669  if (inputs.kind === "workflow") {
670    const read = inputs.step_outputs.read_fixture;
671    if (read) {
672      return { type: "complete", output: { text: read.output } };
673    }
674    return {
675      type: "schedule_step",
676      step_id: "read_fixture",
677      step_name: "read_fixture",
678      input: { path: inputs.input.path },
679      retry: { max_attempts: 1, delay_ms: 0 },
680    };
681  }
682
683  if (inputs.kind === "step" && inputs.step_name === "read_fixture") {
684    return await ctx.read(inputs.input.path);
685  }
686
687  return { error: "unknown invocation" };
688}
689"#;
690
691        let result = executor
692            .execute(
693                DYNAMIC_WORKFLOW_TOOL,
694                &json!({
695                    "source": source,
696                    "input": { "path": "fixture.txt" },
697                    "run_id": "test-dynamic-workflow",
698                    "allowed_tools": ["read"],
699                }),
700            )
701            .await
702            .unwrap();
703
704        assert_eq!(result.exit_code, 0, "{}", result.output);
705        assert!(
706            result.output.contains("hello from fixture"),
707            "{}",
708            result.output
709        );
710        let metadata = result.metadata.unwrap();
711        assert_eq!(
712            metadata["dynamic_workflow"]["run_id"],
713            "test-dynamic-workflow"
714        );
715        assert_eq!(metadata["dynamic_workflow"]["status"], "Completed");
716        assert_eq!(
717            metadata["dynamic_workflow"]["snapshot"]["steps"]["read_fixture"]["status"],
718            "completed"
719        );
720    }
721
722    #[tokio::test]
723    async fn dynamic_workflow_emits_agent_progress_events() {
724        let dir = tempfile::tempdir().unwrap();
725        tokio::fs::write(dir.path().join("fixture.txt"), "hello from fixture")
726            .await
727            .unwrap();
728        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
729        register_dynamic_workflow(executor.registry());
730        let (tx, mut rx) = broadcast::channel(64);
731        let ctx = ToolContext::new(dir.path().to_path_buf())
732            .with_session_id("progress-session")
733            .with_agent_event_tx(tx);
734
735        let source = r#"
736async function run(ctx, inputs) {
737  if (inputs.kind === "workflow") {
738    const read = inputs.step_outputs.read_fixture;
739    if (read) {
740      return { type: "complete", output: { text: read.output } };
741    }
742    return {
743      type: "schedule_step",
744      step_id: "read_fixture",
745      step_name: "read_fixture",
746      input: { path: inputs.input.path, description: "Read fixture" },
747      retry: { max_attempts: 1, delay_ms: 0 },
748    };
749  }
750
751  if (inputs.kind === "step" && inputs.step_name === "read_fixture") {
752    return await ctx.read(inputs.input.path);
753  }
754
755  return { error: "unknown invocation" };
756}
757"#;
758
759        let result = executor
760            .execute_with_context(
761                DYNAMIC_WORKFLOW_TOOL,
762                &json!({
763                    "source": source,
764                    "input": { "path": "fixture.txt" },
765                    "run_id": "test-dynamic-workflow-progress",
766                    "allowed_tools": ["read"],
767                }),
768                &ctx,
769            )
770            .await
771            .unwrap();
772
773        assert_eq!(result.exit_code, 0, "{}", result.output);
774        let mut events = Vec::new();
775        while let Ok(event) = rx.try_recv() {
776            events.push(event);
777        }
778
779        assert!(
780            events
781                .iter()
782                .any(|event| matches!(event, AgentEvent::PlanningStart { .. })),
783            "{events:?}"
784        );
785        assert!(
786            events.iter().any(|event| matches!(
787                event,
788                AgentEvent::TaskUpdated { tasks, .. }
789                    if tasks.iter().any(|task| task.id == "read_fixture")
790            )),
791            "{events:?}"
792        );
793        assert!(
794            events.iter().any(|event| matches!(
795                event,
796                AgentEvent::StepStart { step_id, .. } if step_id == "read_fixture"
797            )),
798            "{events:?}"
799        );
800        assert!(
801            events.iter().any(|event| matches!(
802                event,
803                AgentEvent::StepEnd { step_id, status, .. }
804                    if step_id == "read_fixture" && *status == TaskStatus::Completed
805            )),
806            "{events:?}"
807        );
808    }
809
810    #[tokio::test]
811    async fn dynamic_workflow_step_can_call_host_parallel_task_without_ptc_parallel_task() {
812        let dir = tempfile::tempdir().unwrap();
813        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
814        executor.register_dynamic_tool(Arc::new(FakeParallelTaskTool));
815        register_dynamic_workflow(executor.registry());
816
817        let source = r#"
818async function run(ctx, inputs) {
819  if (inputs.kind === "workflow") {
820    const fanout = inputs.step_outputs.fanout;
821    if (fanout) {
822      return { type: "complete", output: { fanout } };
823    }
824    return {
825      type: "schedule_step",
826      step_id: "fanout",
827      step_name: "parallel_task",
828      input: {
829        tasks: [
830          { agent: "explore", description: "alpha", prompt: "research alpha" },
831          { agent: "explore", description: "beta", prompt: "research beta" },
832        ],
833      },
834    };
835  }
836
837  return { error: "ptc step handler should not run for parallel_task" };
838}
839"#;
840
841        let result = executor
842            .execute(
843                DYNAMIC_WORKFLOW_TOOL,
844                &json!({
845                    "source": source,
846                    "run_id": "test-dynamic-workflow-parallel-step",
847                    "allowed_tools": [],
848                }),
849            )
850            .await
851            .unwrap();
852
853        assert_eq!(result.exit_code, 0, "{}", result.output);
854        assert!(result.output.contains("parallel:2"), "{}", result.output);
855        let metadata = result.metadata.unwrap();
856        assert_eq!(metadata["dynamic_workflow"]["status"], "Completed");
857        let step = &metadata["dynamic_workflow"]["snapshot"]["steps"]["fanout"];
858        assert_eq!(step["status"], "completed");
859        assert_eq!(step["output"]["tool"], PARALLEL_TASK_TOOL);
860        assert_eq!(step["output"]["metadata"]["task_count"], 2);
861    }
862
863    #[tokio::test]
864    async fn dynamic_workflow_ptc_step_can_call_login_registered_runtime_tool_by_default() {
865        let dir = tempfile::tempdir().unwrap();
866        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
867        executor.register_dynamic_tool(Arc::new(FakeRuntimeTool));
868        register_dynamic_workflow(executor.registry());
869
870        let source = r#"
871async function run(ctx, inputs) {
872  if (inputs.kind === "workflow") {
873    const runtime = inputs.step_outputs.runtime_fanout;
874    if (runtime) {
875      return { type: "complete", output: { runtime } };
876    }
877    return {
878      type: "schedule_step",
879      step_id: "runtime_fanout",
880      step_name: "runtime_fanout",
881      input: {
882        worker: "research-worker",
883        tasks: ["alpha", "beta", "gamma"],
884      },
885    };
886  }
887
888  if (inputs.kind === "step" && inputs.step_name === "runtime_fanout") {
889    return await ctx.tool("runtime", inputs.input);
890  }
891
892  return { error: "unknown invocation" };
893}
894"#;
895
896        let result = executor
897            .execute(
898                DYNAMIC_WORKFLOW_TOOL,
899                &json!({
900                    "source": source,
901                    "run_id": "test-dynamic-workflow-runtime-step",
902                }),
903            )
904            .await
905            .unwrap();
906
907        assert_eq!(result.exit_code, 0, "{}", result.output);
908        assert!(result.output.contains("runtime:3"), "{}", result.output);
909        let metadata = result.metadata.unwrap();
910        let step = &metadata["dynamic_workflow"]["snapshot"]["steps"]["runtime_fanout"];
911        assert_eq!(step["status"], "completed");
912        assert_eq!(step["output"]["name"], "runtime");
913        assert_eq!(step["output"]["metadata"]["runtime_tasks"], 3);
914    }
915
916    #[tokio::test]
917    async fn dynamic_workflow_tool_returns_error_when_runtime_step_fails() {
918        let dir = tempfile::tempdir().unwrap();
919        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
920        executor.register_dynamic_tool(Arc::new(FailingRuntimeTool));
921        register_dynamic_workflow(executor.registry());
922
923        let source = r#"
924async function run(ctx, inputs) {
925  if (inputs.kind === "workflow") {
926    const runtime = inputs.step_outputs.runtime_fanout;
927    if (runtime) {
928      return { type: "complete", output: { runtime } };
929    }
930    return {
931      type: "schedule_step",
932      step_id: "runtime_fanout",
933      step_name: "runtime_fanout",
934      input: { worker: "research-worker", tasks: ["alpha"] },
935    };
936  }
937
938  if (inputs.kind === "step" && inputs.step_name === "runtime_fanout") {
939    const result = await ctx.tool("runtime", inputs.input);
940    if (result.exitCode !== 0) {
941      throw new Error(result.output || "runtime failed");
942    }
943    return result;
944  }
945
946  return { error: "unknown invocation" };
947}
948"#;
949
950        let result = executor
951            .execute(
952                DYNAMIC_WORKFLOW_TOOL,
953                &json!({
954                    "source": source,
955                    "run_id": "test-dynamic-workflow-runtime-step-fails",
956                }),
957            )
958            .await
959            .unwrap();
960
961        assert_ne!(result.exit_code, 0, "{}", result.output);
962        assert!(
963            result.output.contains("runtime unavailable"),
964            "{}",
965            result.output
966        );
967        let metadata = result.metadata.unwrap();
968        assert_eq!(metadata["dynamic_workflow"]["status"], "Failed");
969        let step = &metadata["dynamic_workflow"]["snapshot"]["steps"]["runtime_fanout"];
970        assert_eq!(step["status"], "failed");
971    }
972
973    #[test]
974    fn default_allowed_tools_exclude_recursive_program_and_dynamic_workflow_tools() {
975        let dir = tempfile::tempdir().unwrap();
976        let executor = ToolExecutor::new(dir.path().to_string_lossy().to_string());
977        register_dynamic_workflow(executor.registry());
978
979        let tools = default_allowed_tools(executor.registry());
980
981        assert!(!tools.contains(&PROGRAM_TOOL.to_string()));
982        assert!(!tools.contains(&DYNAMIC_WORKFLOW_TOOL.to_string()));
983        assert!(!tools.contains(&PARALLEL_TASK_TOOL.to_string()));
984        assert!(tools.contains(&"read".to_string()));
985    }
986}
a3s_code_core/dynamic_workflow.rs

a3s_code_core/
dynamic_workflow.rs
