Skip to main content

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::llm::LlmStreamCallback;
9use crate::memory::MemoryProvider;
10
11use super::state::StateStore;
12use super::{CancellationToken, CancelledError};
13
14pub type StreamCallback = LlmStreamCallback;
15
16#[derive(Clone, Default)]
17pub struct ExecutionContext {
18    pub cancellation_token: Option<CancellationToken>,
19    pub stream_callback: Option<StreamCallback>,
20    pub state_store: Option<Arc<dyn StateStore>>,
21    pub approval_provider: Option<Arc<dyn ApprovalProvider>>,
22    pub approval_broker: Option<ApprovalBroker>,
23    pub approval_timeout: Option<Duration>,
24    pub memory_providers: Vec<Arc<dyn MemoryProvider>>,
25    pub app_state: Option<Arc<dyn std::any::Any + Send + Sync>>,
26    pub metadata: BTreeMap<String, Value>,
27}
28
29impl std::fmt::Debug for ExecutionContext {
30    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31        formatter
32            .debug_struct("ExecutionContext")
33            .field("has_cancellation_token", &self.cancellation_token.is_some())
34            .field("has_stream_callback", &self.stream_callback.is_some())
35            .field("has_state_store", &self.state_store.is_some())
36            .field("has_approval_provider", &self.approval_provider.is_some())
37            .field("has_approval_broker", &self.approval_broker.is_some())
38            .field("memory_provider_count", &self.memory_providers.len())
39            .field("has_app_state", &self.app_state.is_some())
40            .field("metadata", &self.metadata)
41            .finish()
42    }
43}
44
45impl ExecutionContext {
46    pub fn with_cancellation_token(mut self, cancellation_token: CancellationToken) -> Self {
47        self.cancellation_token = Some(cancellation_token);
48        self
49    }
50
51    pub fn with_stream_callback(mut self, stream_callback: StreamCallback) -> Self {
52        self.stream_callback = Some(stream_callback);
53        self
54    }
55
56    pub fn with_state_store(mut self, state_store: Arc<dyn StateStore>) -> Self {
57        self.state_store = Some(state_store);
58        self
59    }
60
61    pub fn with_approval_provider(mut self, provider: Arc<dyn ApprovalProvider>) -> Self {
62        self.approval_provider = Some(provider);
63        self
64    }
65
66    pub fn with_approval_broker(mut self, broker: ApprovalBroker) -> Self {
67        self.approval_broker = Some(broker);
68        self
69    }
70
71    pub fn with_approval_timeout(mut self, timeout: Duration) -> Self {
72        self.approval_timeout = Some(timeout);
73        self
74    }
75
76    pub fn with_memory_provider(mut self, provider: Arc<dyn MemoryProvider>) -> Self {
77        self.memory_providers.push(provider);
78        self
79    }
80
81    pub fn with_app_state<T>(mut self, app_state: T) -> Self
82    where
83        T: Send + Sync + 'static,
84    {
85        self.app_state = Some(Arc::new(app_state));
86        self
87    }
88
89    pub fn with_app_state_arc(mut self, app_state: Arc<dyn std::any::Any + Send + Sync>) -> Self {
90        self.app_state = Some(app_state);
91        self
92    }
93
94    pub fn with_metadata(mut self, metadata: BTreeMap<String, Value>) -> Self {
95        self.metadata = metadata;
96        self
97    }
98
99    pub fn check_cancelled(&self) -> Result<(), CancelledError> {
100        if let Some(token) = &self.cancellation_token {
101            token.check()
102        } else {
103            Ok(())
104        }
105    }
106}