systemprompt_models/execution/context/
propagation.rs

1use super::{CallSource, RequestContext};
2use http::{HeaderMap, HeaderValue};
3use std::str::FromStr;
4use systemprompt_identifiers::{
5    AgentName, AiToolCallId, ClientId, ContextId, SessionId, TaskId, TraceId, UserId, headers,
6};
7use systemprompt_traits::{
8    ContextPropagation, ContextPropagationError, ContextPropagationResult, InjectContextHeaders,
9};
10
11fn insert_header(headers: &mut HeaderMap, name: &'static str, value: &str) {
12    match HeaderValue::from_str(value) {
13        Ok(val) => {
14            headers.insert(name, val);
15        },
16        Err(e) => {
17            tracing::warn!(
18                header = %name,
19                value = %value,
20                error = %e,
21                "Invalid header value - header not inserted"
22            );
23        },
24    }
25}
26
27fn insert_header_if_present(headers: &mut HeaderMap, name: &'static str, value: Option<&str>) {
28    if let Some(v) = value {
29        insert_header(headers, name, v);
30    }
31}
32
33impl InjectContextHeaders for RequestContext {
34    fn inject_headers(&self, hdrs: &mut HeaderMap) {
35        insert_header(hdrs, headers::SESSION_ID, self.request.session_id.as_str());
36        insert_header(hdrs, headers::TRACE_ID, self.execution.trace_id.as_str());
37        insert_header(hdrs, headers::USER_ID, self.auth.user_id.as_str());
38        insert_header(hdrs, headers::USER_TYPE, self.auth.user_type.as_str());
39        insert_header(
40            hdrs,
41            headers::AGENT_NAME,
42            self.execution.agent_name.as_str(),
43        );
44
45        insert_header(
46            hdrs,
47            headers::CONTEXT_ID,
48            self.execution.context_id.as_str(),
49        );
50
51        insert_header_if_present(
52            hdrs,
53            headers::TASK_ID,
54            self.execution.task_id.as_ref().map(TaskId::as_str),
55        );
56        insert_header_if_present(
57            hdrs,
58            headers::AI_TOOL_CALL_ID,
59            self.execution.ai_tool_call_id.as_ref().map(AsRef::as_ref),
60        );
61        insert_header_if_present(
62            hdrs,
63            headers::CALL_SOURCE,
64            self.execution.call_source.as_ref().map(CallSource::as_str),
65        );
66        insert_header_if_present(
67            hdrs,
68            headers::CLIENT_ID,
69            self.request.client_id.as_ref().map(ClientId::as_str),
70        );
71
72        let auth_token = self.auth.auth_token.as_str();
73        if auth_token.is_empty() {
74            tracing::trace!(user_id = %self.auth.user_id, "No auth_token to inject - Authorization header not added");
75        } else {
76            let auth_value = format!("Bearer {}", auth_token);
77            insert_header(hdrs, headers::AUTHORIZATION, &auth_value);
78            tracing::trace!(user_id = %self.auth.user_id, "Injected Authorization header for proxy");
79        }
80
81        if let Some(user) = &self.user {
82            insert_header(hdrs, headers::PROXY_VERIFIED, "true");
83            let perms = crate::auth::permissions_to_string(&user.permissions);
84            insert_header(hdrs, headers::USER_PERMISSIONS, &perms);
85        }
86    }
87}
88
89impl ContextPropagation for RequestContext {
90    fn from_headers(hdrs: &HeaderMap) -> ContextPropagationResult<Self> {
91        let session_id = hdrs
92            .get(headers::SESSION_ID)
93            .and_then(|v| v.to_str().ok())
94            .ok_or_else(|| {
95                ContextPropagationError::MissingHeader(headers::SESSION_ID.to_string())
96            })?;
97
98        let trace_id = hdrs
99            .get(headers::TRACE_ID)
100            .and_then(|v| v.to_str().ok())
101            .ok_or_else(|| ContextPropagationError::MissingHeader(headers::TRACE_ID.to_string()))?;
102
103        let user_id = hdrs
104            .get(headers::USER_ID)
105            .and_then(|v| v.to_str().ok())
106            .ok_or_else(|| ContextPropagationError::MissingHeader(headers::USER_ID.to_string()))?;
107
108        let context_id = hdrs
109            .get(headers::CONTEXT_ID)
110            .and_then(|v| v.to_str().ok())
111            .filter(|s| !s.is_empty())
112            .and_then(|s| ContextId::try_new(s).ok())
113            .unwrap_or_else(ContextId::generate);
114
115        let agent_name = hdrs
116            .get(headers::AGENT_NAME)
117            .and_then(|v| v.to_str().ok())
118            .ok_or_else(|| {
119                ContextPropagationError::MissingHeader(headers::AGENT_NAME.to_string())
120            })?;
121
122        let task_id = hdrs
123            .get(headers::TASK_ID)
124            .and_then(|v| v.to_str().ok())
125            .map(|s| TaskId::new(s.to_string()));
126
127        let ai_tool_call_id = hdrs
128            .get(headers::AI_TOOL_CALL_ID)
129            .and_then(|v| v.to_str().ok())
130            .map(|s| AiToolCallId::new(s.to_string()));
131
132        let call_source = hdrs
133            .get(headers::CALL_SOURCE)
134            .and_then(|v| v.to_str().ok())
135            .and_then(|s| CallSource::from_str(s).ok());
136
137        let client_id = hdrs
138            .get(headers::CLIENT_ID)
139            .and_then(|v| v.to_str().ok())
140            .map(|s| ClientId::new(s.to_string()));
141
142        let auth_token = hdrs
143            .get(headers::AUTHORIZATION)
144            .and_then(|v| v.to_str().ok())
145            .and_then(|s| s.strip_prefix("Bearer "))
146            .map(ToString::to_string);
147
148        let mut ctx = Self::new(
149            SessionId::new(session_id.to_string()),
150            TraceId::new(trace_id.to_string()),
151            context_id,
152            AgentName::new(agent_name.to_string()),
153        )
154        .with_user_id(UserId::new(user_id.to_string()));
155
156        if let Some(tid) = task_id {
157            ctx = ctx.with_task_id(tid);
158        }
159
160        if let Some(ai_id) = ai_tool_call_id {
161            ctx = ctx.with_ai_tool_call_id(ai_id);
162        }
163
164        if let Some(cs) = call_source {
165            ctx = ctx.with_call_source(cs);
166        }
167
168        if let Some(cid) = client_id {
169            ctx = ctx.with_client_id(cid);
170        }
171
172        if let Some(token) = auth_token {
173            ctx = ctx.with_auth_token(token);
174        }
175
176        let proxy_verified = hdrs
177            .get(headers::PROXY_VERIFIED)
178            .and_then(|v| v.to_str().ok())
179            .is_some_and(|v| v == "true");
180
181        if proxy_verified {
182            if let Some(permissions) = hdrs
183                .get(headers::USER_PERMISSIONS)
184                .and_then(|v| v.to_str().ok())
185                .and_then(|s| crate::auth::parse_permissions(s).ok())
186            {
187                let user_id_uuid = user_id.parse::<uuid::Uuid>().map_err(|e| {
188                    ContextPropagationError::InvalidHeader {
189                        name: headers::USER_ID.to_string(),
190                        message: format!("invalid UUID: {e}"),
191                    }
192                })?;
193                let user = crate::auth::AuthenticatedUser::new(
194                    user_id_uuid,
195                    String::new(),
196                    String::new(),
197                    permissions,
198                );
199                ctx = ctx.with_user(user);
200            }
201        }
202
203        Ok(ctx)
204    }
205
206    fn to_headers(&self) -> HeaderMap {
207        let mut headers = HeaderMap::new();
208        self.inject_headers(&mut headers);
209        headers
210    }
211}
systemprompt_models/execution/context/propagation.rs

systemprompt_models/execution/context/
propagation.rs
