Skip to main content

systemprompt_logging/services/spans/
mod.rs

1//! Typed tracing-span constructors carrying request attribution.
2//!
3//! [`RequestSpan`] and [`SystemSpan`] wrap a `tracing::Span` seeded with the
4//! identifier fields (`user_id`, `session_id`, `trace_id`, and optional
5//! `context_id`/`task_id`/`client_id`) that the database log layer extracts to
6//! attribute each emitted log row. [`RequestSpanBuilder`] assembles a span with
7//! the optional fields populated.
8//!
9//! Copyright (c) systemprompt.io — Business Source License 1.1.
10//! See <https://systemprompt.io> for licensing details.
11
12use systemprompt_identifiers::{ClientId, ContextId, SessionId, TaskId, TraceId, UserId};
13use tracing::Span;
14
15pub struct RequestSpan(Span);
16
17impl std::fmt::Debug for RequestSpan {
18    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19        f.debug_tuple("RequestSpan").finish()
20    }
21}
22
23impl RequestSpan {
24    pub fn new(user_id: &UserId, session_id: &SessionId, trace_id: &TraceId) -> Self {
25        let span = tracing::info_span!(
26            "request",
27            user_id = %user_id.as_str(),
28            session_id = %session_id.as_str(),
29            trace_id = %trace_id.as_str(),
30            context_id = tracing::field::Empty,
31            task_id = tracing::field::Empty,
32            client_id = tracing::field::Empty,
33        );
34
35        Self(span)
36    }
37
38    pub fn enter(&self) -> tracing::span::EnteredSpan {
39        self.0.clone().entered()
40    }
41
42    pub fn record_task_id(&self, task_id: &TaskId) {
43        self.0.record("task_id", task_id.as_str());
44    }
45
46    pub fn record_context_id(&self, context_id: &ContextId) {
47        self.0.record("context_id", context_id.as_str());
48    }
49
50    pub fn record_client_id(&self, client_id: &ClientId) {
51        self.0.record("client_id", client_id.as_str());
52    }
53
54    pub const fn span(&self) -> &Span {
55        &self.0
56    }
57}
58
59pub struct SystemSpan(Span);
60
61impl std::fmt::Debug for SystemSpan {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        f.debug_tuple("SystemSpan").finish()
64    }
65}
66
67impl SystemSpan {
68    pub fn new(component: &str) -> Self {
69        Self(tracing::info_span!(
70            "system",
71            user_id = "system",
72            session_id = "system",
73            trace_id = %TraceId::generate().as_str(),
74            client_id = %format!("system:{component}"),
75            context_id = tracing::field::Empty,
76            task_id = tracing::field::Empty,
77        ))
78    }
79
80    pub fn enter(&self) -> tracing::span::EnteredSpan {
81        self.0.clone().entered()
82    }
83
84    pub fn record_task_id(&self, task_id: &TaskId) {
85        self.0.record("task_id", task_id.as_str());
86    }
87
88    pub fn record_context_id(&self, context_id: &ContextId) {
89        self.0.record("context_id", context_id.as_str());
90    }
91
92    pub const fn span(&self) -> &Span {
93        &self.0
94    }
95
96    pub fn into_span(self) -> Span {
97        self.0
98    }
99}
100
101impl From<SystemSpan> for Span {
102    fn from(system_span: SystemSpan) -> Self {
103        system_span.0
104    }
105}
106
107pub struct RequestSpanBuilder<'a> {
108    user: &'a UserId,
109    session: &'a SessionId,
110    trace: &'a TraceId,
111    context: Option<&'a ContextId>,
112    task: Option<&'a TaskId>,
113    client: Option<&'a ClientId>,
114}
115
116impl std::fmt::Debug for RequestSpanBuilder<'_> {
117    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
118        f.debug_struct("RequestSpanBuilder").finish_non_exhaustive()
119    }
120}
121
122impl<'a> RequestSpanBuilder<'a> {
123    pub const fn new(
124        user_id: &'a UserId,
125        session_id: &'a SessionId,
126        trace_id: &'a TraceId,
127    ) -> Self {
128        Self {
129            user: user_id,
130            session: session_id,
131            trace: trace_id,
132            context: None,
133            task: None,
134            client: None,
135        }
136    }
137
138    #[must_use]
139    pub fn with_context_id(mut self, context_id: &'a ContextId) -> Self {
140        if !context_id.as_str().is_empty() {
141            self.context = Some(context_id);
142        }
143        self
144    }
145
146    #[must_use]
147    pub const fn with_task_id(mut self, task_id: &'a TaskId) -> Self {
148        self.task = Some(task_id);
149        self
150    }
151
152    #[must_use]
153    pub const fn with_client_id(mut self, client_id: &'a ClientId) -> Self {
154        self.client = Some(client_id);
155        self
156    }
157
158    pub fn build(self) -> RequestSpan {
159        let span = RequestSpan::new(self.user, self.session, self.trace);
160
161        if let Some(context_id) = self.context {
162            span.record_context_id(context_id);
163        }
164        if let Some(task_id) = self.task {
165            span.record_task_id(task_id);
166        }
167        if let Some(client_id) = self.client {
168            span.record_client_id(client_id);
169        }
170
171        span
172    }
173}