systemprompt_logging/layer/
proxy.rs

1//! Deferred database-logging tracing layer.
2//!
3//! [`ProxyDatabaseLayer`] is installed in the subscriber stack before a
4//! database pool exists, buffering span attribution into span extensions. Once
5//! [`ProxyDatabaseLayer::attach`] supplies a pool it delegates to the real
6//! `DatabaseLayer` and hands the caller the writer's [`LogWriterHandle`] (only
7//! on the first attach); until then span fields are recorded so attribution is
8//! not lost across the boot window. Clones share one state, so the copy inside
9//! the installed subscriber and the copy a [`crate::LogSink`] holds see the
10//! same sink, instance label, drop count and parked writer. The free functions
11//! build the [`LogEntry`] actor triple by walking the span tree.
12//!
13//! Copyright (c) systemprompt.io — Business Source License 1.1.
14//! See <https://systemprompt.io> for licensing details.
15
16use std::sync::{Arc, Mutex, OnceLock, PoisonError};
17
18use chrono::Utc;
19use tracing::{Event, Subscriber};
20use tracing_subscriber::Layer;
21use tracing_subscriber::layer::Context;
22use tracing_subscriber::registry::LookupSpan;
23
24use super::visitor::{FieldVisitor, SpanContext, SpanFields, SpanVisitor, extract_span_context};
25use super::{DatabaseLayer, LayerShared, LogWriterHandle, LogWriterShutdownError};
26use crate::models::{LogEntry, LogLevel};
27use systemprompt_database::DbPool;
28use systemprompt_identifiers::{
29    ClientId, ContextId, InstanceId, LogId, SessionId, TaskId, TraceId, UserId,
30};
31
32#[derive(Default)]
33struct ProxyState {
34    db: OnceLock<DatabaseLayer>,
35    shared: LayerShared,
36    writer: Mutex<Option<LogWriterHandle>>,
37}
38
39#[derive(Clone, Default)]
40pub struct ProxyDatabaseLayer {
41    inner: Arc<ProxyState>,
42}
43
44impl std::fmt::Debug for ProxyDatabaseLayer {
45    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
46        f.debug_struct("ProxyDatabaseLayer")
47            .field("attached", &self.inner.db.get().is_some())
48            .finish()
49    }
50}
51
52impl ProxyDatabaseLayer {
53    pub fn new() -> Self {
54        Self::default()
55    }
56
57    pub fn attach(&self, db_pool: &DbPool) -> Option<LogWriterHandle> {
58        let mut writer = None;
59        self.inner.db.get_or_init(|| {
60            let (layer, handle) = DatabaseLayer::with_shared(db_pool, self.inner.shared.clone());
61            writer = Some(handle);
62            layer
63        });
64        writer
65    }
66
67    pub(crate) fn attach_owned(&self, db_pool: &DbPool, instance_id: Option<InstanceId>) {
68        if let Some(instance_id) = instance_id {
69            self.inner.shared.instance.set(instance_id);
70        }
71        if let Some(writer) = self.attach(db_pool) {
72            *self
73                .inner
74                .writer
75                .lock()
76                .unwrap_or_else(PoisonError::into_inner) = Some(writer);
77        }
78    }
79
80    pub(crate) fn shared(&self) -> &LayerShared {
81        &self.inner.shared
82    }
83
84    pub(crate) fn enqueue(&self, entry: LogEntry) {
85        match self.inner.db.get() {
86            Some(db) => db.enqueue(entry),
87            None => self.inner.shared.record_drop(),
88        }
89    }
90
91    pub(crate) async fn shutdown(&self) -> Result<(), LogWriterShutdownError> {
92        let writer = self
93            .inner
94            .writer
95            .lock()
96            .unwrap_or_else(PoisonError::into_inner)
97            .take();
98        match writer {
99            Some(writer) => writer.shutdown().await,
100            None => Ok(()),
101        }
102    }
103}
104
105impl<S> Layer<S> for ProxyDatabaseLayer
106where
107    S: Subscriber + for<'a> LookupSpan<'a>,
108{
109    fn on_new_span(
110        &self,
111        attrs: &tracing::span::Attributes<'_>,
112        id: &tracing::span::Id,
113        ctx: Context<'_, S>,
114    ) {
115        if let Some(db) = self.inner.db.get() {
116            db.on_new_span(attrs, id, ctx);
117        } else {
118            record_span_fields(attrs, id, &ctx);
119        }
120    }
121
122    fn on_record(
123        &self,
124        id: &tracing::span::Id,
125        values: &tracing::span::Record<'_>,
126        ctx: Context<'_, S>,
127    ) {
128        if let Some(db) = self.inner.db.get() {
129            db.on_record(id, values, ctx);
130        } else {
131            update_span_fields(id, values, &ctx);
132        }
133    }
134
135    fn on_event(&self, event: &Event<'_>, ctx: Context<'_, S>) {
136        if let Some(db) = self.inner.db.get() {
137            db.on_event(event, ctx);
138        }
139    }
140}
141
142pub(super) fn record_span_fields<S>(
143    attrs: &tracing::span::Attributes<'_>,
144    id: &tracing::span::Id,
145    ctx: &Context<'_, S>,
146) where
147    S: Subscriber + for<'a> LookupSpan<'a>,
148{
149    let Some(span) = ctx.span(id) else {
150        return;
151    };
152    let mut fields = SpanFields::default();
153    let mut context = SpanContext::default();
154    let mut visitor = SpanVisitor {
155        context: &mut context,
156    };
157    attrs.record(&mut visitor);
158
159    fields.user = context.user;
160    fields.session = context.session;
161    fields.task = context.task;
162    fields.trace = context.trace;
163    fields.context = context.context;
164    fields.client = context.client;
165
166    let mut extensions = span.extensions_mut();
167    extensions.insert(fields);
168}
169
170pub(super) fn update_span_fields<S>(
171    id: &tracing::span::Id,
172    values: &tracing::span::Record<'_>,
173    ctx: &Context<'_, S>,
174) where
175    S: Subscriber + for<'a> LookupSpan<'a>,
176{
177    if let Some(span) = ctx.span(id) {
178        let mut extensions = span.extensions_mut();
179        if let Some(fields) = extensions.get_mut::<SpanFields>() {
180            let mut context = SpanContext {
181                user: fields.user.clone(),
182                session: fields.session.clone(),
183                task: fields.task.clone(),
184                trace: fields.trace.clone(),
185                context: fields.context.clone(),
186                client: fields.client.clone(),
187            };
188            let mut visitor = SpanVisitor {
189                context: &mut context,
190            };
191            values.record(&mut visitor);
192
193            fields.user = context.user;
194            fields.session = context.session;
195            fields.task = context.task;
196            fields.trace = context.trace;
197            fields.context = context.context;
198            fields.client = context.client;
199        }
200    }
201}
202
203pub(super) fn build_log_entry<S>(event: &Event<'_>, ctx: &Context<'_, S>) -> Option<LogEntry>
204where
205    S: Subscriber + for<'a> LookupSpan<'a>,
206{
207    let level = *event.metadata().level();
208    let module = event.metadata().target().to_owned();
209
210    let mut visitor = FieldVisitor::default();
211    event.record(&mut visitor);
212
213    let span_context = ctx
214        .current_span()
215        .id()
216        .and_then(|id| ctx.span(id))
217        .map(extract_span_context)?;
218
219    let log_level = match level {
220        tracing::Level::ERROR => LogLevel::Error,
221        tracing::Level::WARN => LogLevel::Warn,
222        tracing::Level::INFO => LogLevel::Info,
223        tracing::Level::DEBUG => LogLevel::Debug,
224        tracing::Level::TRACE => LogLevel::Trace,
225    };
226
227    let user_id = UserId::new(span_context.user.as_ref()?.clone());
228    let session_id = SessionId::new(span_context.session.as_ref()?.clone());
229    let trace_id = TraceId::new(span_context.trace.as_ref()?.clone());
230
231    Some(LogEntry {
232        id: LogId::generate(),
233        timestamp: Utc::now(),
234        level: log_level,
235        module,
236        message: visitor.message,
237        metadata: visitor.fields,
238        user_id,
239        session_id,
240        task_id: span_context.task.as_ref().map(|s| TaskId::new(s.clone())),
241        trace_id,
242        context_id: span_context.context.as_ref().and_then(|s| {
243            ContextId::try_new(s.clone())
244                .map_err(|e| {
245                    tracing::warn!(error = %e, raw = %s, "Skipping non-UUID context_id from span context");
246                    e
247                })
248                .ok()
249        }),
250        client_id: span_context
251            .client
252            .as_ref()
253            .map(|s| ClientId::new(s.clone())),
254        instance_id: None,
255    })
256}
systemprompt_logging/layer/proxy.rs

systemprompt_logging/layer/
proxy.rs
