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