systemprompt_logging/services/spans/
mod.rs1use 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}