Skip to main content

vv_agent/runtime/
context.rs

1use std::collections::BTreeMap;
2use std::sync::Arc;
3
4use serde_json::Value;
5
6use crate::llm::LlmStreamCallback;
7
8use super::state::StateStore;
9use super::CancellationToken;
10
11pub type StreamCallback = LlmStreamCallback;
12
13#[derive(Clone, Default)]
14pub struct ExecutionContext {
15    pub cancellation_token: Option<CancellationToken>,
16    pub stream_callback: Option<StreamCallback>,
17    pub state_store: Option<Arc<dyn StateStore>>,
18    pub metadata: BTreeMap<String, Value>,
19}
20
21impl std::fmt::Debug for ExecutionContext {
22    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        formatter
24            .debug_struct("ExecutionContext")
25            .field("has_cancellation_token", &self.cancellation_token.is_some())
26            .field("has_stream_callback", &self.stream_callback.is_some())
27            .field("has_state_store", &self.state_store.is_some())
28            .field("metadata", &self.metadata)
29            .finish()
30    }
31}
32
33impl ExecutionContext {
34    pub fn with_cancellation_token(mut self, cancellation_token: CancellationToken) -> Self {
35        self.cancellation_token = Some(cancellation_token);
36        self
37    }
38
39    pub fn with_stream_callback(mut self, stream_callback: StreamCallback) -> Self {
40        self.stream_callback = Some(stream_callback);
41        self
42    }
43
44    pub fn with_state_store(mut self, state_store: Arc<dyn StateStore>) -> Self {
45        self.state_store = Some(state_store);
46        self
47    }
48
49    pub fn with_metadata(mut self, metadata: BTreeMap<String, Value>) -> Self {
50        self.metadata = metadata;
51        self
52    }
53
54    pub fn check_cancelled(&self) -> Result<(), String> {
55        if let Some(token) = &self.cancellation_token {
56            token.check()
57        } else {
58            Ok(())
59        }
60    }
61}