Skip to main content

starweaver_model/adapter/
context.rs

1use serde::{Deserialize, Serialize};
2use serde_json::{Map, Value};
3use starweaver_core::{CancellationToken, ConversationId, RunId, TraceContext};
4
5/// Per-request context attached by the runtime.
6#[derive(Clone, Deserialize, Serialize)]
7pub struct ModelRequestContext {
8    /// Run identifier.
9    pub run_id: RunId,
10    /// Conversation identifier.
11    pub conversation_id: ConversationId,
12    /// Trace correlation context propagated from the runtime context.
13    #[serde(default, skip_serializing_if = "TraceContext::is_empty")]
14    pub trace_context: TraceContext,
15    /// Debug metadata for raw provider request/response/event-stream evidence.
16    #[serde(default, skip_serializing_if = "Map::is_empty")]
17    pub llm_trace_metadata: Map<String, Value>,
18    /// Cancellation token for streaming and long-running provider requests.
19    #[serde(skip, default)]
20    pub cancellation_token: CancellationToken,
21}
22
23impl std::fmt::Debug for ModelRequestContext {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        f.debug_struct("ModelRequestContext")
26            .field("run_id", &self.run_id)
27            .field("conversation_id", &self.conversation_id)
28            .field("trace_context", &self.trace_context)
29            .field("cancelled", &self.cancellation_token.is_cancelled())
30            .finish_non_exhaustive()
31    }
32}
33
34impl PartialEq for ModelRequestContext {
35    fn eq(&self, other: &Self) -> bool {
36        self.run_id == other.run_id
37            && self.conversation_id == other.conversation_id
38            && self.trace_context == other.trace_context
39            && self.llm_trace_metadata == other.llm_trace_metadata
40            && self.cancellation_token == other.cancellation_token
41    }
42}
43
44impl Eq for ModelRequestContext {}
45
46impl ModelRequestContext {
47    /// Build request context for one model call.
48    #[must_use]
49    pub fn new(run_id: RunId, conversation_id: ConversationId) -> Self {
50        Self {
51            run_id,
52            conversation_id,
53            trace_context: TraceContext::default(),
54            llm_trace_metadata: Map::new(),
55            cancellation_token: CancellationToken::default(),
56        }
57    }
58
59    /// Attach trace correlation context.
60    #[must_use]
61    pub fn with_trace_context(mut self, trace_context: TraceContext) -> Self {
62        self.trace_context = trace_context;
63        self
64    }
65
66    /// Attach LLM request debug metadata.
67    #[must_use]
68    pub fn with_llm_trace_metadata(mut self, metadata: Map<String, Value>) -> Self {
69        self.llm_trace_metadata = metadata;
70        self
71    }
72
73    /// Attach a cancellation token.
74    #[must_use]
75    pub fn with_cancellation_token(mut self, token: CancellationToken) -> Self {
76        self.cancellation_token = token;
77        self
78    }
79
80    /// Return the shared cancellation token.
81    #[must_use]
82    pub fn cancellation_token(&self) -> CancellationToken {
83        self.cancellation_token.clone()
84    }
85}