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