Skip to main content

vv_agent/runtime/backends/
base.rs

1use super::distributed::DistributedBackend;
2use super::inline::InlineBackend;
3use super::thread::ThreadBackend;
4use crate::budget::{BudgetUsageSnapshot, RunBudgetLimits};
5use crate::runtime::checkpoint_resume::CheckpointController;
6use crate::runtime::CancellationToken;
7use crate::types::{AgentResult, AgentTask, CycleRecord, Message, Metadata};
8
9#[derive(Debug, Clone)]
10#[allow(clippy::large_enum_variant)] // Preserve direct backend construction in the public API.
11pub enum RuntimeExecutionBackend {
12    Inline(InlineBackend),
13    Thread(ThreadBackend),
14    Distributed(DistributedBackend),
15}
16
17impl Default for RuntimeExecutionBackend {
18    fn default() -> Self {
19        Self::Inline(InlineBackend)
20    }
21}
22
23impl From<InlineBackend> for RuntimeExecutionBackend {
24    fn from(backend: InlineBackend) -> Self {
25        Self::Inline(backend)
26    }
27}
28
29impl From<ThreadBackend> for RuntimeExecutionBackend {
30    fn from(backend: ThreadBackend) -> Self {
31        Self::Thread(backend)
32    }
33}
34
35impl From<DistributedBackend> for RuntimeExecutionBackend {
36    fn from(backend: DistributedBackend) -> Self {
37        Self::Distributed(backend)
38    }
39}
40
41impl RuntimeExecutionBackend {
42    pub(crate) fn manages_run_budget(&self) -> bool {
43        matches!(self, Self::Distributed(backend) if backend.runtime_recipe().is_some())
44    }
45
46    pub(crate) fn manages_checkpoint_cycles(&self) -> bool {
47        matches!(self, Self::Distributed(backend) if backend.runtime_recipe().is_some())
48    }
49
50    pub fn execute<F>(
51        &self,
52        task: &AgentTask,
53        initial_messages: Vec<Message>,
54        shared_state: Metadata,
55        cycle_executor: F,
56        cancellation_token: Option<&CancellationToken>,
57        max_cycles: u32,
58    ) -> AgentResult
59    where
60        F: FnMut(
61            u32,
62            &mut Vec<Message>,
63            &mut Vec<CycleRecord>,
64            &mut Metadata,
65            Option<&CancellationToken>,
66        ) -> Option<AgentResult>,
67    {
68        match self {
69            Self::Inline(backend) => backend.execute(
70                task,
71                initial_messages,
72                shared_state,
73                cycle_executor,
74                cancellation_token,
75                max_cycles,
76            ),
77            Self::Thread(backend) => backend.execute(
78                task,
79                initial_messages,
80                shared_state,
81                cycle_executor,
82                cancellation_token,
83                max_cycles,
84            ),
85            Self::Distributed(backend) => backend.execute(
86                task,
87                initial_messages,
88                shared_state,
89                cycle_executor,
90                cancellation_token,
91                max_cycles,
92            ),
93        }
94    }
95
96    #[allow(clippy::too_many_arguments)]
97    pub(crate) fn execute_with_state<F>(
98        &self,
99        task: &AgentTask,
100        initial_messages: Vec<Message>,
101        initial_cycles: Vec<CycleRecord>,
102        shared_state: Metadata,
103        cycle_executor: F,
104        cancellation_token: Option<&CancellationToken>,
105        cycle_index_start: u32,
106        cycle_count: u32,
107        budget_limits: Option<RunBudgetLimits>,
108        initial_budget_usage: Option<BudgetUsageSnapshot>,
109        checkpoint_controller: Option<CheckpointController>,
110    ) -> AgentResult
111    where
112        F: FnMut(
113            u32,
114            &mut Vec<Message>,
115            &mut Vec<CycleRecord>,
116            &mut Metadata,
117            Option<&CancellationToken>,
118        ) -> Option<AgentResult>,
119    {
120        match self {
121            Self::Inline(backend) => backend.execute_with_state(
122                task,
123                initial_messages,
124                initial_cycles,
125                shared_state,
126                cycle_executor,
127                cancellation_token,
128                cycle_index_start,
129                cycle_count,
130            ),
131            Self::Thread(backend) => backend.execute_with_state(
132                task,
133                initial_messages,
134                initial_cycles,
135                shared_state,
136                cycle_executor,
137                cancellation_token,
138                cycle_index_start,
139                cycle_count,
140            ),
141            Self::Distributed(backend) => backend.execute_with_state(
142                task,
143                initial_messages,
144                initial_cycles,
145                shared_state,
146                cycle_executor,
147                cancellation_token,
148                cycle_index_start,
149                cycle_count,
150                budget_limits,
151                initial_budget_usage,
152                checkpoint_controller,
153            ),
154        }
155    }
156
157    pub fn parallel_map<T, R, F>(&self, function: F, items: Vec<T>) -> Vec<R>
158    where
159        T: Send + 'static,
160        R: Send + 'static,
161        F: Fn(T) -> R + Send + Sync + 'static,
162    {
163        match self {
164            Self::Inline(backend) => backend.parallel_map(function, items),
165            Self::Thread(backend) => backend.parallel_map(function, items),
166            Self::Distributed(backend) => backend.parallel_map(function, items),
167        }
168    }
169}