systemprompt_models/execution/context/
mod.rs

1//! Request context for execution tracking.
2
3mod call_source;
4mod context_error;
5mod context_types;
6
7pub use call_source::CallSource;
8pub use context_error::{ContextExtractionError, ContextIdSource, TASK_BASED_CONTEXT_MARKER};
9pub use context_types::{
10    AuthContext, ExecutionContext, ExecutionSettings, RequestMetadata, UserInteractionMode,
11};
12
13use crate::ai::ToolModelConfig;
14use crate::auth::{AuthenticatedUser, RateLimitTier, UserType};
15use anyhow::anyhow;
16use http::{HeaderMap, HeaderValue};
17use serde::{Deserialize, Serialize};
18use std::str::FromStr;
19use std::time::{Duration, Instant};
20use systemprompt_identifiers::{
21    headers, AgentName, AiToolCallId, ClientId, ContextId, JwtToken, McpExecutionId, SessionId,
22    TaskId, TraceId, UserId,
23};
24use systemprompt_traits::{ContextPropagation, InjectContextHeaders};
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct RequestContext {
28    pub auth: AuthContext,
29    pub request: RequestMetadata,
30    pub execution: ExecutionContext,
31    pub settings: ExecutionSettings,
32
33    #[serde(skip)]
34    pub user: Option<AuthenticatedUser>,
35
36    #[serde(skip, default = "Instant::now")]
37    pub start_time: Instant,
38}
39
40impl RequestContext {
41    /// Creates a new `RequestContext` - the ONLY way to construct a context.
42    ///
43    /// This is the single constructor for `RequestContext`. All contexts must
44    /// be created through this method, ensuring consistent initialization.
45    ///
46    /// # Required Fields
47    /// - `session_id`: Identifies the user session
48    /// - `trace_id`: For distributed tracing
49    /// - `context_id`: Conversation/execution context (empty string for
50    ///   user-level contexts)
51    /// - `agent_name`: The agent handling this request (use
52    ///   `AgentName::system()` for system operations)
53    ///
54    /// # Optional Fields
55    /// Use builder methods to set optional fields:
56    /// - `.with_user_id()` - Set the authenticated user
57    /// - `.with_auth_token()` - Set the JWT token
58    /// - `.with_user_type()` - Set user type (Admin, Standard, Anon)
59    /// - `.with_task_id()` - Set task ID for AI operations
60    /// - `.with_client_id()` - Set client ID
61    /// - `.with_call_source()` - Set call source (Agentic, Direct, Ephemeral)
62    ///
63    /// # Example
64    /// ```
65    /// # use systemprompt_models::execution::context::RequestContext;
66    /// # use systemprompt_identifiers::{SessionId, TraceId, ContextId, AgentName, UserId};
67    /// # use systemprompt_models::auth::UserType;
68    /// let ctx = RequestContext::new(
69    ///     SessionId::new("sess_123".to_string()),
70    ///     TraceId::new("trace_456".to_string()),
71    ///     ContextId::new("ctx_789".to_string()),
72    ///     AgentName::new("my-agent".to_string()),
73    /// )
74    /// .with_user_id(UserId::new("user_123".to_string()))
75    /// .with_auth_token("jwt_token_here")
76    /// .with_user_type(UserType::User);
77    /// ```
78    pub fn new(
79        session_id: SessionId,
80        trace_id: TraceId,
81        context_id: ContextId,
82        agent_name: AgentName,
83    ) -> Self {
84        Self {
85            auth: AuthContext {
86                auth_token: JwtToken::new(""),
87                user_id: UserId::anonymous(),
88                user_type: UserType::Anon,
89            },
90            request: RequestMetadata {
91                session_id,
92                timestamp: Instant::now(),
93                client_id: None,
94                is_tracked: true,
95                fingerprint_hash: None,
96            },
97            execution: ExecutionContext {
98                trace_id,
99                context_id,
100                task_id: None,
101                ai_tool_call_id: None,
102                mcp_execution_id: None,
103                call_source: None,
104                agent_name,
105                tool_model_config: None,
106            },
107            settings: ExecutionSettings::default(),
108            user: None,
109            start_time: Instant::now(),
110        }
111    }
112
113    pub fn with_user(mut self, user: AuthenticatedUser) -> Self {
114        self.auth.user_id = UserId::new(user.id.to_string());
115        self.user = Some(user);
116        self
117    }
118
119    pub fn with_user_id(mut self, user_id: UserId) -> Self {
120        self.auth.user_id = user_id;
121        self
122    }
123
124    pub fn with_agent_name(mut self, agent_name: AgentName) -> Self {
125        self.execution.agent_name = agent_name;
126        self
127    }
128
129    pub fn with_context_id(mut self, context_id: ContextId) -> Self {
130        self.execution.context_id = context_id;
131        self
132    }
133
134    pub fn with_task_id(mut self, task_id: TaskId) -> Self {
135        self.execution.task_id = Some(task_id);
136        self
137    }
138
139    pub fn with_task(mut self, task_id: TaskId, call_source: CallSource) -> Self {
140        self.execution.task_id = Some(task_id);
141        self.execution.call_source = Some(call_source);
142        self
143    }
144
145    pub fn with_ai_tool_call_id(mut self, ai_tool_call_id: AiToolCallId) -> Self {
146        self.execution.ai_tool_call_id = Some(ai_tool_call_id);
147        self
148    }
149
150    pub fn with_mcp_execution_id(mut self, mcp_execution_id: McpExecutionId) -> Self {
151        self.execution.mcp_execution_id = Some(mcp_execution_id);
152        self
153    }
154
155    pub fn with_client_id(mut self, client_id: ClientId) -> Self {
156        self.request.client_id = Some(client_id);
157        self
158    }
159
160    pub const fn with_user_type(mut self, user_type: UserType) -> Self {
161        self.auth.user_type = user_type;
162        self
163    }
164
165    pub fn with_auth_token(mut self, token: impl Into<String>) -> Self {
166        self.auth.auth_token = JwtToken::new(token.into());
167        self
168    }
169
170    pub const fn with_call_source(mut self, call_source: CallSource) -> Self {
171        self.execution.call_source = Some(call_source);
172        self
173    }
174
175    pub const fn with_budget(mut self, cents: i32) -> Self {
176        self.settings.max_budget_cents = Some(cents);
177        self
178    }
179
180    pub const fn with_interaction_mode(mut self, mode: UserInteractionMode) -> Self {
181        self.settings.user_interaction_mode = Some(mode);
182        self
183    }
184
185    pub const fn with_tracked(mut self, is_tracked: bool) -> Self {
186        self.request.is_tracked = is_tracked;
187        self
188    }
189
190    pub fn with_fingerprint_hash(mut self, hash: impl Into<String>) -> Self {
191        self.request.fingerprint_hash = Some(hash.into());
192        self
193    }
194
195    pub fn fingerprint_hash(&self) -> Option<&str> {
196        self.request.fingerprint_hash.as_deref()
197    }
198
199    pub fn with_tool_model_config(mut self, config: ToolModelConfig) -> Self {
200        self.execution.tool_model_config = Some(config);
201        self
202    }
203
204    pub const fn tool_model_config(&self) -> Option<&ToolModelConfig> {
205        self.execution.tool_model_config.as_ref()
206    }
207
208    pub const fn session_id(&self) -> &SessionId {
209        &self.request.session_id
210    }
211
212    pub const fn user_id(&self) -> &UserId {
213        &self.auth.user_id
214    }
215
216    pub const fn trace_id(&self) -> &TraceId {
217        &self.execution.trace_id
218    }
219
220    pub const fn context_id(&self) -> &ContextId {
221        &self.execution.context_id
222    }
223
224    pub const fn agent_name(&self) -> &AgentName {
225        &self.execution.agent_name
226    }
227
228    pub const fn auth_token(&self) -> &JwtToken {
229        &self.auth.auth_token
230    }
231
232    pub const fn user_type(&self) -> UserType {
233        self.auth.user_type
234    }
235
236    pub const fn rate_limit_tier(&self) -> RateLimitTier {
237        self.auth.user_type.rate_tier()
238    }
239
240    pub const fn task_id(&self) -> Option<&TaskId> {
241        self.execution.task_id.as_ref()
242    }
243
244    pub const fn client_id(&self) -> Option<&ClientId> {
245        self.request.client_id.as_ref()
246    }
247
248    pub const fn ai_tool_call_id(&self) -> Option<&AiToolCallId> {
249        self.execution.ai_tool_call_id.as_ref()
250    }
251
252    pub const fn mcp_execution_id(&self) -> Option<&McpExecutionId> {
253        self.execution.mcp_execution_id.as_ref()
254    }
255
256    pub const fn call_source(&self) -> Option<CallSource> {
257        self.execution.call_source
258    }
259
260    pub const fn is_authenticated(&self) -> bool {
261        self.user.is_some()
262    }
263
264    pub fn is_system(&self) -> bool {
265        self.auth.user_id.is_system() && self.execution.context_id.is_system()
266    }
267
268    pub fn elapsed(&self) -> Duration {
269        self.start_time.elapsed()
270    }
271
272    pub fn validate_task_execution(&self) -> Result<(), String> {
273        if self.execution.task_id.is_none() {
274            return Err("Missing task_id for task execution".to_string());
275        }
276        if self.execution.context_id.as_str().is_empty() {
277            return Err("Missing context_id for task execution".to_string());
278        }
279        Ok(())
280    }
281
282    pub fn validate_authenticated(&self) -> Result<(), String> {
283        if self.auth.auth_token.as_str().is_empty() {
284            return Err("Missing authentication token".to_string());
285        }
286        if self.auth.user_id.is_anonymous() {
287            return Err("User is not authenticated".to_string());
288        }
289        Ok(())
290    }
291}
292
293fn insert_header(headers: &mut HeaderMap, name: &'static str, value: &str) {
294    if let Ok(val) = HeaderValue::from_str(value) {
295        headers.insert(name, val);
296    }
297}
298
299fn insert_header_if_present(headers: &mut HeaderMap, name: &'static str, value: Option<&str>) {
300    if let Some(v) = value {
301        insert_header(headers, name, v);
302    }
303}
304
305impl InjectContextHeaders for RequestContext {
306    fn inject_headers(&self, hdrs: &mut HeaderMap) {
307        insert_header(hdrs, headers::SESSION_ID, self.request.session_id.as_str());
308        insert_header(hdrs, headers::TRACE_ID, self.execution.trace_id.as_str());
309        insert_header(hdrs, headers::USER_ID, self.auth.user_id.as_str());
310        insert_header(hdrs, headers::USER_TYPE, self.auth.user_type.as_str());
311        insert_header(
312            hdrs,
313            headers::AGENT_NAME,
314            self.execution.agent_name.as_str(),
315        );
316
317        let context_id = self.execution.context_id.as_str();
318        if !context_id.is_empty() {
319            insert_header(hdrs, headers::CONTEXT_ID, context_id);
320        }
321
322        insert_header_if_present(
323            hdrs,
324            headers::TASK_ID,
325            self.execution.task_id.as_ref().map(TaskId::as_str),
326        );
327        insert_header_if_present(
328            hdrs,
329            headers::AI_TOOL_CALL_ID,
330            self.execution.ai_tool_call_id.as_ref().map(AsRef::as_ref),
331        );
332        insert_header_if_present(
333            hdrs,
334            headers::CALL_SOURCE,
335            self.execution.call_source.as_ref().map(CallSource::as_str),
336        );
337        insert_header_if_present(
338            hdrs,
339            headers::CLIENT_ID,
340            self.request.client_id.as_ref().map(ClientId::as_str),
341        );
342
343        let auth_token = self.auth.auth_token.as_str();
344        if auth_token.is_empty() {
345            tracing::trace!(user_id = %self.auth.user_id, "No auth_token to inject - Authorization header not added");
346        } else {
347            let auth_value = format!("Bearer {}", auth_token);
348            insert_header(hdrs, headers::AUTHORIZATION, &auth_value);
349            tracing::trace!(user_id = %self.auth.user_id, "Injected Authorization header for proxy");
350        }
351    }
352}
353
354impl ContextPropagation for RequestContext {
355    fn from_headers(hdrs: &HeaderMap) -> anyhow::Result<Self> {
356        let session_id = hdrs
357            .get(headers::SESSION_ID)
358            .and_then(|v| v.to_str().ok())
359            .ok_or_else(|| anyhow!("Missing {} header", headers::SESSION_ID))?;
360
361        let trace_id = hdrs
362            .get(headers::TRACE_ID)
363            .and_then(|v| v.to_str().ok())
364            .ok_or_else(|| anyhow!("Missing {} header", headers::TRACE_ID))?;
365
366        let user_id = hdrs
367            .get(headers::USER_ID)
368            .and_then(|v| v.to_str().ok())
369            .ok_or_else(|| anyhow!("Missing {} header", headers::USER_ID))?;
370
371        let context_id = hdrs
372            .get(headers::CONTEXT_ID)
373            .and_then(|v| v.to_str().ok())
374            .map_or_else(
375                || ContextId::new(String::new()),
376                |s| ContextId::new(s.to_string()),
377            );
378
379        let agent_name = hdrs
380            .get(headers::AGENT_NAME)
381            .and_then(|v| v.to_str().ok())
382            .ok_or_else(|| {
383                anyhow!(
384                    "Missing {} header - all requests must have agent context",
385                    headers::AGENT_NAME
386                )
387            })?;
388
389        let task_id = hdrs
390            .get(headers::TASK_ID)
391            .and_then(|v| v.to_str().ok())
392            .map(|s| TaskId::new(s.to_string()));
393
394        let ai_tool_call_id = hdrs
395            .get(headers::AI_TOOL_CALL_ID)
396            .and_then(|v| v.to_str().ok())
397            .map(|s| AiToolCallId::from(s.to_string()));
398
399        let call_source = hdrs
400            .get(headers::CALL_SOURCE)
401            .and_then(|v| v.to_str().ok())
402            .and_then(|s| CallSource::from_str(s).ok());
403
404        let client_id = hdrs
405            .get(headers::CLIENT_ID)
406            .and_then(|v| v.to_str().ok())
407            .map(|s| ClientId::new(s.to_string()));
408
409        let mut ctx = Self::new(
410            SessionId::new(session_id.to_string()),
411            TraceId::new(trace_id.to_string()),
412            context_id,
413            AgentName::new(agent_name.to_string()),
414        )
415        .with_user_id(UserId::new(user_id.to_string()));
416
417        if let Some(tid) = task_id {
418            ctx = ctx.with_task_id(tid);
419        }
420
421        if let Some(ai_id) = ai_tool_call_id {
422            ctx = ctx.with_ai_tool_call_id(ai_id);
423        }
424
425        if let Some(cs) = call_source {
426            ctx = ctx.with_call_source(cs);
427        }
428
429        if let Some(cid) = client_id {
430            ctx = ctx.with_client_id(cid);
431        }
432
433        Ok(ctx)
434    }
435
436    fn to_headers(&self) -> HeaderMap {
437        let mut headers = HeaderMap::new();
438        self.inject_headers(&mut headers);
439        headers
440    }
441}
systemprompt_models/execution/context/mod.rs

systemprompt_models/execution/context/
mod.rs
