systemprompt_models/execution/context/
propagation.rs

1use super::{CallSource, RequestContext};
2use anyhow::anyhow;
3use http::{HeaderMap, HeaderValue};
4use std::str::FromStr;
5use systemprompt_identifiers::{
6    AgentName, AiToolCallId, ClientId, ContextId, SessionId, TaskId, TraceId, UserId, headers,
7};
8use systemprompt_traits::{ContextPropagation, InjectContextHeaders};
9
10fn insert_header(headers: &mut HeaderMap, name: &'static str, value: &str) {
11    match HeaderValue::from_str(value) {
12        Ok(val) => {
13            headers.insert(name, val);
14        },
15        Err(e) => {
16            tracing::warn!(
17                header = %name,
18                value = %value,
19                error = %e,
20                "Invalid header value - header not inserted"
21            );
22        },
23    }
24}
25
26fn insert_header_if_present(headers: &mut HeaderMap, name: &'static str, value: Option<&str>) {
27    if let Some(v) = value {
28        insert_header(headers, name, v);
29    }
30}
31
32impl InjectContextHeaders for RequestContext {
33    fn inject_headers(&self, hdrs: &mut HeaderMap) {
34        insert_header(hdrs, headers::SESSION_ID, self.request.session_id.as_str());
35        insert_header(hdrs, headers::TRACE_ID, self.execution.trace_id.as_str());
36        insert_header(hdrs, headers::USER_ID, self.auth.user_id.as_str());
37        insert_header(hdrs, headers::USER_TYPE, self.auth.user_type.as_str());
38        insert_header(
39            hdrs,
40            headers::AGENT_NAME,
41            self.execution.agent_name.as_str(),
42        );
43
44        let context_id = self.execution.context_id.as_str();
45        if !context_id.is_empty() {
46            insert_header(hdrs, headers::CONTEXT_ID, context_id);
47        }
48
49        insert_header_if_present(
50            hdrs,
51            headers::TASK_ID,
52            self.execution.task_id.as_ref().map(TaskId::as_str),
53        );
54        insert_header_if_present(
55            hdrs,
56            headers::AI_TOOL_CALL_ID,
57            self.execution.ai_tool_call_id.as_ref().map(AsRef::as_ref),
58        );
59        insert_header_if_present(
60            hdrs,
61            headers::CALL_SOURCE,
62            self.execution.call_source.as_ref().map(CallSource::as_str),
63        );
64        insert_header_if_present(
65            hdrs,
66            headers::CLIENT_ID,
67            self.request.client_id.as_ref().map(ClientId::as_str),
68        );
69
70        let auth_token = self.auth.auth_token.as_str();
71        if auth_token.is_empty() {
72            tracing::trace!(user_id = %self.auth.user_id, "No auth_token to inject - Authorization header not added");
73        } else {
74            let auth_value = format!("Bearer {}", auth_token);
75            insert_header(hdrs, headers::AUTHORIZATION, &auth_value);
76            tracing::trace!(user_id = %self.auth.user_id, "Injected Authorization header for proxy");
77        }
78
79        if let Some(user) = &self.user {
80            insert_header(hdrs, headers::PROXY_VERIFIED, "true");
81            let perms = crate::auth::permissions_to_string(&user.permissions);
82            insert_header(hdrs, headers::USER_PERMISSIONS, &perms);
83        }
84    }
85}
86
87impl ContextPropagation for RequestContext {
88    fn from_headers(hdrs: &HeaderMap) -> anyhow::Result<Self> {
89        let session_id = hdrs
90            .get(headers::SESSION_ID)
91            .and_then(|v| v.to_str().ok())
92            .ok_or_else(|| anyhow!("Missing {} header", headers::SESSION_ID))?;
93
94        let trace_id = hdrs
95            .get(headers::TRACE_ID)
96            .and_then(|v| v.to_str().ok())
97            .ok_or_else(|| anyhow!("Missing {} header", headers::TRACE_ID))?;
98
99        let user_id = hdrs
100            .get(headers::USER_ID)
101            .and_then(|v| v.to_str().ok())
102            .ok_or_else(|| anyhow!("Missing {} header", headers::USER_ID))?;
103
104        let context_id = hdrs
105            .get(headers::CONTEXT_ID)
106            .and_then(|v| v.to_str().ok())
107            .map_or_else(
108                || ContextId::new(String::new()),
109                |s| ContextId::new(s.to_string()),
110            );
111
112        let agent_name = hdrs
113            .get(headers::AGENT_NAME)
114            .and_then(|v| v.to_str().ok())
115            .ok_or_else(|| {
116                anyhow!(
117                    "Missing {} header - all requests must have agent context",
118                    headers::AGENT_NAME
119                )
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::from(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
188                    .parse::<uuid::Uuid>()
189                    .map_err(|e| anyhow!("Invalid UUID in {} header: {}", headers::USER_ID, e))?;
190                let user = crate::auth::AuthenticatedUser::new(
191                    user_id_uuid,
192                    String::new(),
193                    String::new(),
194                    permissions,
195                );
196                ctx = ctx.with_user(user);
197            }
198        }
199
200        Ok(ctx)
201    }
202
203    fn to_headers(&self) -> HeaderMap {
204        let mut headers = HeaderMap::new();
205        self.inject_headers(&mut headers);
206        headers
207    }
208}
systemprompt_models/execution/context/propagation.rs

systemprompt_models/execution/context/
propagation.rs
