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;
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 metadata: BTreeMap<String, Value>,
26}
27
28impl std::fmt::Debug for ExecutionContext {
29    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
30        formatter
31            .debug_struct("ExecutionContext")
32            .field("has_cancellation_token", &self.cancellation_token.is_some())
33            .field("has_stream_callback", &self.stream_callback.is_some())
34            .field("has_state_store", &self.state_store.is_some())
35            .field("has_approval_provider", &self.approval_provider.is_some())
36            .field("has_approval_broker", &self.approval_broker.is_some())
37            .field("memory_provider_count", &self.memory_providers.len())
38            .field("metadata", &self.metadata)
39            .finish()
40    }
41}
42
43impl ExecutionContext {
44    pub fn with_cancellation_token(mut self, cancellation_token: CancellationToken) -> Self {
45        self.cancellation_token = Some(cancellation_token);
46        self
47    }
48
49    pub fn with_stream_callback(mut self, stream_callback: StreamCallback) -> Self {
50        self.stream_callback = Some(stream_callback);
51        self
52    }
53
54    pub fn with_state_store(mut self, state_store: Arc<dyn StateStore>) -> Self {
55        self.state_store = Some(state_store);
56        self
57    }
58
59    pub fn with_approval_provider(mut self, provider: Arc<dyn ApprovalProvider>) -> Self {
60        self.approval_provider = Some(provider);
61        self
62    }
63
64    pub fn with_approval_broker(mut self, broker: ApprovalBroker) -> Self {
65        self.approval_broker = Some(broker);
66        self
67    }
68
69    pub fn with_approval_timeout(mut self, timeout: Duration) -> Self {
70        self.approval_timeout = Some(timeout);
71        self
72    }
73
74    pub fn with_memory_provider(mut self, provider: Arc<dyn MemoryProvider>) -> Self {
75        self.memory_providers.push(provider);
76        self
77    }
78
79    pub fn with_metadata(mut self, metadata: BTreeMap<String, Value>) -> Self {
80        self.metadata = metadata;
81        self
82    }
83
84    pub fn check_cancelled(&self) -> Result<(), String> {
85        if let Some(token) = &self.cancellation_token {
86            token.check()
87        } else {
88            Ok(())
89        }
90    }
91}