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