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::{
8    registry_tool_invoker, Tool, ToolContext, ToolInvocation, ToolInvoker, ToolOutput,
9    ToolRegistry, ToolResult,
10};
11use crate::{
12    agent::AgentEvent,
13    planning::{Complexity, ExecutionPlan, Task, TaskStatus},
14};
15use a3s_flow::{
16    FlowEngine, FlowEvent, FlowEventEnvelope, FlowEventObserver, FlowEventStore, FlowRuntime,
17    InMemoryEventStore, LocalFileEventStore, RuntimeCommand, StepInvocation, WorkflowInvocation,
18    WorkflowRunStatus, WorkflowSpec,
19};
20use anyhow::{Context, Result};
21use async_trait::async_trait;
22use serde::{Deserialize, Serialize};
23use serde_json::{json, Map, Value};
24use std::collections::BTreeSet;
25use std::path::{Path, PathBuf};
26use std::sync::Arc;
27use tokio::sync::{broadcast, Mutex};
28
29const DYNAMIC_WORKFLOW_TOOL: &str = "dynamic_workflow";
30const PROGRAM_TOOL: &str = "program";
31const PARALLEL_TASK_TOOL: &str = "parallel_task";
32
33/// Project-relative directory used for durable dynamic workflow history.
34pub const DYNAMIC_WORKFLOW_STORE_RELATIVE_PATH: &str = ".a3s/workflow";
35
36/// Resolve the durable dynamic workflow history directory for a local workspace.
37pub fn dynamic_workflow_store_path(workspace_root: impl AsRef<Path>) -> PathBuf {
38    workspace_root
39        .as_ref()
40        .join(DYNAMIC_WORKFLOW_STORE_RELATIVE_PATH)
41}
42
43/// Limits forwarded to the underlying PTC `program` tool.
44#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
45#[serde(rename_all = "camelCase")]
46pub struct DynamicWorkflowScriptLimits {
47    #[serde(skip_serializing_if = "Option::is_none")]
48    pub timeout_ms: Option<u64>,
49    #[serde(skip_serializing_if = "Option::is_none")]
50    pub max_tool_calls: Option<usize>,
51    #[serde(skip_serializing_if = "Option::is_none")]
52    pub max_output_bytes: Option<usize>,
53}
54
55/// Runs A3S Flow workflow and step invocations through a sandboxed PTC script.
56#[derive(Clone)]
57pub struct DynamicWorkflowRuntime {
58    invoker: Arc<dyn ToolInvoker>,
59    context: ToolContext,
60    source: Arc<str>,
61    allowed_tools: Vec<String>,
62    limits: DynamicWorkflowScriptLimits,
63}
64
65impl DynamicWorkflowRuntime {
66    pub fn new(
67        registry: Arc<ToolRegistry>,
68        context: ToolContext,
69        source: impl Into<String>,
70    ) -> Self {
71        let allowed_tools = default_allowed_tools(&registry);
72        // Session/agent callers install the governed gateway in ToolContext.
73        // The raw registry adapter is retained only for explicit low-level
74        // callers that construct this public runtime outside an AgentSession.
75        let invoker = context
76            .tool_invoker()
77            .unwrap_or_else(|| registry_tool_invoker(registry));
78        Self {
79            invoker,
80            context,
81            source: Arc::from(source.into()),
82            allowed_tools,
83            limits: DynamicWorkflowScriptLimits::default(),
84        }
85    }
86
87    pub fn with_allowed_tools(mut self, allowed_tools: impl IntoIterator<Item = String>) -> Self {
88        self.allowed_tools = sanitize_allowed_tools(allowed_tools);
89        self
90    }
91
92    pub fn with_limits(mut self, limits: DynamicWorkflowScriptLimits) -> Self {
93        self.limits = limits;
94        self
95    }
96
97    async fn run_script(&self, payload: Value) -> a3s_flow::Result<ToolResult> {
98        let mut args = json!({
99            "type": "script",
100            "language": "javascript",
101            "source": self.source.as_ref(),
102            "inputs": payload,
103            "allowed_tools": self.allowed_tools,
104        });
105        if let Some(object) = args.as_object_mut() {
106            if let Ok(Value::Object(limits)) = serde_json::to_value(&self.limits) {
107                if !limits.is_empty() {
108                    object.insert("limits".to_string(), Value::Object(limits));
109                }
110            }
111        }
112
113        let result = self
114            .invoker
115            .invoke(ToolInvocation::nested(PROGRAM_TOOL, args), &self.context)
116            .await;
117        if result.exit_code != 0 {
118            return Err(a3s_flow::FlowError::Runtime(result.output));
119        }
120        Ok(result)
121    }
122
123    async fn run_tool_step(&self, tool_name: &str, args: Value) -> a3s_flow::Result<Value> {
124        let result = self
125            .invoker
126            .invoke(
127                ToolInvocation::nested(tool_name.to_string(), args),
128                &self.context,
129            )
130            .await;
131        if result.exit_code != 0 {
132            return Err(a3s_flow::FlowError::Runtime(result.output));
133        }
134        Ok(json!({
135            "tool": result.name,
136            "output": result.output,
137            "exit_code": result.exit_code,
138            "metadata": result.metadata,
139        }))
140    }
141}
142
143#[async_trait]
144impl FlowRuntime for DynamicWorkflowRuntime {
145    async fn run_workflow(
146        &self,
147        invocation: WorkflowInvocation,
148    ) -> a3s_flow::Result<RuntimeCommand> {
149        let payload = invocation_payload("workflow", &invocation.run_id, &invocation.history)
150            .with("input", invocation.input);
151        let result = self.run_script(payload.into_value()).await?;
152        serde_json::from_value(script_result(&result)?).map_err(a3s_flow::FlowError::from)
153    }
154
155    async fn run_step(&self, invocation: StepInvocation) -> a3s_flow::Result<Value> {
156        if invocation.step_name == PARALLEL_TASK_TOOL {
157            return self
158                .run_tool_step(PARALLEL_TASK_TOOL, invocation.input)
159                .await;
160        }
161
162        let payload = invocation_payload("step", &invocation.run_id, &invocation.history)
163            .with("step_id", invocation.step_id)
164            .with("step_name", invocation.step_name)
165            .with("input", invocation.input);
166        let result = self.run_script(payload.into_value()).await?;
167        script_result(&result)
168    }
169}
170
171struct WorkflowProgressState {
172    tasks: Vec<Task>,
173}
174
175impl WorkflowProgressState {
176    fn new() -> Self {
177        Self { tasks: Vec::new() }
178    }
179
180    fn upsert_step(
181        &mut self,
182        step_id: &str,
183        step_name: &str,
184        input: Option<&Value>,
185        status: TaskStatus,
186    ) {
187        let content = workflow_step_description(step_id, step_name, input);
188        if let Some(task) = self.tasks.iter_mut().find(|task| task.id == step_id) {
189            task.content = content;
190            task.status = status;
191            task.tool = Some(step_name.to_string());
192        } else {
193            self.tasks
194                .push(Task::new(step_id.to_string(), content).with_tool(step_name));
195            if let Some(task) = self.tasks.last_mut() {
196                task.status = status;
197            }
198        }
199    }
200
201    fn mark_status(&mut self, step_id: &str, status: TaskStatus) {
202        if let Some(task) = self.tasks.iter_mut().find(|task| task.id == step_id) {
203            task.status = status;
204        }
205    }
206
207    fn step_position(&self, step_id: &str) -> (usize, usize) {
208        let total = self.tasks.len().max(1);
209        let number = self
210            .tasks
211            .iter()
212            .position(|task| task.id == step_id)
213            .map(|idx| idx + 1)
214            .unwrap_or(total);
215        (number, total)
216    }
217
218    fn step_description(&self, step_id: &str) -> String {
219        self.tasks
220            .iter()
221            .find(|task| task.id == step_id)
222            .map(|task| task.content.clone())
223            .unwrap_or_else(|| step_id.to_string())
224    }
225}
226
227struct AgentEventFlowObserver {
228    tx: broadcast::Sender<AgentEvent>,
229    session_id: String,
230    state: Mutex<WorkflowProgressState>,
231}
232
233impl AgentEventFlowObserver {
234    fn new(tx: broadcast::Sender<AgentEvent>, session_id: String) -> Self {
235        Self {
236            tx,
237            session_id,
238            state: Mutex::new(WorkflowProgressState::new()),
239        }
240    }
241
242    fn emit_task_update(&self, tasks: &[Task]) {
243        let _ = self.tx.send(AgentEvent::TaskUpdated {
244            session_id: self.session_id.clone(),
245            tasks: tasks.to_vec(),
246        });
247    }
248}
249
250#[async_trait]
251impl FlowEventObserver for AgentEventFlowObserver {
252    async fn observe(&self, envelope: FlowEventEnvelope) {
253        match envelope.event {
254            FlowEvent::RunStarted => {
255                let _ = self.tx.send(AgentEvent::PlanningStart {
256                    prompt: "dynamic_workflow".to_string(),
257                });
258            }
259            FlowEvent::StepCreated {
260                step_id,
261                step_name,
262                input,
263                ..
264            } => {
265                let mut state = self.state.lock().await;
266                state.upsert_step(&step_id, &step_name, Some(&input), TaskStatus::Pending);
267                self.emit_task_update(&state.tasks);
268                let mut plan = ExecutionPlan::new("dynamic workflow", Complexity::Medium);
269                for task in state.tasks.iter().cloned() {
270                    plan.add_step(task);
271                }
272                let _ = self.tx.send(AgentEvent::PlanningEnd {
273                    estimated_steps: plan.steps.len(),
274                    plan,
275                });
276            }
277            FlowEvent::StepStarted { step_id, .. } => {
278                let mut state = self.state.lock().await;
279                state.mark_status(&step_id, TaskStatus::InProgress);
280                self.emit_task_update(&state.tasks);
281                let (step_number, total_steps) = state.step_position(&step_id);
282                let _ = self.tx.send(AgentEvent::StepStart {
283                    description: state.step_description(&step_id),
284                    step_id,
285                    step_number,
286                    total_steps,
287                });
288            }
289            FlowEvent::StepCompleted { step_id, .. } => {
290                let mut state = self.state.lock().await;
291                state.mark_status(&step_id, TaskStatus::Completed);
292                self.emit_task_update(&state.tasks);
293                let (step_number, total_steps) = state.step_position(&step_id);
294                let _ = self.tx.send(AgentEvent::StepEnd {
295                    step_id,
296                    status: TaskStatus::Completed,
297                    step_number,
298                    total_steps,
299                });
300            }
301            FlowEvent::StepRetrying { step_id, .. } => {
302                let mut state = self.state.lock().await;
303                state.mark_status(&step_id, TaskStatus::InProgress);
304                self.emit_task_update(&state.tasks);
305            }
306            FlowEvent::StepFailed { step_id, .. } => {
307                let mut state = self.state.lock().await;
308                state.mark_status(&step_id, TaskStatus::Failed);
309                self.emit_task_update(&state.tasks);
310                let (step_number, total_steps) = state.step_position(&step_id);
311                let _ = self.tx.send(AgentEvent::StepEnd {
312                    step_id,
313                    status: TaskStatus::Failed,
314                    step_number,
315                    total_steps,
316                });
317            }
318            FlowEvent::RunFailed { .. } => {
319                let mut state = self.state.lock().await;
320                for task in &mut state.tasks {
321                    if task.status.is_active() {
322                        task.status = TaskStatus::Failed;
323                    }
324                }
325                self.emit_task_update(&state.tasks);
326            }
327            FlowEvent::RunCancelled { .. } => {
328                let mut state = self.state.lock().await;
329                for task in &mut state.tasks {
330                    if task.status.is_active() {
331                        task.status = TaskStatus::Cancelled;
332                    }
333                }
334                self.emit_task_update(&state.tasks);
335            }
336            _ => {}
337        }
338    }
339}
340
341fn workflow_step_description(step_id: &str, step_name: &str, input: Option<&Value>) -> String {
342    if step_name == PARALLEL_TASK_TOOL {
343        let count = input
344            .and_then(|value| value.get("tasks"))
345            .and_then(Value::as_array)
346            .map(Vec::len)
347            .unwrap_or(0);
348        if count > 0 {
349            return format!("Fan out {count} parallel subagent task(s)");
350        }
351    }
352
353    input
354        .and_then(|value| value.get("description").or_else(|| value.get("title")))
355        .and_then(Value::as_str)
356        .map(ToString::to_string)
357        .unwrap_or_else(|| {
358            if step_name == step_id {
359                step_id.to_string()
360            } else {
361                format!("{step_name}: {step_id}")
362            }
363        })
364}
365
366/// Model-visible tool that executes a dynamic workflow through A3S Flow.
367pub struct DynamicWorkflowTool {
368    registry: Arc<ToolRegistry>,
369}
370
371impl DynamicWorkflowTool {
372    pub fn new(registry: Arc<ToolRegistry>) -> Self {
373        Self { registry }
374    }
375}
376
377#[async_trait]
378impl Tool for DynamicWorkflowTool {
379    fn name(&self) -> &str {
380        DYNAMIC_WORKFLOW_TOOL
381    }
382
383    fn description(&self) -> &str {
384        "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."
385    }
386
387    fn parameters(&self) -> Value {
388        json!({
389            "type": "object",
390            "additionalProperties": false,
391            "properties": {
392                "source": {
393                    "type": "string",
394                    "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."
395                },
396                "input": {
397                    "type": "object",
398                    "description": "Initial workflow input."
399                },
400                "run_id": {
401                    "type": "string",
402                    "description": "Optional durable run id. Reusing it with the same source and input is idempotent."
403                },
404                "allowed_tools": {
405                    "type": "array",
406                    "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.",
407                    "items": { "type": "string" }
408                },
409                "limits": {
410                    "type": "object",
411                    "additionalProperties": false,
412                    "properties": {
413                        "timeoutMs": { "type": "integer", "minimum": 1 },
414                        "maxToolCalls": { "type": "integer", "minimum": 1 },
415                        "maxOutputBytes": { "type": "integer", "minimum": 1 }
416                    }
417                }
418            },
419            "required": ["source"]
420        })
421    }
422
423    async fn execute(&self, args: &Value, ctx: &ToolContext) -> Result<ToolOutput> {
424        let Some(source) = args.get("source").and_then(Value::as_str) else {
425            return Ok(ToolOutput::error("dynamic_workflow requires source"));
426        };
427        let input = args.get("input").cloned().unwrap_or_else(|| json!({}));
428        let allowed_tools = args
429            .get("allowed_tools")
430            .and_then(Value::as_array)
431            .map(|items| {
432                items
433                    .iter()
434                    .filter_map(Value::as_str)
435                    .map(ToString::to_string)
436                    .collect::<Vec<_>>()
437            })
438            .unwrap_or_else(|| default_allowed_tools(&self.registry));
439        let limits = args
440            .get("limits")
441            .cloned()
442            .and_then(|value| serde_json::from_value(value).ok())
443            .unwrap_or_default();
444
445        let runtime = Arc::new(
446            DynamicWorkflowRuntime::new(Arc::clone(&self.registry), ctx.clone(), source)
447                .with_allowed_tools(allowed_tools)
448                .with_limits(limits),
449        );
450        let requested_run_id = args.get("run_id").and_then(Value::as_str);
451        let store = match flow_store_for_context(ctx, requested_run_id).await {
452            Ok(store) => store,
453            Err(error) => return Ok(ToolOutput::error(error.to_string())),
454        };
455        let engine = match ctx.agent_event_tx.clone() {
456            Some(tx) => FlowEngine::builder(runtime)
457                .with_store(store)
458                .with_observer(Arc::new(AgentEventFlowObserver::new(
459                    tx,
460                    ctx.session_id.clone().unwrap_or_default(),
461                )))
462                .build(),
463            None => FlowEngine::new(store, runtime),
464        };
465        let source_hash = source_hash(source);
466        let spec = WorkflowSpec::rust_embedded(
467            "a3s-code.dynamic-workflow",
468            source_hash.as_str(),
469            "ptc",
470            "run",
471        );
472
473        let run_id = match requested_run_id {
474            Some(run_id) => match engine.start_with_id(run_id, spec, input).await {
475                Ok(run_id) => run_id,
476                Err(err) => return Ok(ToolOutput::error(err.to_string())),
477            },
478            None => match engine.start(spec, input).await {
479                Ok(run_id) => run_id,
480                Err(err) => return Ok(ToolOutput::error(err.to_string())),
481            },
482        };
483
484        let snapshot = match engine.snapshot(&run_id).await {
485            Ok(snapshot) => snapshot,
486            Err(err) => return Ok(ToolOutput::error(err.to_string())),
487        };
488        let history = match engine.history(&run_id).await {
489            Ok(history) => history,
490            Err(err) => return Ok(ToolOutput::error(err.to_string())),
491        };
492
493        let output = match &snapshot.output {
494            Some(output) => {
495                serde_json::to_string_pretty(output).unwrap_or_else(|_| output.to_string())
496            }
497            None => snapshot
498                .error
499                .clone()
500                .unwrap_or_else(|| format!("workflow status: {:?}", snapshot.status)),
501        };
502
503        let status = snapshot.status;
504        let metadata = json!({
505            "dynamic_workflow": {
506                "run_id": run_id,
507                "status": format!("{:?}", snapshot.status),
508                "last_sequence": snapshot.last_sequence,
509                "source_hash": source_hash,
510                "snapshot": snapshot,
511                "history": history,
512            }
513        });
514        let output = match status {
515            WorkflowRunStatus::Completed => ToolOutput::success(output),
516            WorkflowRunStatus::Failed | WorkflowRunStatus::Cancelled => ToolOutput::error(output),
517            _ => ToolOutput::error(format!(
518                "dynamic_workflow ended without a terminal result: {status:?}; {output}"
519            )),
520        };
521
522        Ok(output.with_metadata(metadata))
523    }
524}
525
526pub fn register_dynamic_workflow(registry: &Arc<ToolRegistry>) {
527    registry.register(Arc::new(DynamicWorkflowTool::new(Arc::clone(registry))));
528}
529
530async fn flow_store_for_context(
531    ctx: &ToolContext,
532    requested_run_id: Option<&str>,
533) -> Result<Arc<dyn FlowEventStore>> {
534    match ctx.workspace_services.local_root() {
535        Some(root) => {
536            let store = dynamic_workflow_store_path(root);
537            validate_dynamic_workflow_directory(&root.join(".a3s"), ".a3s").await?;
538            validate_dynamic_workflow_directory(&store, ".a3s/workflow").await?;
539            if let Some(run_id) = requested_run_id.filter(|run_id| safe_workflow_run_id(run_id)) {
540                validate_dynamic_workflow_log(&store.join(format!("{run_id}.jsonl"))).await?;
541            }
542            Ok(Arc::new(LocalFileEventStore::new(store)))
543        }
544        None => Ok(Arc::new(InMemoryEventStore::new())),
545    }
546}
547
548async fn validate_dynamic_workflow_directory(path: &Path, label: &str) -> Result<()> {
549    match tokio::fs::symlink_metadata(path).await {
550        Ok(metadata) if metadata.file_type().is_symlink() => {
551            anyhow::bail!("refusing to use symlinked dynamic workflow directory {label}")
552        }
553        Ok(metadata) if !metadata.is_dir() => {
554            anyhow::bail!("dynamic workflow path {label} exists but is not a directory")
555        }
556        Ok(_) => Ok(()),
557        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
558        Err(error) => Err(error).with_context(|| format!("inspect dynamic workflow path {label}")),
559    }
560}
561
562async fn validate_dynamic_workflow_log(path: &Path) -> Result<()> {
563    match tokio::fs::symlink_metadata(path).await {
564        Ok(metadata) if metadata.file_type().is_symlink() => anyhow::bail!(
565            "refusing to read or append symlinked dynamic workflow history {}",
566            path.display()
567        ),
568        Ok(metadata) if !metadata.is_file() => anyhow::bail!(
569            "dynamic workflow history path {} exists but is not a file",
570            path.display()
571        ),
572        Ok(_) => Ok(()),
573        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
574        Err(error) => Err(error)
575            .with_context(|| format!("inspect dynamic workflow history {}", path.display())),
576    }
577}
578
579fn safe_workflow_run_id(run_id: &str) -> bool {
580    !run_id.is_empty()
581        && run_id
582            .chars()
583            .all(|ch| ch.is_ascii_alphanumeric() || ch == '-' || ch == '_')
584}
585
586struct PayloadBuilder {
587    value: Map<String, Value>,
588}
589
590impl PayloadBuilder {
591    fn with(mut self, key: &str, value: impl Serialize) -> Self {
592        self.value.insert(
593            key.to_string(),
594            serde_json::to_value(value).unwrap_or(Value::Null),
595        );
596        self
597    }
598
599    fn into_value(self) -> Value {
600        Value::Object(self.value)
601    }
602}
603
604fn invocation_payload(kind: &str, run_id: &str, history: &[FlowEventEnvelope]) -> PayloadBuilder {
605    let mut value = Map::new();
606    value.insert("kind".to_string(), json!(kind));
607    value.insert("run_id".to_string(), json!(run_id));
608    value.insert("history".to_string(), json!(history));
609    value.insert("step_outputs".to_string(), completed_step_outputs(history));
610    value.insert("step_failures".to_string(), failed_step_outputs(history));
611    PayloadBuilder { value }
612}
613
614fn completed_step_outputs(history: &[FlowEventEnvelope]) -> Value {
615    let mut outputs = Map::new();
616    for envelope in history {
617        if let FlowEvent::StepCompleted { step_id, output } = &envelope.event {
618            outputs.insert(step_id.clone(), output.clone());
619        }
620    }
621    Value::Object(outputs)
622}
623
624fn failed_step_outputs(history: &[FlowEventEnvelope]) -> Value {
625    let mut outputs = Map::new();
626    for envelope in history {
627        if let FlowEvent::StepFailed {
628            step_id,
629            attempt,
630            error,
631        } = &envelope.event
632        {
633            outputs.insert(
634                step_id.clone(),
635                json!({
636                    "attempt": attempt,
637                    "error": error,
638                }),
639            );
640        }
641    }
642    Value::Object(outputs)
643}
644
645fn script_result(result: &ToolResult) -> a3s_flow::Result<Value> {
646    result
647        .metadata
648        .as_ref()
649        .and_then(|metadata| metadata.get("script_result"))
650        .cloned()
651        .ok_or_else(|| {
652            a3s_flow::FlowError::Runtime(
653                "PTC program result did not include script_result metadata".to_string(),
654            )
655        })
656}
657
658fn default_allowed_tools(registry: &ToolRegistry) -> Vec<String> {
659    sanitize_allowed_tools(registry.list())
660}
661
662fn sanitize_allowed_tools(items: impl IntoIterator<Item = String>) -> Vec<String> {
663    let mut tools = items.into_iter().collect::<BTreeSet<_>>();
664    tools.remove(PROGRAM_TOOL);
665    tools.remove(DYNAMIC_WORKFLOW_TOOL);
666    tools.remove(PARALLEL_TASK_TOOL);
667    tools.into_iter().collect()
668}
669
670fn source_hash(source: &str) -> String {
671    sha256::digest(source.as_bytes())
672}
673
674#[cfg(test)]
675#[path = "dynamic_workflow/tests.rs"]
676mod tests;
a3s_code_core/dynamic_workflow.rs

a3s_code_core/
dynamic_workflow.rs
