treeship_core/session/
context.rs

1//! Cross-tool and cross-host context propagation.
2//!
3//! Propagates session identity through environment variables, HTTP headers,
4//! and CLI wrappers so that spawned processes and remote agents inherit
5//! session context.
6
7use serde::{Deserialize, Serialize};
8
9use super::event::{generate_span_id, generate_trace_id};
10
11/// Environment variable prefix for Treeship context.
12const ENV_PREFIX: &str = "TREESHIP_";
13
14/// HTTP header prefix for Treeship context.
15const HEADER_PREFIX: &str = "x-treeship-";
16
17/// Field names used for both env vars and headers.
18const FIELD_SESSION_ID: &str = "SESSION_ID";
19const FIELD_TRACE_ID: &str = "TRACE_ID";
20const FIELD_SPAN_ID: &str = "SPAN_ID";
21const FIELD_PARENT_SPAN_ID: &str = "PARENT_SPAN_ID";
22const FIELD_AGENT_ID: &str = "AGENT_ID";
23const FIELD_AGENT_INSTANCE_ID: &str = "AGENT_INSTANCE_ID";
24const FIELD_WORKSPACE_ID: &str = "WORKSPACE_ID";
25const FIELD_MISSION_ID: &str = "MISSION_ID";
26const FIELD_HOST_ID: &str = "HOST_ID";
27const FIELD_TOOL_RUNTIME_ID: &str = "TOOL_RUNTIME_ID";
28
29/// Context propagated across tool and host boundaries.
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct PropagationContext {
32    pub session_id: String,
33    pub trace_id: String,
34    pub span_id: String,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub parent_span_id: Option<String>,
37    pub agent_id: String,
38    pub agent_instance_id: String,
39    #[serde(skip_serializing_if = "Option::is_none")]
40    pub workspace_id: Option<String>,
41    #[serde(skip_serializing_if = "Option::is_none")]
42    pub mission_id: Option<String>,
43    pub host_id: String,
44    #[serde(skip_serializing_if = "Option::is_none")]
45    pub tool_runtime_id: Option<String>,
46}
47
48impl PropagationContext {
49    /// Read context from environment variables.
50    ///
51    /// Returns `None` if required fields (`TREESHIP_SESSION_ID`, `TREESHIP_TRACE_ID`)
52    /// are not present.
53    pub fn from_env() -> Option<Self> {
54        let session_id = std::env::var(format!("{ENV_PREFIX}{FIELD_SESSION_ID}")).ok()?;
55        let trace_id = std::env::var(format!("{ENV_PREFIX}{FIELD_TRACE_ID}"))
56            .unwrap_or_else(|_| generate_trace_id());
57
58        Some(Self {
59            session_id,
60            trace_id,
61            span_id: std::env::var(format!("{ENV_PREFIX}{FIELD_SPAN_ID}"))
62                .unwrap_or_else(|_| generate_span_id()),
63            parent_span_id: std::env::var(format!("{ENV_PREFIX}{FIELD_PARENT_SPAN_ID}")).ok(),
64            agent_id: std::env::var(format!("{ENV_PREFIX}{FIELD_AGENT_ID}"))
65                .unwrap_or_else(|_| "agent://unknown".into()),
66            agent_instance_id: std::env::var(format!("{ENV_PREFIX}{FIELD_AGENT_INSTANCE_ID}"))
67                .unwrap_or_else(|_| "ai_unknown".into()),
68            workspace_id: std::env::var(format!("{ENV_PREFIX}{FIELD_WORKSPACE_ID}")).ok(),
69            mission_id: std::env::var(format!("{ENV_PREFIX}{FIELD_MISSION_ID}")).ok(),
70            host_id: std::env::var(format!("{ENV_PREFIX}{FIELD_HOST_ID}"))
71                .unwrap_or_else(|_| default_host_id()),
72            tool_runtime_id: std::env::var(format!("{ENV_PREFIX}{FIELD_TOOL_RUNTIME_ID}")).ok(),
73        })
74    }
75
76    /// Inject context as environment variables on a Command builder.
77    pub fn inject_env(&self, cmd: &mut std::process::Command) {
78        cmd.env(format!("{ENV_PREFIX}{FIELD_SESSION_ID}"), &self.session_id);
79        cmd.env(format!("{ENV_PREFIX}{FIELD_TRACE_ID}"), &self.trace_id);
80        cmd.env(format!("{ENV_PREFIX}{FIELD_SPAN_ID}"), &self.span_id);
81        if let Some(ref psid) = self.parent_span_id {
82            cmd.env(format!("{ENV_PREFIX}{FIELD_PARENT_SPAN_ID}"), psid);
83        }
84        cmd.env(format!("{ENV_PREFIX}{FIELD_AGENT_ID}"), &self.agent_id);
85        cmd.env(
86            format!("{ENV_PREFIX}{FIELD_AGENT_INSTANCE_ID}"),
87            &self.agent_instance_id,
88        );
89        if let Some(ref wid) = self.workspace_id {
90            cmd.env(format!("{ENV_PREFIX}{FIELD_WORKSPACE_ID}"), wid);
91        }
92        if let Some(ref mid) = self.mission_id {
93            cmd.env(format!("{ENV_PREFIX}{FIELD_MISSION_ID}"), mid);
94        }
95        cmd.env(format!("{ENV_PREFIX}{FIELD_HOST_ID}"), &self.host_id);
96        if let Some(ref trid) = self.tool_runtime_id {
97            cmd.env(format!("{ENV_PREFIX}{FIELD_TOOL_RUNTIME_ID}"), trid);
98        }
99    }
100
101    /// Produce HTTP header pairs for outbound requests.
102    pub fn to_headers(&self) -> Vec<(String, String)> {
103        let mut h = vec![
104            (
105                format!("{HEADER_PREFIX}session-id"),
106                self.session_id.clone(),
107            ),
108            (format!("{HEADER_PREFIX}trace-id"), self.trace_id.clone()),
109            (format!("{HEADER_PREFIX}span-id"), self.span_id.clone()),
110            (format!("{HEADER_PREFIX}agent-id"), self.agent_id.clone()),
111            (
112                format!("{HEADER_PREFIX}agent-instance-id"),
113                self.agent_instance_id.clone(),
114            ),
115            (format!("{HEADER_PREFIX}host-id"), self.host_id.clone()),
116        ];
117        if let Some(ref psid) = self.parent_span_id {
118            h.push((format!("{HEADER_PREFIX}parent-span-id"), psid.clone()));
119        }
120        if let Some(ref wid) = self.workspace_id {
121            h.push((format!("{HEADER_PREFIX}workspace-id"), wid.clone()));
122        }
123        if let Some(ref mid) = self.mission_id {
124            h.push((format!("{HEADER_PREFIX}mission-id"), mid.clone()));
125        }
126        if let Some(ref trid) = self.tool_runtime_id {
127            h.push((format!("{HEADER_PREFIX}tool-runtime-id"), trid.clone()));
128        }
129        h
130    }
131
132    /// Parse context from HTTP header pairs.
133    pub fn from_headers(headers: &[(String, String)]) -> Option<Self> {
134        let get = |name: &str| -> Option<String> {
135            let key = format!("{HEADER_PREFIX}{name}");
136            headers
137                .iter()
138                .find(|(k, _)| k.eq_ignore_ascii_case(&key))
139                .map(|(_, v)| v.clone())
140        };
141
142        let session_id = get("session-id")?;
143        let trace_id = get("trace-id").unwrap_or_else(generate_trace_id);
144
145        Some(Self {
146            session_id,
147            trace_id,
148            span_id: get("span-id").unwrap_or_else(generate_span_id),
149            parent_span_id: get("parent-span-id"),
150            agent_id: get("agent-id").unwrap_or_else(|| "agent://unknown".into()),
151            agent_instance_id: get("agent-instance-id").unwrap_or_else(|| "ai_unknown".into()),
152            workspace_id: get("workspace-id"),
153            mission_id: get("mission-id"),
154            host_id: get("host-id").unwrap_or_else(default_host_id),
155            tool_runtime_id: get("tool-runtime-id"),
156        })
157    }
158
159    /// Generate a child span context: new span_id, current span_id becomes parent.
160    pub fn child_span(&self) -> Self {
161        Self {
162            session_id: self.session_id.clone(),
163            trace_id: self.trace_id.clone(),
164            span_id: generate_span_id(),
165            parent_span_id: Some(self.span_id.clone()),
166            agent_id: self.agent_id.clone(),
167            agent_instance_id: self.agent_instance_id.clone(),
168            workspace_id: self.workspace_id.clone(),
169            mission_id: self.mission_id.clone(),
170            host_id: self.host_id.clone(),
171            tool_runtime_id: self.tool_runtime_id.clone(),
172        }
173    }
174
175    /// Generate a W3C traceparent header value.
176    pub fn to_traceparent(&self) -> String {
177        // Pad trace_id to 32 chars, span_id to 16 chars
178        let tid = format!("{:0>32}", &self.trace_id);
179        let sid = format!("{:0>16}", &self.span_id);
180        format!("00-{tid}-{sid}-01")
181    }
182}
183
184/// Default host ID derived from hostname.
185///
186/// WASM builds have no OS-level hostname available, so they fall back to
187/// "host_unknown". The receipts emitted from WASM contexts (edge workers,
188/// browser verifiers, Lambda) don't author session manifests directly --
189/// they consume receipts authored elsewhere -- so the fallback is safe.
190#[cfg(not(target_family = "wasm"))]
191fn default_host_id() -> String {
192    hostname::get()
193        .ok()
194        .and_then(|h| h.into_string().ok())
195        .map(|h| format!("host_{}", h.replace('.', "_")))
196        .unwrap_or_else(|| "host_unknown".into())
197}
198
199#[cfg(target_family = "wasm")]
200fn default_host_id() -> String {
201    "host_unknown".into()
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207
208    #[test]
209    fn child_span_preserves_trace() {
210        let ctx = PropagationContext {
211            session_id: "ssn_001".into(),
212            trace_id: "abcd1234abcd1234abcd1234abcd1234".into(),
213            span_id: "1111222233334444".into(),
214            parent_span_id: None,
215            agent_id: "agent://test".into(),
216            agent_instance_id: "ai_1".into(),
217            workspace_id: None,
218            mission_id: None,
219            host_id: "host_local".into(),
220            tool_runtime_id: None,
221        };
222
223        let child = ctx.child_span();
224        assert_eq!(child.trace_id, ctx.trace_id);
225        assert_eq!(child.parent_span_id.as_deref(), Some("1111222233334444"));
226        assert_ne!(child.span_id, ctx.span_id);
227    }
228
229    #[test]
230    fn headers_roundtrip() {
231        let ctx = PropagationContext {
232            session_id: "ssn_002".into(),
233            trace_id: "abcd".into(),
234            span_id: "ef01".into(),
235            parent_span_id: Some("0000".into()),
236            agent_id: "agent://claude".into(),
237            agent_instance_id: "ai_cc_1".into(),
238            workspace_id: Some("ws_1".into()),
239            mission_id: None,
240            host_id: "host_mac".into(),
241            tool_runtime_id: Some("rt_1".into()),
242        };
243
244        let headers = ctx.to_headers();
245        let back = PropagationContext::from_headers(&headers).unwrap();
246        assert_eq!(back.session_id, "ssn_002");
247        assert_eq!(back.parent_span_id.as_deref(), Some("0000"));
248        assert_eq!(back.workspace_id.as_deref(), Some("ws_1"));
249    }
250
251    #[test]
252    fn traceparent_format() {
253        let ctx = PropagationContext {
254            session_id: "ssn_001".into(),
255            trace_id: "abcd1234abcd1234abcd1234abcd1234".into(),
256            span_id: "1111222233334444".into(),
257            parent_span_id: None,
258            agent_id: "agent://test".into(),
259            agent_instance_id: "ai_1".into(),
260            workspace_id: None,
261            mission_id: None,
262            host_id: "host_local".into(),
263            tool_runtime_id: None,
264        };
265        let tp = ctx.to_traceparent();
266        assert_eq!(
267            tp,
268            "00-abcd1234abcd1234abcd1234abcd1234-1111222233334444-01"
269        );
270    }
271}
treeship_core/session/context.rs

treeship_core/session/
context.rs
