vv_agent/runtime/
context.rs1use 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}