vv_agent/runtime/
context.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3use std::time::Duration;
4
5use serde_json::Value;
6
7use crate::approval::{ApprovalBroker, ApprovalProvider};
8use crate::events::RunEvent;
9use crate::memory::MemoryProvider;
10
11use super::model_calls::{ModelCallCoordinator, ModelCallLedger};
12use super::state::CheckpointStore;
13use super::{CancellationToken, CancelledError};
14
15pub type RunEventHandler = Arc<dyn Fn(&RunEvent) + Send + Sync + 'static>;
16
17#[doc(hidden)]
18#[derive(Clone, Default)]
19pub struct ExecutionRuntimeState {
20    pub(crate) model_call_ledger: ModelCallLedger,
21    pub(crate) model_call_coordinator: Option<ModelCallCoordinator>,
22}
23
24#[derive(Default)]
25pub struct ExecutionContext {
26    pub cancellation_token: Option<CancellationToken>,
27    pub event_handler: Option<RunEventHandler>,
28    pub checkpoint_store: Option<Arc<dyn CheckpointStore>>,
29    pub approval_provider: Option<Arc<dyn ApprovalProvider>>,
30    pub approval_broker: Option<ApprovalBroker>,
31    pub approval_timeout: Option<Duration>,
32    pub memory_providers: Vec<Arc<dyn MemoryProvider>>,
33    pub app_state: Option<Arc<dyn std::any::Any + Send + Sync>>,
34    pub metadata: BTreeMap<String, Value>,
35    #[doc(hidden)]
36    pub runtime_state: ExecutionRuntimeState,
37}
38
39impl Clone for ExecutionContext {
40    fn clone(&self) -> Self {
41        Self {
42            cancellation_token: self.cancellation_token.clone(),
43            event_handler: self.event_handler.clone(),
44            checkpoint_store: self.checkpoint_store.clone(),
45            approval_provider: self.approval_provider.clone(),
46            approval_broker: self.approval_broker.clone(),
47            approval_timeout: self.approval_timeout,
48            memory_providers: self.memory_providers.clone(),
49            app_state: self.app_state.clone(),
50            metadata: self.metadata.clone(),
51            runtime_state: ExecutionRuntimeState::default(),
52        }
53    }
54}
55
56impl std::fmt::Debug for ExecutionContext {
57    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
58        formatter
59            .debug_struct("ExecutionContext")
60            .field("has_cancellation_token", &self.cancellation_token.is_some())
61            .field("has_event_handler", &self.event_handler.is_some())
62            .field("has_checkpoint_store", &self.checkpoint_store.is_some())
63            .field("has_approval_provider", &self.approval_provider.is_some())
64            .field("has_approval_broker", &self.approval_broker.is_some())
65            .field("memory_provider_count", &self.memory_providers.len())
66            .field("has_app_state", &self.app_state.is_some())
67            .field(
68                "model_call_count",
69                &self.runtime_state.model_call_ledger.records().len(),
70            )
71            .field("metadata", &self.metadata)
72            .finish()
73    }
74}
75
76impl ExecutionContext {
77    pub fn with_cancellation_token(mut self, cancellation_token: CancellationToken) -> Self {
78        self.cancellation_token = Some(cancellation_token);
79        self
80    }
81
82    pub fn with_event_handler(mut self, event_handler: RunEventHandler) -> Self {
83        self.event_handler = Some(event_handler);
84        self
85    }
86
87    pub fn with_checkpoint_store(mut self, checkpoint_store: Arc<dyn CheckpointStore>) -> Self {
88        self.checkpoint_store = Some(checkpoint_store);
89        self
90    }
91
92    pub fn with_approval_provider(mut self, provider: Arc<dyn ApprovalProvider>) -> Self {
93        self.approval_provider = Some(provider);
94        self
95    }
96
97    pub fn with_approval_broker(mut self, broker: ApprovalBroker) -> Self {
98        self.approval_broker = Some(broker);
99        self
100    }
101
102    pub fn with_approval_timeout(mut self, timeout: Duration) -> Self {
103        self.approval_timeout = Some(timeout);
104        self
105    }
106
107    pub fn with_memory_provider(mut self, provider: Arc<dyn MemoryProvider>) -> Self {
108        self.memory_providers.push(provider);
109        self
110    }
111
112    pub fn with_app_state<T>(mut self, app_state: T) -> Self
113    where
114        T: Send + Sync + 'static,
115    {
116        self.app_state = Some(Arc::new(app_state));
117        self
118    }
119
120    pub fn with_app_state_arc(mut self, app_state: Arc<dyn std::any::Any + Send + Sync>) -> Self {
121        self.app_state = Some(app_state);
122        self
123    }
124
125    pub fn with_metadata(mut self, metadata: BTreeMap<String, Value>) -> Self {
126        self.metadata = metadata;
127        self
128    }
129
130    pub fn check_cancelled(&self) -> Result<(), CancelledError> {
131        if let Some(token) = &self.cancellation_token {
132            token.check()
133        } else {
134            Ok(())
135        }
136    }
137}
vv_agent/runtime/context.rs

vv_agent/runtime/
context.rs
