Skip to main content

vv_agent/runtime/backends/distributed/
execution.rs

1use crate::budget::{BudgetUsageSnapshot, RunBudgetLimits};
2use crate::runtime::checkpoint_resume::CheckpointController;
3use crate::runtime::CancellationToken;
4use crate::types::{AgentResult, AgentTask, CycleRecord, Message, Metadata};
5
6use super::super::{execute_cycle_loop, execute_cycle_loop_with_state, failed_backend_result};
7use super::backend::DistributedBackend;
8
9impl DistributedBackend {
10    pub fn execute<F>(
11        &self,
12        _task: &AgentTask,
13        initial_messages: Vec<Message>,
14        shared_state: Metadata,
15        cycle_executor: F,
16        cancellation_token: Option<&CancellationToken>,
17        max_cycles: u32,
18    ) -> AgentResult
19    where
20        F: FnMut(
21            u32,
22            &mut Vec<Message>,
23            &mut Vec<CycleRecord>,
24            &mut Metadata,
25            Option<&CancellationToken>,
26        ) -> Option<AgentResult>,
27    {
28        if self.runtime_recipe.is_some() {
29            return failed_backend_result(
30                initial_messages,
31                Vec::new(),
32                shared_state,
33                "DistributedBackend requires RunConfig.checkpoint_config".to_string(),
34            );
35        }
36        execute_cycle_loop(
37            initial_messages,
38            shared_state,
39            cycle_executor,
40            cancellation_token,
41            max_cycles,
42        )
43    }
44
45    #[allow(clippy::too_many_arguments)]
46    pub(crate) fn execute_with_state<F>(
47        &self,
48        task: &AgentTask,
49        initial_messages: Vec<Message>,
50        initial_cycles: Vec<CycleRecord>,
51        shared_state: Metadata,
52        cycle_executor: F,
53        cancellation_token: Option<&CancellationToken>,
54        cycle_index_start: u32,
55        cycle_count: u32,
56        budget_limits: Option<RunBudgetLimits>,
57        _initial_budget_usage: Option<BudgetUsageSnapshot>,
58        checkpoint_controller: Option<CheckpointController>,
59    ) -> AgentResult
60    where
61        F: FnMut(
62            u32,
63            &mut Vec<Message>,
64            &mut Vec<CycleRecord>,
65            &mut Metadata,
66            Option<&CancellationToken>,
67        ) -> Option<AgentResult>,
68    {
69        if self.runtime_recipe.is_some()
70            && checkpoint_controller.is_none()
71            && (!initial_cycles.is_empty() || cycle_index_start != 1)
72        {
73            return failed_backend_result(
74                initial_messages,
75                initial_cycles,
76                shared_state,
77                "Checkpoint resume cannot recursively dispatch a distributed backend".to_string(),
78            );
79        }
80        if self.runtime_recipe.is_some() {
81            if let Some(checkpoint_controller) = checkpoint_controller {
82                return self.execute_distributed(
83                    task,
84                    cycle_index_start,
85                    cycle_count,
86                    budget_limits,
87                    cancellation_token,
88                    checkpoint_controller,
89                );
90            }
91            return failed_backend_result(
92                initial_messages,
93                initial_cycles,
94                shared_state,
95                "DistributedBackend requires RunConfig.checkpoint_config".to_string(),
96            );
97        }
98        execute_cycle_loop_with_state(
99            initial_messages,
100            initial_cycles,
101            shared_state,
102            cycle_executor,
103            cancellation_token,
104            cycle_index_start,
105            cycle_count,
106        )
107    }
108}