systemprompt_logging/services/spans/
mod.rs

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

systemprompt_logging/services/spans/
mod.rs
