Skip to main content

systemprompt_models/execution/context/
propagation.rs

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