Skip to main content

systemprompt_models/execution/context/
mod.rs

1//! Request context for execution tracking.
2
3mod call_source;
4mod context_error;
5mod context_types;
6mod propagation;
7
8pub use call_source::CallSource;
9pub use context_error::{ContextExtractionError, ContextIdSource, TASK_BASED_CONTEXT_MARKER};
10pub use context_types::{
11    AuthContext, ExecutionContext, ExecutionSettings, RequestMetadata, UserInteractionMode,
12};
13
14use crate::ai::ToolModelConfig;
15use crate::auth::{AuthenticatedUser, RateLimitTier, UserType};
16use serde::{Deserialize, Serialize};
17use std::time::{Duration, Instant};
18use systemprompt_identifiers::{
19    AgentName, AiToolCallId, ClientId, ContextId, JwtToken, McpExecutionId, SessionId, TaskId,
20    TraceId, UserId,
21};
22
23#[derive(Debug, Clone, Serialize, Deserialize)]
24pub struct RequestContext {
25    pub auth: AuthContext,
26    pub request: RequestMetadata,
27    pub execution: ExecutionContext,
28    pub settings: ExecutionSettings,
29
30    #[serde(skip)]
31    pub user: Option<AuthenticatedUser>,
32
33    #[serde(skip, default = "Instant::now")]
34    pub start_time: Instant,
35}
36
37impl RequestContext {
38    pub fn new(
39        session_id: SessionId,
40        trace_id: TraceId,
41        context_id: ContextId,
42        agent_name: AgentName,
43    ) -> Self {
44        Self {
45            auth: AuthContext {
46                auth_token: JwtToken::new(""),
47                user_id: UserId::anonymous(),
48                user_type: UserType::Anon,
49            },
50            request: RequestMetadata {
51                session_id,
52                timestamp: Instant::now(),
53                client_id: None,
54                is_tracked: true,
55                fingerprint_hash: None,
56            },
57            execution: ExecutionContext {
58                trace_id,
59                context_id,
60                task_id: None,
61                ai_tool_call_id: None,
62                mcp_execution_id: None,
63                call_source: None,
64                agent_name,
65                tool_model_config: None,
66            },
67            settings: ExecutionSettings::default(),
68            user: None,
69            start_time: Instant::now(),
70        }
71    }
72
73    pub fn with_user(mut self, user: AuthenticatedUser) -> Self {
74        self.auth.user_id = UserId::new(user.id.to_string());
75        self.user = Some(user);
76        self
77    }
78
79    pub fn with_user_id(mut self, user_id: UserId) -> Self {
80        self.auth.user_id = user_id;
81        self
82    }
83
84    pub fn with_agent_name(mut self, agent_name: AgentName) -> Self {
85        self.execution.agent_name = agent_name;
86        self
87    }
88
89    pub fn with_context_id(mut self, context_id: ContextId) -> Self {
90        self.execution.context_id = context_id;
91        self
92    }
93
94    pub fn with_task_id(mut self, task_id: TaskId) -> Self {
95        self.execution.task_id = Some(task_id);
96        self
97    }
98
99    pub fn with_task(mut self, task_id: TaskId, call_source: CallSource) -> Self {
100        self.execution.task_id = Some(task_id);
101        self.execution.call_source = Some(call_source);
102        self
103    }
104
105    pub fn with_ai_tool_call_id(mut self, ai_tool_call_id: AiToolCallId) -> Self {
106        self.execution.ai_tool_call_id = Some(ai_tool_call_id);
107        self
108    }
109
110    pub fn with_mcp_execution_id(mut self, mcp_execution_id: McpExecutionId) -> Self {
111        self.execution.mcp_execution_id = Some(mcp_execution_id);
112        self
113    }
114
115    pub fn with_client_id(mut self, client_id: ClientId) -> Self {
116        self.request.client_id = Some(client_id);
117        self
118    }
119
120    pub const fn with_user_type(mut self, user_type: UserType) -> Self {
121        self.auth.user_type = user_type;
122        self
123    }
124
125    pub fn with_auth_token(mut self, token: impl Into<String>) -> Self {
126        self.auth.auth_token = JwtToken::new(token.into());
127        self
128    }
129
130    pub const fn with_call_source(mut self, call_source: CallSource) -> Self {
131        self.execution.call_source = Some(call_source);
132        self
133    }
134
135    pub const fn with_budget(mut self, cents: i32) -> Self {
136        self.settings.max_budget_cents = Some(cents);
137        self
138    }
139
140    pub const fn with_interaction_mode(mut self, mode: UserInteractionMode) -> Self {
141        self.settings.user_interaction_mode = Some(mode);
142        self
143    }
144
145    pub const fn with_tracked(mut self, is_tracked: bool) -> Self {
146        self.request.is_tracked = is_tracked;
147        self
148    }
149
150    pub fn with_fingerprint_hash(mut self, hash: impl Into<String>) -> Self {
151        self.request.fingerprint_hash = Some(hash.into());
152        self
153    }
154
155    pub fn fingerprint_hash(&self) -> Option<&str> {
156        self.request.fingerprint_hash.as_deref()
157    }
158
159    pub fn with_tool_model_config(mut self, config: ToolModelConfig) -> Self {
160        self.execution.tool_model_config = Some(config);
161        self
162    }
163
164    pub const fn tool_model_config(&self) -> Option<&ToolModelConfig> {
165        self.execution.tool_model_config.as_ref()
166    }
167
168    pub const fn session_id(&self) -> &SessionId {
169        &self.request.session_id
170    }
171
172    pub const fn user_id(&self) -> &UserId {
173        &self.auth.user_id
174    }
175
176    pub const fn trace_id(&self) -> &TraceId {
177        &self.execution.trace_id
178    }
179
180    pub const fn context_id(&self) -> &ContextId {
181        &self.execution.context_id
182    }
183
184    pub const fn agent_name(&self) -> &AgentName {
185        &self.execution.agent_name
186    }
187
188    pub const fn auth_token(&self) -> &JwtToken {
189        &self.auth.auth_token
190    }
191
192    pub const fn user_type(&self) -> UserType {
193        self.auth.user_type
194    }
195
196    pub const fn rate_limit_tier(&self) -> RateLimitTier {
197        self.auth.user_type.rate_tier()
198    }
199
200    pub const fn task_id(&self) -> Option<&TaskId> {
201        self.execution.task_id.as_ref()
202    }
203
204    pub const fn client_id(&self) -> Option<&ClientId> {
205        self.request.client_id.as_ref()
206    }
207
208    pub const fn ai_tool_call_id(&self) -> Option<&AiToolCallId> {
209        self.execution.ai_tool_call_id.as_ref()
210    }
211
212    pub const fn mcp_execution_id(&self) -> Option<&McpExecutionId> {
213        self.execution.mcp_execution_id.as_ref()
214    }
215
216    pub const fn call_source(&self) -> Option<CallSource> {
217        self.execution.call_source
218    }
219
220    pub const fn is_authenticated(&self) -> bool {
221        self.user.is_some()
222    }
223
224    pub fn is_system(&self) -> bool {
225        self.auth.user_id.is_system() && self.execution.context_id.is_system()
226    }
227
228    pub fn elapsed(&self) -> Duration {
229        self.start_time.elapsed()
230    }
231
232    pub fn validate_task_execution(&self) -> Result<(), String> {
233        if self.execution.task_id.is_none() {
234            return Err("Missing task_id for task execution".to_string());
235        }
236        if self.execution.context_id.as_str().is_empty() {
237            return Err("Missing context_id for task execution".to_string());
238        }
239        Ok(())
240    }
241
242    pub fn validate_authenticated(&self) -> Result<(), String> {
243        if self.auth.auth_token.as_str().is_empty() {
244            return Err("Missing authentication token".to_string());
245        }
246        if self.auth.user_id.is_anonymous() {
247            return Err("User is not authenticated".to_string());
248        }
249        Ok(())
250    }
251}