Skip to main content

vv_agent/
run_config.rs

1use std::path::PathBuf;
2use std::sync::Arc;
3use std::time::Duration;
4
5use serde_json::Value;
6
7use crate::approval::{ApprovalBroker, ApprovalProvider};
8use crate::budget::{HostCostMeter, RunBudgetLimits};
9use crate::checkpoint::{CheckpointConfig, CheckpointExtension, ReconciliationProvider};
10use crate::context_providers::ContextProvider;
11use crate::event_store::RunEventStore;
12use crate::execution_mode::ExecutionMode;
13use crate::memory::MemoryProvider;
14use crate::model::{ModelProvider, ModelRef};
15use crate::model_settings::ModelSettings;
16use crate::runtime::backends::RuntimeExecutionBackend;
17use crate::runtime::{
18    AfterCycleHook, BeforeCycleMessageProvider, CancellationToken, InterruptionMessageProvider,
19    RunEventHandler, RuntimeHook, SubTaskManager,
20};
21use crate::sessions::Session;
22use crate::tools::{ToolPolicy, ToolRegistry};
23use crate::tracing::TraceSink;
24use crate::types::{Message, Metadata, NoToolPolicy};
25use crate::workspace::WorkspaceBackend;
26
27pub type ToolRegistryFactory = Arc<dyn Fn() -> ToolRegistry + Send + Sync + 'static>;
28
29pub(crate) const INITIAL_BUDGET_USAGE_METADATA_KEY: &str = "_vv_agent_initial_budget_usage";
30
31pub(crate) const MAX_CYCLES_RANGE_ERROR: &str = "max_cycles must be between 1 and 4294967295";
32
33pub(crate) fn validate_max_cycles(max_cycles: u32) -> Result<u32, String> {
34    if max_cycles == 0 {
35        return Err(MAX_CYCLES_RANGE_ERROR.to_string());
36    }
37    Ok(max_cycles)
38}
39
40#[derive(Clone, Default)]
41pub struct RunConfig {
42    pub model: Option<ModelRef>,
43    pub model_provider: Option<Arc<dyn ModelProvider>>,
44    pub model_settings: Option<ModelSettings>,
45    pub workspace: Option<PathBuf>,
46    pub workspace_backend: Option<Arc<dyn WorkspaceBackend>>,
47    pub session: Option<Arc<dyn Session>>,
48    pub session_memory_enabled: bool,
49    pub initial_messages: Option<Vec<Message>>,
50    pub max_cycles: Option<u32>,
51    pub max_handoffs: Option<u32>,
52    pub no_tool_policy: Option<NoToolPolicy>,
53    pub tool_policy: ToolPolicy,
54    pub execution_backend: Option<RuntimeExecutionBackend>,
55    pub cancellation_token: Option<CancellationToken>,
56    pub hooks: Vec<Arc<dyn RuntimeHook>>,
57    pub after_cycle_hooks: Vec<Arc<dyn AfterCycleHook>>,
58    pub trace_sink: Option<Arc<dyn TraceSink>>,
59    pub trace_id: Option<String>,
60    pub workflow_name: Option<String>,
61    pub event_store: Option<Arc<dyn RunEventStore>>,
62    pub event_store_fail_closed: bool,
63    pub approval_provider: Option<Arc<dyn ApprovalProvider>>,
64    pub approval_timeout: Option<Duration>,
65    pub approval_broker: Option<ApprovalBroker>,
66    pub context_providers: Vec<Arc<dyn ContextProvider>>,
67    pub max_context_chars: Option<usize>,
68    pub memory_providers: Vec<Arc<dyn MemoryProvider>>,
69    pub app_state: Option<Arc<dyn std::any::Any + Send + Sync>>,
70    pub initial_shared_state: Metadata,
71    pub tool_registry_factory: Option<ToolRegistryFactory>,
72    pub log_preview_chars: Option<usize>,
73    pub debug_dump_dir: Option<PathBuf>,
74    pub before_cycle_messages: Option<BeforeCycleMessageProvider>,
75    pub interruption_messages: Option<InterruptionMessageProvider>,
76    pub sub_task_manager: Option<SubTaskManager>,
77    pub stream: Option<RunEventHandler>,
78    pub budget_limits: Option<RunBudgetLimits>,
79    pub host_cost_meter: Option<Arc<dyn HostCostMeter>>,
80    pub checkpoint_config: Option<CheckpointConfig>,
81    pub checkpoint_extensions: Vec<Arc<dyn CheckpointExtension>>,
82    pub reconciliation_provider: Option<Arc<dyn ReconciliationProvider>>,
83    pub metadata: Metadata,
84}
85
86impl RunConfig {
87    pub fn builder() -> RunConfigBuilder {
88        RunConfigBuilder::default()
89    }
90
91    pub(crate) fn for_background_child(&self, shared_state: Metadata) -> Self {
92        let mut child = self.clone();
93        child.model = None;
94        child.model_settings = None;
95        child.session = None;
96        child.initial_messages = None;
97        child.cancellation_token = self
98            .cancellation_token
99            .as_ref()
100            .map(CancellationToken::child);
101        child.initial_shared_state = shared_state;
102        child.before_cycle_messages = None;
103        child.interruption_messages = None;
104        child.sub_task_manager = None;
105        child.stream = None;
106        child.host_cost_meter = None;
107        child.checkpoint_config = None;
108        child.checkpoint_extensions.clear();
109        child.reconciliation_provider = None;
110        child.metadata.remove(INITIAL_BUDGET_USAGE_METADATA_KEY);
111        child
112    }
113}
114
115#[derive(Default)]
116pub struct RunConfigBuilder {
117    config: RunConfig,
118}
119
120impl RunConfigBuilder {
121    pub fn model(mut self, model: ModelRef) -> Self {
122        self.config.model = Some(model);
123        self
124    }
125
126    pub fn model_provider(mut self, provider: impl ModelProvider + 'static) -> Self {
127        self.config.model_provider = Some(Arc::new(provider));
128        self
129    }
130
131    pub fn model_provider_arc(mut self, provider: Arc<dyn ModelProvider>) -> Self {
132        self.config.model_provider = Some(provider);
133        self
134    }
135
136    pub fn model_settings(mut self, settings: ModelSettings) -> Self {
137        self.config.model_settings = Some(settings);
138        self
139    }
140
141    pub fn workspace(mut self, workspace: impl Into<PathBuf>) -> Self {
142        self.config.workspace = Some(workspace.into());
143        self
144    }
145
146    pub fn workspace_backend(mut self, backend: Arc<dyn WorkspaceBackend>) -> Self {
147        self.config.workspace_backend = Some(backend);
148        self
149    }
150
151    pub fn session(mut self, session: impl Session + 'static) -> Self {
152        self.config.session = Some(Arc::new(session));
153        self
154    }
155
156    pub fn session_arc(mut self, session: Arc<dyn Session>) -> Self {
157        self.config.session = Some(session);
158        self
159    }
160
161    pub fn session_memory_enabled(mut self, enabled: bool) -> Self {
162        self.config.session_memory_enabled = enabled;
163        self
164    }
165
166    pub fn initial_messages(mut self, messages: Vec<Message>) -> Self {
167        self.config.initial_messages = Some(messages);
168        self
169    }
170
171    pub fn max_cycles(mut self, max_cycles: u32) -> Self {
172        self.config.max_cycles = Some(max_cycles);
173        self
174    }
175
176    pub fn max_handoffs(mut self, max_handoffs: u32) -> Self {
177        self.config.max_handoffs = Some(max_handoffs);
178        self
179    }
180
181    pub fn no_tool_policy(mut self, policy: NoToolPolicy) -> Self {
182        self.config.no_tool_policy = Some(policy);
183        self
184    }
185
186    pub fn tool_policy(mut self, tool_policy: ToolPolicy) -> Self {
187        self.config.tool_policy = tool_policy;
188        self
189    }
190
191    pub fn execution_backend(mut self, execution_backend: RuntimeExecutionBackend) -> Self {
192        self.config.execution_backend = Some(execution_backend);
193        self
194    }
195
196    pub fn execution_mode(mut self, execution_mode: ExecutionMode) -> Self {
197        self.config.execution_backend = Some(execution_mode.into());
198        self
199    }
200
201    pub fn cancellation_token(mut self, cancellation_token: CancellationToken) -> Self {
202        self.config.cancellation_token = Some(cancellation_token);
203        self
204    }
205
206    pub fn hook(mut self, hook: Arc<dyn RuntimeHook>) -> Self {
207        self.config.hooks.push(hook);
208        self
209    }
210
211    pub fn after_cycle_hook(mut self, hook: impl AfterCycleHook + 'static) -> Self {
212        self.config.after_cycle_hooks.push(Arc::new(hook));
213        self
214    }
215
216    pub fn after_cycle_hook_arc(mut self, hook: Arc<dyn AfterCycleHook>) -> Self {
217        self.config.after_cycle_hooks.push(hook);
218        self
219    }
220
221    pub fn trace_sink(mut self, sink: Arc<dyn TraceSink>) -> Self {
222        self.config.trace_sink = Some(sink);
223        self
224    }
225
226    pub fn trace_id(mut self, trace_id: impl Into<String>) -> Self {
227        self.config.trace_id = Some(trace_id.into());
228        self
229    }
230
231    pub fn workflow_name(mut self, workflow_name: impl Into<String>) -> Self {
232        self.config.workflow_name = Some(workflow_name.into());
233        self
234    }
235
236    pub fn event_store(mut self, store: Arc<dyn RunEventStore>) -> Self {
237        self.config.event_store = Some(store);
238        self
239    }
240
241    pub fn event_store_fail_closed(mut self, fail_closed: bool) -> Self {
242        self.config.event_store_fail_closed = fail_closed;
243        self
244    }
245
246    pub fn approval_provider(mut self, provider: Arc<dyn ApprovalProvider>) -> Self {
247        self.config.approval_provider = Some(provider);
248        self
249    }
250
251    pub fn approval_timeout(mut self, timeout: Duration) -> Self {
252        self.config.approval_timeout = Some(timeout);
253        self
254    }
255
256    pub fn approval_broker(mut self, broker: ApprovalBroker) -> Self {
257        self.config.approval_broker = Some(broker);
258        self
259    }
260
261    pub fn context_provider(mut self, provider: Arc<dyn ContextProvider>) -> Self {
262        self.config.context_providers.push(provider);
263        self
264    }
265
266    pub fn max_context_chars(mut self, max_chars: usize) -> Self {
267        self.config.max_context_chars = Some(max_chars);
268        self
269    }
270
271    pub fn memory_provider(mut self, provider: Arc<dyn MemoryProvider>) -> Self {
272        self.config.memory_providers.push(provider);
273        self
274    }
275
276    pub fn app_state<T>(mut self, app_state: T) -> Self
277    where
278        T: Send + Sync + 'static,
279    {
280        self.config.app_state = Some(Arc::new(app_state));
281        self
282    }
283
284    pub fn app_state_arc(mut self, app_state: Arc<dyn std::any::Any + Send + Sync>) -> Self {
285        self.config.app_state = Some(app_state);
286        self
287    }
288
289    pub fn initial_shared_state(mut self, shared_state: Metadata) -> Self {
290        self.config.initial_shared_state = shared_state;
291        self
292    }
293
294    pub fn shared_state(self, shared_state: Metadata) -> Self {
295        self.initial_shared_state(shared_state)
296    }
297
298    pub fn tool_registry_factory(
299        mut self,
300        factory: impl Fn() -> ToolRegistry + Send + Sync + 'static,
301    ) -> Self {
302        self.config.tool_registry_factory = Some(Arc::new(factory));
303        self
304    }
305
306    pub fn tool_registry_factory_arc(mut self, factory: ToolRegistryFactory) -> Self {
307        self.config.tool_registry_factory = Some(factory);
308        self
309    }
310
311    pub fn log_preview_chars(mut self, max_chars: usize) -> Self {
312        self.config.log_preview_chars = Some(max_chars);
313        self
314    }
315
316    pub fn debug_dump_dir(mut self, path: impl Into<PathBuf>) -> Self {
317        self.config.debug_dump_dir = Some(path.into());
318        self
319    }
320
321    pub fn before_cycle_messages(
322        mut self,
323        provider: impl Fn(u32, &[Message], &Metadata) -> Vec<Message> + Send + Sync + 'static,
324    ) -> Self {
325        self.config.before_cycle_messages = Some(Arc::new(provider));
326        self
327    }
328
329    pub fn before_cycle_messages_arc(mut self, provider: BeforeCycleMessageProvider) -> Self {
330        self.config.before_cycle_messages = Some(provider);
331        self
332    }
333
334    pub fn interruption_messages(
335        mut self,
336        provider: impl Fn() -> Vec<Message> + Send + Sync + 'static,
337    ) -> Self {
338        self.config.interruption_messages = Some(Arc::new(provider));
339        self
340    }
341
342    pub fn interruption_messages_arc(mut self, provider: InterruptionMessageProvider) -> Self {
343        self.config.interruption_messages = Some(provider);
344        self
345    }
346
347    pub fn sub_task_manager(mut self, manager: SubTaskManager) -> Self {
348        self.config.sub_task_manager = Some(manager);
349        self
350    }
351
352    pub fn stream(
353        mut self,
354        observer: impl Fn(&crate::events::RunEvent) + Send + Sync + 'static,
355    ) -> Self {
356        self.config.stream = Some(Arc::new(observer));
357        self
358    }
359
360    pub fn stream_arc(mut self, observer: RunEventHandler) -> Self {
361        self.config.stream = Some(observer);
362        self
363    }
364
365    pub fn budget_limits(mut self, limits: RunBudgetLimits) -> Self {
366        self.config.budget_limits = Some(limits);
367        self
368    }
369
370    pub fn host_cost_meter(mut self, meter: impl HostCostMeter + 'static) -> Self {
371        self.config.host_cost_meter = Some(Arc::new(meter));
372        self
373    }
374
375    pub fn host_cost_meter_arc(mut self, meter: Arc<dyn HostCostMeter>) -> Self {
376        self.config.host_cost_meter = Some(meter);
377        self
378    }
379
380    pub fn checkpoint_config(mut self, checkpoint_config: CheckpointConfig) -> Self {
381        self.config.checkpoint_config = Some(checkpoint_config);
382        self
383    }
384
385    pub fn checkpoint_extension(mut self, extension: impl CheckpointExtension + 'static) -> Self {
386        self.config.checkpoint_extensions.push(Arc::new(extension));
387        self
388    }
389
390    pub fn checkpoint_extension_arc(mut self, extension: Arc<dyn CheckpointExtension>) -> Self {
391        self.config.checkpoint_extensions.push(extension);
392        self
393    }
394
395    pub fn reconciliation_provider(
396        mut self,
397        provider: impl ReconciliationProvider + 'static,
398    ) -> Self {
399        self.config.reconciliation_provider = Some(Arc::new(provider));
400        self
401    }
402
403    pub fn reconciliation_provider_arc(
404        mut self,
405        provider: Arc<dyn ReconciliationProvider>,
406    ) -> Self {
407        self.config.reconciliation_provider = Some(provider);
408        self
409    }
410
411    pub fn metadata(mut self, key: impl Into<String>, value: Value) -> Self {
412        self.config.metadata.insert(key.into(), value);
413        self
414    }
415
416    pub fn build(self) -> RunConfig {
417        self.config
418    }
419}