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    headers, AgentName, AiToolCallId, ClientId, ContextId, SessionId, TaskId, TraceId, UserId,
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}
80
81impl ContextPropagation for RequestContext {
82    fn from_headers(hdrs: &HeaderMap) -> anyhow::Result<Self> {
83        let session_id = hdrs
84            .get(headers::SESSION_ID)
85            .and_then(|v| v.to_str().ok())
86            .ok_or_else(|| anyhow!("Missing {} header", headers::SESSION_ID))?;
87
88        let trace_id = hdrs
89            .get(headers::TRACE_ID)
90            .and_then(|v| v.to_str().ok())
91            .ok_or_else(|| anyhow!("Missing {} header", headers::TRACE_ID))?;
92
93        let user_id = hdrs
94            .get(headers::USER_ID)
95            .and_then(|v| v.to_str().ok())
96            .ok_or_else(|| anyhow!("Missing {} header", headers::USER_ID))?;
97
98        let context_id = hdrs
99            .get(headers::CONTEXT_ID)
100            .and_then(|v| v.to_str().ok())
101            .map_or_else(
102                || ContextId::new(String::new()),
103                |s| ContextId::new(s.to_string()),
104            );
105
106        let agent_name = hdrs
107            .get(headers::AGENT_NAME)
108            .and_then(|v| v.to_str().ok())
109            .ok_or_else(|| {
110                anyhow!(
111                    "Missing {} header - all requests must have agent context",
112                    headers::AGENT_NAME
113                )
114            })?;
115
116        let task_id = hdrs
117            .get(headers::TASK_ID)
118            .and_then(|v| v.to_str().ok())
119            .map(|s| TaskId::new(s.to_string()));
120
121        let ai_tool_call_id = hdrs
122            .get(headers::AI_TOOL_CALL_ID)
123            .and_then(|v| v.to_str().ok())
124            .map(|s| AiToolCallId::from(s.to_string()));
125
126        let call_source = hdrs
127            .get(headers::CALL_SOURCE)
128            .and_then(|v| v.to_str().ok())
129            .and_then(|s| CallSource::from_str(s).ok());
130
131        let client_id = hdrs
132            .get(headers::CLIENT_ID)
133            .and_then(|v| v.to_str().ok())
134            .map(|s| ClientId::new(s.to_string()));
135
136        let auth_token = hdrs
137            .get(headers::AUTHORIZATION)
138            .and_then(|v| v.to_str().ok())
139            .and_then(|s| s.strip_prefix("Bearer "))
140            .map(ToString::to_string);
141
142        let mut ctx = Self::new(
143            SessionId::new(session_id.to_string()),
144            TraceId::new(trace_id.to_string()),
145            context_id,
146            AgentName::new(agent_name.to_string()),
147        )
148        .with_user_id(UserId::new(user_id.to_string()));
149
150        if let Some(tid) = task_id {
151            ctx = ctx.with_task_id(tid);
152        }
153
154        if let Some(ai_id) = ai_tool_call_id {
155            ctx = ctx.with_ai_tool_call_id(ai_id);
156        }
157
158        if let Some(cs) = call_source {
159            ctx = ctx.with_call_source(cs);
160        }
161
162        if let Some(cid) = client_id {
163            ctx = ctx.with_client_id(cid);
164        }
165
166        if let Some(token) = auth_token {
167            ctx = ctx.with_auth_token(token);
168        }
169
170        Ok(ctx)
171    }
172
173    fn to_headers(&self) -> HeaderMap {
174        let mut headers = HeaderMap::new();
175        self.inject_headers(&mut headers);
176        headers
177    }
178}
systemprompt_models/execution/context/propagation.rs

systemprompt_models/execution/context/
propagation.rs
