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        let context_id = self.execution.context_id.as_str();
46        if !context_id.is_empty() {
47            insert_header(hdrs, headers::CONTEXT_ID, context_id);
48        }
49
50        insert_header_if_present(
51            hdrs,
52            headers::TASK_ID,
53            self.execution.task_id.as_ref().map(TaskId::as_str),
54        );
55        insert_header_if_present(
56            hdrs,
57            headers::AI_TOOL_CALL_ID,
58            self.execution.ai_tool_call_id.as_ref().map(AsRef::as_ref),
59        );
60        insert_header_if_present(
61            hdrs,
62            headers::CALL_SOURCE,
63            self.execution.call_source.as_ref().map(CallSource::as_str),
64        );
65        insert_header_if_present(
66            hdrs,
67            headers::CLIENT_ID,
68            self.request.client_id.as_ref().map(ClientId::as_str),
69        );
70
71        let auth_token = self.auth.auth_token.as_str();
72        if auth_token.is_empty() {
73            tracing::trace!(user_id = %self.auth.user_id, "No auth_token to inject - Authorization header not added");
74        } else {
75            let auth_value = format!("Bearer {}", auth_token);
76            insert_header(hdrs, headers::AUTHORIZATION, &auth_value);
77            tracing::trace!(user_id = %self.auth.user_id, "Injected Authorization header for proxy");
78        }
79
80        if let Some(user) = &self.user {
81            insert_header(hdrs, headers::PROXY_VERIFIED, "true");
82            let perms = crate::auth::permissions_to_string(&user.permissions);
83            insert_header(hdrs, headers::USER_PERMISSIONS, &perms);
84        }
85    }
86}
87
88impl ContextPropagation for RequestContext {
89    fn from_headers(hdrs: &HeaderMap) -> ContextPropagationResult<Self> {
90        let session_id = hdrs
91            .get(headers::SESSION_ID)
92            .and_then(|v| v.to_str().ok())
93            .ok_or_else(|| {
94                ContextPropagationError::MissingHeader(headers::SESSION_ID.to_string())
95            })?;
96
97        let trace_id = hdrs
98            .get(headers::TRACE_ID)
99            .and_then(|v| v.to_str().ok())
100            .ok_or_else(|| ContextPropagationError::MissingHeader(headers::TRACE_ID.to_string()))?;
101
102        let user_id = hdrs
103            .get(headers::USER_ID)
104            .and_then(|v| v.to_str().ok())
105            .ok_or_else(|| ContextPropagationError::MissingHeader(headers::USER_ID.to_string()))?;
106
107        let context_id = hdrs
108            .get(headers::CONTEXT_ID)
109            .and_then(|v| v.to_str().ok())
110            .map_or_else(
111                || ContextId::new(String::new()),
112                |s| ContextId::new(s.to_string()),
113            );
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
