1use std::time::{Duration, Instant};
2
3use anyhow::{Context, Result};
4use async_trait::async_trait;
5use serde_json::Value;
6use tokio::task::JoinSet;
7use tokio::time::MissedTickBehavior;
8
9use crate::workflow_definition::HostExecutorNode;
10use crate::workflow_definition::SubagentNode;
11use crate::workflow_orchestrator::OrchestratorAction;
12use crate::{
13 EventDraft, EventLog, RunSnapshotDTO, RunStatus, WorkflowActor, WorkflowDefinition,
14 WorkflowNode, decide_next_actions, read_run_snapshot,
15};
16
17mod completion;
18mod dispatch;
19mod helpers;
20pub mod r#loop;
21
22#[cfg(test)]
23mod test_common;
24#[cfg(test)]
25mod tests_cr_loop;
26#[cfg(test)]
27mod tests_loop;
28#[cfg(test)]
29mod tests_real_cr;
30#[cfg(test)]
31mod tests_recovery;
32#[cfg(test)]
33mod tests_run_loop;
34#[cfg(test)]
35mod tests_run_tick;
36
37pub use completion::{
39 complete_node_failed, complete_node_succeeded, complete_run_failed, complete_run_succeeded,
40};
41pub use dispatch::{dispatch_gate, dispatch_work};
42pub use helpers::{derive_workflow_idempotency_key, get_host_executor_provider_meta};
43pub use r#loop::{finish_loop, finish_loop_iteration, start_loop, start_loop_iteration};
44
45use helpers::write_json_blob;
46
47#[derive(Debug, Clone)]
48pub struct WorkflowRuntimeContext {
49 pub log: EventLog,
50 pub def: WorkflowDefinition,
51 pub runs_base_dir: std::path::PathBuf,
52}
53
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct WorkflowDispatchSession {
56 pub session_id: String,
57 pub bot_name: String,
58 pub started_at: u64,
59 pub ended_at: Option<u64>,
60 pub cli_session_id: Option<String>,
61 pub lark_app_id: Option<String>,
62 pub cli_id: Option<String>,
63 pub working_dir: Option<String>,
64 pub web_port: Option<u16>,
65 pub log_path: Option<String>,
66}
67
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub enum WorkflowDispatchOutcome {
70 Succeeded {
71 output: Value,
72 session: Option<WorkflowDispatchSession>,
73 },
74 Failed {
75 error_code: String,
76 error_class: String,
77 error_message: String,
78 session: Option<WorkflowDispatchSession>,
79 },
80 Cancelled {
81 cancel_origin_event_id: String,
82 session: Option<WorkflowDispatchSession>,
83 },
84}
85
86#[derive(Debug, Clone)]
87pub struct WorkflowDispatchRun<'a> {
88 pub run_id: &'a str,
89 pub workflow_id: &'a str,
90 pub revision_id: &'a str,
91 pub activity_id: &'a str,
92 pub attempt_id: &'a str,
93 pub node_id: &'a str,
94}
95
96#[derive(Debug, Clone)]
101pub struct HostExecutorPrepareResult {
102 pub parsed_input: Value,
104 pub canonical_input: Value,
107 pub provider: String,
109 pub idempotency_ttl_ms: u64,
111}
112
113#[async_trait]
114pub trait WorkflowExecutionHooks {
115 async fn execute_subagent(
116 &mut self,
117 ctx: WorkflowDispatchRun<'_>,
118 node: &SubagentNode,
119 resolved_prompt: String,
120 ) -> Result<WorkflowDispatchOutcome>;
121
122 async fn execute_host_executor(
123 &mut self,
124 ctx: WorkflowDispatchRun<'_>,
125 node: &HostExecutorNode,
126 parsed_input: Value,
128 ) -> Result<WorkflowDispatchOutcome>;
129
130 fn prepare_host_executor(
139 &self,
140 executor_name: &str,
141 resolved_input: &Value,
142 ) -> Result<HostExecutorPrepareResult> {
143 let (provider, idempotency_ttl_ms) = get_host_executor_provider_meta(executor_name);
144 Ok(HostExecutorPrepareResult {
145 parsed_input: resolved_input.clone(),
146 canonical_input: resolved_input.clone(),
147 provider: provider.to_string(),
148 idempotency_ttl_ms,
149 })
150 }
151
152 async fn recover_dangling_effects(
166 &mut self,
167 _log: &mut EventLog,
168 snapshot: &RunSnapshotDTO,
169 ) -> Result<RecoveryResult> {
170 Ok(RecoveryResult {
171 had_progress: false,
172 has_remaining_dangling: !snapshot.dangling.effect_attempted.is_empty(),
173 })
174 }
175
176 async fn on_activities_cancelled(
183 &mut self,
184 _activity_ids: &[String],
185 _node_ids: &[String],
186 _run_id: &str,
187 ) {
188 }
189}
190
191#[derive(Debug, Clone, PartialEq, Eq)]
192pub struct RunTickResult {
193 pub actions: usize,
194 pub snapshot: RunSnapshotDTO,
195}
196
197#[derive(Debug, Clone)]
198struct ScheduledAction {
199 action: OrchestratorAction,
200}
201
202#[derive(Debug, Clone, PartialEq, Eq)]
203pub enum RunLoopStopReason {
204 Terminal,
205 NoProgress,
206 AwaitingWait,
207 MaxTicks,
208}
209
210#[derive(Debug, Clone, PartialEq, Eq)]
212pub struct RecoveryResult {
213 pub had_progress: bool,
215 pub has_remaining_dangling: bool,
217}
218
219#[derive(Debug, Clone, PartialEq, Eq)]
220pub struct RunLoopResult {
221 pub reason: RunLoopStopReason,
222 pub ticks: usize,
223 pub last_snapshot: RunSnapshotDTO,
224}
225
226pub async fn run_tick<H: WorkflowExecutionHooks + Clone + Send + 'static>(
227 rt: &mut WorkflowRuntimeContext,
228 hooks: &mut H,
229 max_concurrency: usize,
230) -> Result<RunTickResult> {
231 let snapshot = read_snapshot(rt).await?;
232 if matches!(
233 snapshot.run.status,
234 RunStatus::Succeeded | RunStatus::Failed | RunStatus::Cancelled
235 ) {
236 return Ok(RunTickResult {
237 actions: 0,
238 snapshot,
239 });
240 }
241
242 if snapshot_has_pending_cancel(&snapshot) {
243 return Ok(RunTickResult {
244 actions: 0,
245 snapshot,
246 });
247 }
248
249 let actions = select_tick_actions(
250 decide_next_actions(&snapshot, &rt.def),
251 &rt.def,
252 max_concurrency,
253 );
254 if actions.is_empty() {
255 return Ok(RunTickResult {
256 actions: 0,
257 snapshot,
258 });
259 }
260
261 let mut join_set: JoinSet<Result<()>> = JoinSet::new();
262 for scheduled in actions.into_iter() {
263 let mut rt_clone = rt.clone();
264 let mut hooks_clone = hooks.clone();
265 join_set.spawn(async move {
266 apply_orchestrator_action(&mut rt_clone, &mut hooks_clone, scheduled.action).await
267 });
268 }
269
270 let mut applied = 0usize;
271 let mut cancel_poll = tokio::time::interval(Duration::from_millis(20));
272 cancel_poll.set_missed_tick_behavior(MissedTickBehavior::Skip);
273 let mut cancel_seen = false;
274 let mut cancel_abort_deadline: Option<Instant> = None;
275
276 while !join_set.is_empty() {
277 tokio::select! {
278 result = join_set.join_next() => {
279 let Some(result) = result else {
280 break;
281 };
282 match result {
283 Ok(Ok(())) => {
284 applied += 1;
285 let snapshot = read_snapshot(rt).await?;
286 if snapshot_has_pending_cancel(&snapshot)
287 || matches!(
288 snapshot.run.status,
289 RunStatus::Succeeded | RunStatus::Failed | RunStatus::Cancelled
290 )
291 {
292 if snapshot_has_pending_cancel(&snapshot) && !cancel_seen {
293 cancel_seen = true;
294 cancel_abort_deadline =
295 Some(Instant::now() + Duration::from_millis(120));
296 }
297 if matches!(
298 snapshot.run.status,
299 RunStatus::Succeeded | RunStatus::Failed | RunStatus::Cancelled
300 ) {
301 join_set.abort_all();
302 return Ok(RunTickResult {
303 actions: applied,
304 snapshot,
305 });
306 }
307 }
308 }
309 Ok(Err(err)) => {
310 if cancel_seen {
311 continue;
312 }
313 join_set.abort_all();
314 return Err(err);
315 }
316 Err(err) => {
317 if cancel_seen {
318 continue;
319 }
320 join_set.abort_all();
321 return Err(anyhow::anyhow!(err));
322 }
323 }
324 }
325 _ = cancel_poll.tick(), if !cancel_seen => {
326 let snapshot = read_snapshot(rt).await?;
327 if snapshot_has_pending_cancel(&snapshot) {
328 cancel_seen = true;
329 cancel_abort_deadline = Some(Instant::now() + Duration::from_millis(120));
330 join_set.abort_all();
331 }
332 }
333 _ = tokio::time::sleep(Duration::from_millis(20)), if cancel_seen => {
334 if let Some(deadline) = cancel_abort_deadline {
335 if Instant::now() >= deadline {
336 join_set.abort_all();
337 }
338 }
339 }
340 }
341 }
342
343 let snapshot = read_snapshot(rt).await?;
344 Ok(RunTickResult {
345 actions: applied,
346 snapshot,
347 })
348}
349
350pub async fn run_loop<H: WorkflowExecutionHooks + Clone + Send + 'static>(
351 rt: &mut WorkflowRuntimeContext,
352 hooks: &mut H,
353 max_ticks: usize,
354 max_concurrency: usize,
355) -> Result<RunLoopResult> {
356 let mut ticks = 0usize;
357 loop {
358 if ticks >= max_ticks {
359 let snapshot = read_snapshot(rt).await?;
360 return Ok(RunLoopResult {
361 reason: RunLoopStopReason::MaxTicks,
362 ticks,
363 last_snapshot: snapshot,
364 });
365 }
366
367 check_pending_cancels(rt, hooks).await?;
368
369 let pre_recovery_snapshot = read_snapshot(rt).await?;
376 if matches!(
377 pre_recovery_snapshot.run.status,
378 RunStatus::Succeeded | RunStatus::Failed | RunStatus::Cancelled
379 ) {
380 return Ok(RunLoopResult {
381 reason: RunLoopStopReason::Terminal,
382 ticks,
383 last_snapshot: pre_recovery_snapshot,
384 });
385 }
386
387 if !pre_recovery_snapshot.dangling.effect_attempted.is_empty() {
388 let recovery = hooks
389 .recover_dangling_effects(&mut rt.log, &pre_recovery_snapshot)
390 .await?;
391 if recovery.had_progress {
392 continue;
395 }
396 }
401
402 if !pre_recovery_snapshot.dangling.wait_resolutions.is_empty() {
408 let had_progress = resolve_wait_terminals(rt, &pre_recovery_snapshot).await?;
409 if had_progress {
410 continue;
411 }
412 }
413
414 let tick = run_tick(rt, hooks, max_concurrency).await?;
415 ticks += 1;
416 if tick.snapshot.run.status == RunStatus::Succeeded
417 || tick.snapshot.run.status == RunStatus::Failed
418 || tick.snapshot.run.status == RunStatus::Cancelled
419 {
420 return Ok(RunLoopResult {
421 reason: RunLoopStopReason::Terminal,
422 ticks,
423 last_snapshot: tick.snapshot,
424 });
425 }
426 if tick.actions == 0 {
427 let has_waits = !tick.snapshot.dangling.waits.is_empty()
428 && tick
429 .snapshot
430 .dangling
431 .waits
432 .iter()
433 .any(|w| !tick.snapshot.dangling.cancels.contains(w));
434 let reason = if has_waits {
435 RunLoopStopReason::AwaitingWait
436 } else {
437 RunLoopStopReason::NoProgress
438 };
439 return Ok(RunLoopResult {
440 reason,
441 ticks,
442 last_snapshot: tick.snapshot,
443 });
444 }
445 }
446}
447
448async fn check_pending_cancels<H: WorkflowExecutionHooks + Send>(
449 rt: &mut WorkflowRuntimeContext,
450 hooks: &mut H,
451) -> Result<()> {
452 let _events = rt.log.read_all()?;
453 let snapshot = read_snapshot(rt).await?;
454 let mut cancelled_activities: Vec<String> = Vec::new();
455 let mut cancelled_nodes: Vec<String> = Vec::new();
456
457 for activity_id in &snapshot.dangling.cancels {
458 cancelled_activities.push(activity_id.clone());
459 let attempt_id = snapshot
460 .activities
461 .iter()
462 .find(|a| &a.activity_id == activity_id)
463 .and_then(|a| a.current_attempt_id.clone())
464 .unwrap_or_else(|| format!("{}-attempt-1", activity_id));
465 let origin = snapshot
466 .run
467 .cancelled_run_intent
468 .as_ref()
469 .map(|i| i.cancel_origin_event_id.clone())
470 .or_else(|| {
471 snapshot
472 .run
473 .cancelled_node_intents
474 .values()
475 .next()
476 .map(|i| i.cancel_origin_event_id.clone())
477 })
478 .unwrap_or_default();
479 let _ = crate::complete_activity_cancel(
480 &mut rt.log,
481 crate::CompleteActivityCancelInput {
482 activity_id: activity_id.clone(),
483 attempt_id,
484 cancel_origin_event_id: origin,
485 },
486 WorkflowActor::Scheduler,
487 )
488 .await;
489 }
490
491 if let Some(ref intent) = snapshot.run.cancelled_run_intent {
492 if snapshot.run.status != RunStatus::Cancelled {
493 let _ = crate::complete_run_cancel(
494 &mut rt.log,
495 crate::CompleteRunCancelInput {
496 cancel_origin_event_id: intent.cancel_origin_event_id.clone(),
497 },
498 WorkflowActor::Scheduler,
499 )
500 .await;
501 }
502 }
503
504 if !snapshot.run.cancelled_node_intents.is_empty() {
505 for (node_id, intent) in &snapshot.run.cancelled_node_intents {
506 cancelled_nodes.push(node_id.clone());
507 let _ = crate::complete_node_cancel(
508 &mut rt.log,
509 crate::CompleteNodeCancelInput {
510 node_id: node_id.clone(),
511 cancel_origin_event_id: intent.cancel_origin_event_id.clone(),
512 },
513 WorkflowActor::Scheduler,
514 )
515 .await;
516 }
517 }
518
519 if !cancelled_activities.is_empty() || !cancelled_nodes.is_empty() {
521 hooks
522 .on_activities_cancelled(&cancelled_activities, &cancelled_nodes, &rt.log.run_id)
523 .await;
524 }
525
526 Ok(())
527}
528
529fn snapshot_has_pending_cancel(snapshot: &RunSnapshotDTO) -> bool {
530 snapshot.run.cancelled_run_intent.is_some() || !snapshot.run.cancelled_node_intents.is_empty()
531}
532
533fn select_tick_actions(
534 actions: Vec<OrchestratorAction>,
535 def: &WorkflowDefinition,
536 max_concurrency: usize,
537) -> Vec<ScheduledAction> {
538 let limit = max_concurrency.max(1);
539 let mut selected = Vec::new();
540 let mut seen = std::collections::HashSet::new();
541 let mut dispatch_count: usize = 0;
542 for action in actions {
543 let serialization_key = action_serialization_key(def, &action);
544 if seen.insert(serialization_key.clone()) {
545 let is_dispatch = action.is_dispatch();
546 if !is_dispatch || dispatch_count < limit {
551 if is_dispatch {
552 dispatch_count += 1;
553 }
554 selected.push(ScheduledAction { action });
555 }
556 }
557 }
558 selected
559}
560
561fn action_serialization_key(_def: &WorkflowDefinition, action: &OrchestratorAction) -> String {
562 match action {
563 OrchestratorAction::DispatchWork { node_id, node, .. } => {
564 let bot_key = match node {
565 WorkflowNode::Subagent(node) => Some(format!("bot:{}", node.bot)),
566 WorkflowNode::HostExecutor(node) => Some(format!("executor:{}", node.executor)),
567 WorkflowNode::Loop(_) | WorkflowNode::Decision(_) => None,
568 };
569 bot_key.unwrap_or_else(|| format!("node:{node_id}"))
570 }
571 OrchestratorAction::DispatchGate { node_id, .. } => format!("gate:{node_id}"),
572 OrchestratorAction::CompleteNodeSucceeded { node_id, .. }
573 | OrchestratorAction::CompleteNodeFailed { node_id, .. } => {
574 format!("node:{node_id}")
575 }
576 OrchestratorAction::CompleteRunSucceeded { sink_node_id, .. } => {
577 format!("run:{sink_node_id}:succeeded")
578 }
579 OrchestratorAction::CompleteRunFailed { failed_node_id } => {
580 format!("run:{failed_node_id}:failed")
581 }
582 OrchestratorAction::StartLoop { node_id, .. } => {
583 format!("loop:start:{node_id}")
584 }
585 OrchestratorAction::StartLoopIteration { node_id, .. } => {
586 format!("loop:iter-start:{node_id}")
587 }
588 OrchestratorAction::FinishLoopIteration { node_id, .. } => {
589 format!("loop:iter-finish:{node_id}")
590 }
591 OrchestratorAction::FinishLoop { node_id, .. } => {
592 format!("loop:finish:{node_id}")
593 }
594 }
595}
596
597async fn apply_orchestrator_action<H: WorkflowExecutionHooks>(
598 rt: &mut WorkflowRuntimeContext,
599 hooks: &mut H,
600 action: OrchestratorAction,
601) -> Result<()> {
602 match action {
603 OrchestratorAction::DispatchGate { .. } => dispatch_gate(rt, &action).await?,
604 OrchestratorAction::DispatchWork { .. } => {
605 let _ = dispatch_work(rt, hooks, &action).await?;
606 }
607 OrchestratorAction::CompleteNodeSucceeded { .. } => {
608 complete_node_succeeded(&mut rt.log, &action).await?
609 }
610 OrchestratorAction::CompleteNodeFailed { .. } => {
611 complete_node_failed(&mut rt.log, &action).await?
612 }
613 OrchestratorAction::CompleteRunSucceeded { .. } => {
614 complete_run_succeeded(&mut rt.log, &action).await?
615 }
616 OrchestratorAction::CompleteRunFailed { .. } => {
617 complete_run_failed(&mut rt.log, &action).await?
618 }
619 OrchestratorAction::StartLoop { .. } => start_loop(&mut rt.log, &action).await?,
620 OrchestratorAction::StartLoopIteration { .. } => {
621 start_loop_iteration(&mut rt.log, &action).await?
622 }
623 OrchestratorAction::FinishLoopIteration { .. } => {
624 finish_loop_iteration(&mut rt.log, &action).await?
625 }
626 OrchestratorAction::FinishLoop { .. } => finish_loop(&mut rt.log, &action).await?,
627 }
628 Ok(())
629}
630
631async fn resolve_wait_terminals(
639 rt: &mut WorkflowRuntimeContext,
640 snapshot: &RunSnapshotDTO,
641) -> Result<bool> {
642 let mut had_progress = false;
643 for activity_id in &snapshot.dangling.wait_resolutions {
644 let Some(activity) = snapshot
645 .activities
646 .iter()
647 .find(|a| &a.activity_id == activity_id)
648 else {
649 continue;
650 };
651 let Some(latest) = activity.attempts.last() else {
652 continue;
653 };
654 let Some(wait) = latest.wait.as_ref() else {
655 continue;
656 };
657 let Some(resolution) = wait.resolution.as_ref() else {
658 continue;
659 };
660 let attempt_id = &latest.attempt_id;
661
662 match resolution.kind.as_str() {
663 "resolved" => {
664 if matches!(resolution.resolution.as_deref(), Some("rejected")) {
665 rt.log.append(EventDraft {
667 event_type: "activityFailed".to_string(),
668 actor: WorkflowActor::Scheduler,
669 payload: serde_json::json!({
670 "activityId": activity_id,
671 "attemptId": attempt_id,
672 "error": {
673 "errorCode": "InputValidationFailed",
674 "errorClass": "userFault",
675 "errorMessage": format!(
676 "Recovered wait terminal: rejected by {}{}",
677 resolution.by.clone().unwrap_or_default(),
678 resolution.comment.as_ref()
679 .map(|c| format!(": {}", c))
680 .unwrap_or_default()
681 ),
682 }
683 }),
684 timestamp: None,
685 payload_hash: None,
686 })?;
687 had_progress = true;
688 } else {
689 let external_refs = serde_json::json!({
691 "resolution": resolution.resolution,
692 "by": resolution.by,
693 "comment": resolution.comment,
694 });
695 let output_ref = write_json_blob(&mut rt.log, external_refs.clone())?;
696 rt.log.append(EventDraft {
697 event_type: "activitySucceeded".to_string(),
698 actor: WorkflowActor::Scheduler,
699 payload: serde_json::json!({
700 "activityId": activity_id,
701 "attemptId": attempt_id,
702 "outputRef": output_ref,
703 "externalRefs": external_refs,
704 }),
705 timestamp: None,
706 payload_hash: None,
707 })?;
708 had_progress = true;
709 }
710 }
711 "deadlineExceeded" => {
712 if matches!(wait.on_timeout.as_deref(), Some("success")) {
713 let external_refs = serde_json::json!({
714 "defaultedToTimeout": true,
715 "deadlineAt": resolution.deadline_at,
716 });
717 let output_ref = write_json_blob(&mut rt.log, external_refs.clone())?;
718 rt.log.append(EventDraft {
719 event_type: "activitySucceeded".to_string(),
720 actor: WorkflowActor::Scheduler,
721 payload: serde_json::json!({
722 "activityId": activity_id,
723 "attemptId": attempt_id,
724 "outputRef": output_ref,
725 "externalRefs": external_refs,
726 }),
727 timestamp: None,
728 payload_hash: None,
729 })?;
730 had_progress = true;
731 } else {
732 rt.log.append(EventDraft {
734 event_type: "activityFailed".to_string(),
735 actor: WorkflowActor::Scheduler,
736 payload: serde_json::json!({
737 "activityId": activity_id,
738 "attemptId": attempt_id,
739 "error": {
740 "errorCode": "WaitDeadlineExceeded",
741 "errorClass": "userFault",
742 "errorMessage": "Recovered wait terminal: deadline exceeded",
743 }
744 }),
745 timestamp: None,
746 payload_hash: None,
747 })?;
748 had_progress = true;
749 }
750 }
751 _ => {}
752 }
753 }
754 Ok(had_progress)
755}
756
757pub(crate) async fn read_snapshot(rt: &WorkflowRuntimeContext) -> Result<RunSnapshotDTO> {
758 read_run_snapshot(&rt.log.run_dir)
759 .await?
760 .context("workflow runtime requires an existing run snapshot")
761}