Skip to main content

systemprompt_logging/trace/
service.rs

1//! Read-side query facade over the tracing and audit tables.
2//!
3//! [`TraceQueryService`] is the single entry point for reconstructing a trace
4//! from its constituent rows (logs, AI requests, MCP tool executions, task
5//! execution steps) and for the log/audit browsing surfaces. Each public method
6//! delegates to a focused query module in this directory;
7//! [`get_all_trace_data`] fans the per-source fetches out concurrently.
8//!
9//! [`get_all_trace_data`]: TraceQueryService::get_all_trace_data
10//!
11//! Copyright (c) systemprompt.io — Business Source License 1.1.
12//! See <https://systemprompt.io> for licensing details.
13
14use chrono::{DateTime, Utc};
15use sqlx::PgPool;
16use std::sync::Arc;
17use systemprompt_identifiers::{AiRequestId, TaskId, TraceId};
18
19use crate::models::{LogEntry, LoggingError};
20
21pub(super) type Result<T> = std::result::Result<T, LoggingError>;
22
23use super::models::{
24    AiRequestDetail, AiRequestListItem, AiRequestStats, AiRequestSummary, AuditLookupResult,
25    AuditToolCallRow, ConversationMessage, ExecutionStepSummary, LevelCount, LinkedMcpCall,
26    LogSearchItem, LogTimeRange, McpExecutionSummary, ModuleCount, ToolExecutionFilter,
27    ToolExecutionItem, TraceEvent, TraceListFilter, TraceListItem,
28};
29use super::{
30    audit_queries, list_queries, log_lookup_queries, log_search_queries, log_summary_queries,
31    queries, request_queries, tool_queries,
32};
33
34#[derive(Debug, Clone)]
35pub struct TraceQueryService {
36    pool: Arc<PgPool>,
37}
38
39impl TraceQueryService {
40    pub const fn new(pool: Arc<PgPool>) -> Self {
41        Self { pool }
42    }
43
44    pub async fn get_log_events(&self, trace_id: &TraceId) -> Result<Vec<TraceEvent>> {
45        queries::fetch_log_events(&self.pool, trace_id).await
46    }
47
48    pub async fn get_ai_request_summary(&self, trace_id: &TraceId) -> Result<AiRequestSummary> {
49        queries::fetch_ai_request_summary(&self.pool, trace_id).await
50    }
51
52    pub async fn get_ai_request_events(&self, trace_id: &TraceId) -> Result<Vec<TraceEvent>> {
53        queries::fetch_ai_request_events(&self.pool, trace_id).await
54    }
55
56    pub async fn get_mcp_execution_summary(
57        &self,
58        trace_id: &TraceId,
59    ) -> Result<McpExecutionSummary> {
60        queries::fetch_mcp_execution_summary(&self.pool, trace_id).await
61    }
62
63    pub async fn get_mcp_execution_events(&self, trace_id: &TraceId) -> Result<Vec<TraceEvent>> {
64        queries::fetch_mcp_execution_events(&self.pool, trace_id).await
65    }
66
67    pub async fn get_task_id(&self, trace_id: &TraceId) -> Result<Option<TaskId>> {
68        Ok(queries::fetch_task_id_for_trace(&self.pool, trace_id)
69            .await?
70            .map(TaskId::new))
71    }
72
73    pub async fn get_execution_step_summary(
74        &self,
75        trace_id: &TraceId,
76    ) -> Result<ExecutionStepSummary> {
77        queries::fetch_execution_step_summary(&self.pool, trace_id).await
78    }
79
80    pub async fn get_execution_step_events(&self, trace_id: &TraceId) -> Result<Vec<TraceEvent>> {
81        queries::fetch_execution_step_events(&self.pool, trace_id).await
82    }
83
84    pub async fn get_all_trace_data(
85        &self,
86        trace_id: &TraceId,
87    ) -> Result<(
88        Vec<TraceEvent>,
89        Vec<TraceEvent>,
90        Vec<TraceEvent>,
91        Vec<TraceEvent>,
92        AiRequestSummary,
93        McpExecutionSummary,
94        ExecutionStepSummary,
95        Option<TaskId>,
96    )> {
97        tokio::try_join!(
98            self.get_log_events(trace_id),
99            self.get_ai_request_events(trace_id),
100            self.get_mcp_execution_events(trace_id),
101            self.get_execution_step_events(trace_id),
102            self.get_ai_request_summary(trace_id),
103            self.get_mcp_execution_summary(trace_id),
104            self.get_execution_step_summary(trace_id),
105            self.get_task_id(trace_id),
106        )
107    }
108
109    pub async fn list_traces(&self, filter: &TraceListFilter) -> Result<Vec<TraceListItem>> {
110        list_queries::list_traces(&self.pool, filter).await
111    }
112
113    pub async fn list_tool_executions(
114        &self,
115        filter: &ToolExecutionFilter,
116    ) -> Result<Vec<ToolExecutionItem>> {
117        tool_queries::list_tool_executions(&self.pool, filter).await
118    }
119
120    pub async fn search_logs(
121        &self,
122        pattern: &str,
123        since: Option<DateTime<Utc>>,
124        level: Option<&str>,
125        limit: i64,
126    ) -> Result<Vec<LogSearchItem>> {
127        log_search_queries::search_logs(&self.pool, pattern, since, level, limit).await
128    }
129
130    pub async fn search_tool_executions(
131        &self,
132        pattern: &str,
133        since: Option<DateTime<Utc>>,
134        limit: i64,
135    ) -> Result<Vec<ToolExecutionItem>> {
136        log_search_queries::search_tool_executions(&self.pool, pattern, since, limit).await
137    }
138
139    pub async fn list_ai_requests(
140        &self,
141        since: Option<DateTime<Utc>>,
142        model: Option<&str>,
143        provider: Option<&str>,
144        limit: i64,
145    ) -> Result<Vec<AiRequestListItem>> {
146        request_queries::list_ai_requests(&self.pool, since, model, provider, limit).await
147    }
148
149    pub async fn get_ai_request_stats(
150        &self,
151        since: Option<DateTime<Utc>>,
152    ) -> Result<AiRequestStats> {
153        request_queries::get_ai_request_stats(&self.pool, since).await
154    }
155
156    pub async fn find_ai_request_detail(&self, id: &str) -> Result<Option<AiRequestDetail>> {
157        request_queries::find_ai_request_detail(&self.pool, id).await
158    }
159
160    pub async fn find_ai_request_for_audit(&self, id: &str) -> Result<Option<AuditLookupResult>> {
161        audit_queries::find_ai_request_for_audit(&self.pool, id).await
162    }
163
164    pub async fn list_audit_messages(
165        &self,
166        request_id: &AiRequestId,
167    ) -> Result<Vec<ConversationMessage>> {
168        audit_queries::list_audit_messages(&self.pool, request_id).await
169    }
170
171    pub async fn list_audit_tool_calls(
172        &self,
173        request_id: &AiRequestId,
174    ) -> Result<Vec<AuditToolCallRow>> {
175        audit_queries::list_audit_tool_calls(&self.pool, request_id).await
176    }
177
178    pub async fn list_linked_mcp_calls(
179        &self,
180        request_id: &AiRequestId,
181    ) -> Result<Vec<LinkedMcpCall>> {
182        audit_queries::list_linked_mcp_calls(&self.pool, request_id).await
183    }
184
185    pub async fn find_log_by_id(&self, id: &str) -> Result<Option<LogEntry>> {
186        log_lookup_queries::find_log_by_id(&self.pool, id).await
187    }
188
189    pub async fn find_log_by_partial_id(&self, id_prefix: &str) -> Result<Option<LogEntry>> {
190        log_lookup_queries::find_log_by_partial_id(&self.pool, id_prefix).await
191    }
192
193    pub async fn find_logs_by_trace_id(&self, trace_id: &TraceId) -> Result<Vec<LogEntry>> {
194        log_lookup_queries::find_logs_by_trace_id(&self.pool, trace_id).await
195    }
196
197    pub async fn list_logs_filtered(
198        &self,
199        since: Option<DateTime<Utc>>,
200        level: Option<&str>,
201        limit: i64,
202    ) -> Result<Vec<LogEntry>> {
203        log_lookup_queries::list_logs_filtered(&self.pool, since, level, limit).await
204    }
205
206    pub async fn count_logs_by_level(
207        &self,
208        since: Option<DateTime<Utc>>,
209    ) -> Result<Vec<LevelCount>> {
210        log_summary_queries::count_logs_by_level(&self.pool, since).await
211    }
212
213    pub async fn top_modules(
214        &self,
215        since: Option<DateTime<Utc>>,
216        limit: i64,
217    ) -> Result<Vec<ModuleCount>> {
218        log_summary_queries::top_modules(&self.pool, since, limit).await
219    }
220
221    pub async fn log_time_range(&self, since: Option<DateTime<Utc>>) -> Result<LogTimeRange> {
222        log_summary_queries::log_time_range(&self.pool, since).await
223    }
224
225    pub async fn total_log_count(&self) -> Result<i64> {
226        log_summary_queries::total_log_count(&self.pool).await
227    }
228}