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 && Instant::now() >= deadline {
336 join_set.abort_all();
337 }
338 }
339 }
340 }
341
342 let snapshot = read_snapshot(rt).await?;
343 Ok(RunTickResult {
344 actions: applied,
345 snapshot,
346 })
347}
348
349pub async fn run_loop<H: WorkflowExecutionHooks + Clone + Send + 'static>(
350 rt: &mut WorkflowRuntimeContext,
351 hooks: &mut H,
352 max_ticks: usize,
353 max_concurrency: usize,
354) -> Result<RunLoopResult> {
355 let mut ticks = 0usize;
356 loop {
357 if ticks >= max_ticks {
358 let snapshot = read_snapshot(rt).await?;
359 return Ok(RunLoopResult {
360 reason: RunLoopStopReason::MaxTicks,
361 ticks,
362 last_snapshot: snapshot,
363 });
364 }
365
366 check_pending_cancels(rt, hooks).await?;
367
368 let pre_recovery_snapshot = read_snapshot(rt).await?;
375 if matches!(
376 pre_recovery_snapshot.run.status,
377 RunStatus::Succeeded | RunStatus::Failed | RunStatus::Cancelled
378 ) {
379 return Ok(RunLoopResult {
380 reason: RunLoopStopReason::Terminal,
381 ticks,
382 last_snapshot: pre_recovery_snapshot,
383 });
384 }
385
386 if !pre_recovery_snapshot.dangling.effect_attempted.is_empty() {
387 let recovery = hooks
388 .recover_dangling_effects(&mut rt.log, &pre_recovery_snapshot)
389 .await?;
390 if recovery.had_progress {
391 continue;
394 }
395 }
400
401 if !pre_recovery_snapshot.dangling.wait_resolutions.is_empty() {
407 let had_progress = resolve_wait_terminals(rt, &pre_recovery_snapshot).await?;
408 if had_progress {
409 continue;
410 }
411 }
412
413 let tick = run_tick(rt, hooks, max_concurrency).await?;
414 ticks += 1;
415 if tick.snapshot.run.status == RunStatus::Succeeded
416 || tick.snapshot.run.status == RunStatus::Failed
417 || tick.snapshot.run.status == RunStatus::Cancelled
418 {
419 return Ok(RunLoopResult {
420 reason: RunLoopStopReason::Terminal,
421 ticks,
422 last_snapshot: tick.snapshot,
423 });
424 }
425 if tick.actions == 0 {
426 let has_waits = !tick.snapshot.dangling.waits.is_empty()
427 && tick
428 .snapshot
429 .dangling
430 .waits
431 .iter()
432 .any(|w| !tick.snapshot.dangling.cancels.contains(w));
433 let reason = if has_waits {
434 RunLoopStopReason::AwaitingWait
435 } else {
436 RunLoopStopReason::NoProgress
437 };
438 return Ok(RunLoopResult {
439 reason,
440 ticks,
441 last_snapshot: tick.snapshot,
442 });
443 }
444 }
445}
446
447async fn check_pending_cancels<H: WorkflowExecutionHooks + Send>(
448 rt: &mut WorkflowRuntimeContext,
449 hooks: &mut H,
450) -> Result<()> {
451 let _events = rt.log.read_all()?;
452 let snapshot = read_snapshot(rt).await?;
453 let mut cancelled_activities: Vec<String> = Vec::new();
454 let mut cancelled_nodes: Vec<String> = Vec::new();
455
456 for activity_id in &snapshot.dangling.cancels {
457 cancelled_activities.push(activity_id.clone());
458 let attempt_id = snapshot
459 .activities
460 .iter()
461 .find(|a| &a.activity_id == activity_id)
462 .and_then(|a| a.current_attempt_id.clone())
463 .unwrap_or_else(|| format!("{}-attempt-1", activity_id));
464 let origin = snapshot
465 .run
466 .cancelled_run_intent
467 .as_ref()
468 .map(|i| i.cancel_origin_event_id.clone())
469 .or_else(|| {
470 snapshot
471 .run
472 .cancelled_node_intents
473 .values()
474 .next()
475 .map(|i| i.cancel_origin_event_id.clone())
476 })
477 .unwrap_or_default();
478 let _ = crate::complete_activity_cancel(
479 &mut rt.log,
480 crate::CompleteActivityCancelInput {
481 activity_id: activity_id.clone(),
482 attempt_id,
483 cancel_origin_event_id: origin,
484 },
485 WorkflowActor::Scheduler,
486 )
487 .await;
488 }
489
490 if let Some(ref intent) = snapshot.run.cancelled_run_intent
491 && snapshot.run.status != RunStatus::Cancelled
492 {
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 if !snapshot.run.cancelled_node_intents.is_empty() {
504 for (node_id, intent) in &snapshot.run.cancelled_node_intents {
505 cancelled_nodes.push(node_id.clone());
506 let _ = crate::complete_node_cancel(
507 &mut rt.log,
508 crate::CompleteNodeCancelInput {
509 node_id: node_id.clone(),
510 cancel_origin_event_id: intent.cancel_origin_event_id.clone(),
511 },
512 WorkflowActor::Scheduler,
513 )
514 .await;
515 }
516 }
517
518 if !cancelled_activities.is_empty() || !cancelled_nodes.is_empty() {
520 hooks
521 .on_activities_cancelled(&cancelled_activities, &cancelled_nodes, &rt.log.run_id)
522 .await;
523 }
524
525 Ok(())
526}
527
528fn snapshot_has_pending_cancel(snapshot: &RunSnapshotDTO) -> bool {
529 snapshot.run.cancelled_run_intent.is_some() || !snapshot.run.cancelled_node_intents.is_empty()
530}
531
532fn select_tick_actions(
533 actions: Vec<OrchestratorAction>,
534 def: &WorkflowDefinition,
535 max_concurrency: usize,
536) -> Vec<ScheduledAction> {
537 let limit = max_concurrency.max(1);
538 let mut selected = Vec::new();
539 let mut seen = std::collections::HashSet::new();
540 let mut dispatch_count: usize = 0;
541 for action in actions {
542 let serialization_key = action_serialization_key(def, &action);
543 if seen.insert(serialization_key.clone()) {
544 let is_dispatch = action.is_dispatch();
545 if !is_dispatch || dispatch_count < limit {
550 if is_dispatch {
551 dispatch_count += 1;
552 }
553 selected.push(ScheduledAction { action });
554 }
555 }
556 }
557 selected
558}
559
560fn action_serialization_key(_def: &WorkflowDefinition, action: &OrchestratorAction) -> String {
561 match action {
562 OrchestratorAction::DispatchWork { node_id, node, .. } => {
563 let bot_key = match node.as_ref() {
564 WorkflowNode::Subagent(node) => Some(format!("bot:{}", node.bot)),
565 WorkflowNode::HostExecutor(node) => Some(format!("executor:{}", node.executor)),
566 WorkflowNode::Loop(_) | WorkflowNode::Decision(_) => None,
567 };
568 bot_key.unwrap_or_else(|| format!("node:{node_id}"))
569 }
570 OrchestratorAction::DispatchGate { node_id, .. } => format!("gate:{node_id}"),
571 OrchestratorAction::CompleteNodeSucceeded { node_id, .. }
572 | OrchestratorAction::CompleteNodeFailed { node_id, .. } => {
573 format!("node:{node_id}")
574 }
575 OrchestratorAction::CompleteRunSucceeded { sink_node_id, .. } => {
576 format!("run:{sink_node_id}:succeeded")
577 }
578 OrchestratorAction::CompleteRunFailed { failed_node_id } => {
579 format!("run:{failed_node_id}:failed")
580 }
581 OrchestratorAction::StartLoop { node_id, .. } => {
582 format!("loop:start:{node_id}")
583 }
584 OrchestratorAction::StartLoopIteration { node_id, .. } => {
585 format!("loop:iter-start:{node_id}")
586 }
587 OrchestratorAction::FinishLoopIteration { node_id, .. } => {
588 format!("loop:iter-finish:{node_id}")
589 }
590 OrchestratorAction::FinishLoop { node_id, .. } => {
591 format!("loop:finish:{node_id}")
592 }
593 }
594}
595
596async fn apply_orchestrator_action<H: WorkflowExecutionHooks>(
597 rt: &mut WorkflowRuntimeContext,
598 hooks: &mut H,
599 action: OrchestratorAction,
600) -> Result<()> {
601 match action {
602 OrchestratorAction::DispatchGate { .. } => dispatch_gate(rt, &action).await?,
603 OrchestratorAction::DispatchWork { .. } => {
604 let _ = dispatch_work(rt, hooks, &action).await?;
605 }
606 OrchestratorAction::CompleteNodeSucceeded { .. } => {
607 complete_node_succeeded(&mut rt.log, &action).await?
608 }
609 OrchestratorAction::CompleteNodeFailed { .. } => {
610 complete_node_failed(&mut rt.log, &action).await?
611 }
612 OrchestratorAction::CompleteRunSucceeded { .. } => {
613 complete_run_succeeded(&mut rt.log, &action).await?
614 }
615 OrchestratorAction::CompleteRunFailed { .. } => {
616 complete_run_failed(&mut rt.log, &action).await?
617 }
618 OrchestratorAction::StartLoop { .. } => start_loop(&mut rt.log, &action).await?,
619 OrchestratorAction::StartLoopIteration { .. } => {
620 start_loop_iteration(&mut rt.log, &action).await?
621 }
622 OrchestratorAction::FinishLoopIteration { .. } => {
623 finish_loop_iteration(&mut rt.log, &action).await?
624 }
625 OrchestratorAction::FinishLoop { .. } => finish_loop(&mut rt.log, &action).await?,
626 }
627 Ok(())
628}
629
630async fn resolve_wait_terminals(
638 rt: &mut WorkflowRuntimeContext,
639 snapshot: &RunSnapshotDTO,
640) -> Result<bool> {
641 let mut had_progress = false;
642 for activity_id in &snapshot.dangling.wait_resolutions {
643 let Some(activity) = snapshot
644 .activities
645 .iter()
646 .find(|a| &a.activity_id == activity_id)
647 else {
648 continue;
649 };
650 let Some(latest) = activity.attempts.last() else {
651 continue;
652 };
653 let Some(wait) = latest.wait.as_ref() else {
654 continue;
655 };
656 let Some(resolution) = wait.resolution.as_ref() else {
657 continue;
658 };
659 let attempt_id = &latest.attempt_id;
660
661 match resolution.kind.as_str() {
662 "resolved" => {
663 if matches!(resolution.resolution.as_deref(), Some("rejected")) {
664 rt.log.append(EventDraft {
666 event_type: "activityFailed".to_string(),
667 actor: WorkflowActor::Scheduler,
668 payload: serde_json::json!({
669 "activityId": activity_id,
670 "attemptId": attempt_id,
671 "error": {
672 "errorCode": "InputValidationFailed",
673 "errorClass": "userFault",
674 "errorMessage": format!(
675 "Recovered wait terminal: rejected by {}{}",
676 resolution.by.clone().unwrap_or_default(),
677 resolution.comment.as_ref()
678 .map(|c| format!(": {}", c))
679 .unwrap_or_default()
680 ),
681 }
682 }),
683 timestamp: None,
684 payload_hash: None,
685 })?;
686 had_progress = true;
687 } else {
688 let external_refs = serde_json::json!({
690 "resolution": resolution.resolution,
691 "by": resolution.by,
692 "comment": resolution.comment,
693 });
694 let output_ref = write_json_blob(&mut rt.log, external_refs.clone())?;
695 rt.log.append(EventDraft {
696 event_type: "activitySucceeded".to_string(),
697 actor: WorkflowActor::Scheduler,
698 payload: serde_json::json!({
699 "activityId": activity_id,
700 "attemptId": attempt_id,
701 "outputRef": output_ref,
702 "externalRefs": external_refs,
703 }),
704 timestamp: None,
705 payload_hash: None,
706 })?;
707 had_progress = true;
708 }
709 }
710 "deadlineExceeded" => {
711 if matches!(wait.on_timeout.as_deref(), Some("success")) {
712 let external_refs = serde_json::json!({
713 "defaultedToTimeout": true,
714 "deadlineAt": resolution.deadline_at,
715 });
716 let output_ref = write_json_blob(&mut rt.log, external_refs.clone())?;
717 rt.log.append(EventDraft {
718 event_type: "activitySucceeded".to_string(),
719 actor: WorkflowActor::Scheduler,
720 payload: serde_json::json!({
721 "activityId": activity_id,
722 "attemptId": attempt_id,
723 "outputRef": output_ref,
724 "externalRefs": external_refs,
725 }),
726 timestamp: None,
727 payload_hash: None,
728 })?;
729 had_progress = true;
730 } else {
731 rt.log.append(EventDraft {
733 event_type: "activityFailed".to_string(),
734 actor: WorkflowActor::Scheduler,
735 payload: serde_json::json!({
736 "activityId": activity_id,
737 "attemptId": attempt_id,
738 "error": {
739 "errorCode": "WaitDeadlineExceeded",
740 "errorClass": "userFault",
741 "errorMessage": "Recovered wait terminal: deadline exceeded",
742 }
743 }),
744 timestamp: None,
745 payload_hash: None,
746 })?;
747 had_progress = true;
748 }
749 }
750 _ => {}
751 }
752 }
753 Ok(had_progress)
754}
755
756pub(crate) async fn read_snapshot(rt: &WorkflowRuntimeContext) -> Result<RunSnapshotDTO> {
757 read_run_snapshot(&rt.log.run_dir)
758 .await?
759 .context("workflow runtime requires an existing run snapshot")
760}