systemprompt_runtime/trace/repository/
mcp_trace.rs

1//! MCP tool-call trace queries.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use systemprompt_identifiers::{
7    AiRequestId, ArtifactId, ContextId, McpExecutionId, McpServerId, McpToolName, TaskId,
8};
9
10use super::{Result, TraceRepository};
11use crate::trace::models::{AiRequestInfo, McpToolExecution, TaskArtifact, ToolLogEntry};
12
13impl TraceRepository {
14    pub async fn fetch_mcp_executions(
15        &self,
16        task_id: &TaskId,
17        context_id: &ContextId,
18    ) -> Result<Vec<McpToolExecution>> {
19        let rows = sqlx::query!(
20            r#"SELECT mcp_execution_id, tool_name, server_name, status, execution_time_ms,
21                  error_message, input, output
22           FROM mcp_tool_executions
23           WHERE task_id = $1 OR context_id = $2
24           ORDER BY started_at"#,
25            task_id.as_str(),
26            context_id.as_str()
27        )
28        .fetch_all(&*self.pool)
29        .await?;
30
31        Ok(rows
32            .into_iter()
33            .map(|r| McpToolExecution {
34                mcp_execution_id: McpExecutionId::new(r.mcp_execution_id),
35                tool_name: McpToolName::new(r.tool_name),
36                server_name: McpServerId::new(r.server_name),
37                status: r.status,
38                execution_time_ms: r.execution_time_ms,
39                error_message: r.error_message,
40                input: r.input,
41                output: r.output,
42            })
43            .collect())
44    }
45
46    pub async fn fetch_mcp_linked_ai_requests(
47        &self,
48        mcp_execution_id: &McpExecutionId,
49    ) -> Result<Vec<AiRequestInfo>> {
50        let rows = sqlx::query!(
51        r#"SELECT id, model, provider, max_tokens, input_tokens, output_tokens, cost_microdollars, latency_ms
52           FROM ai_requests
53           WHERE mcp_execution_id = $1
54           ORDER BY created_at"#,
55        mcp_execution_id.as_str()
56    )
57    .fetch_all(&*self.pool)
58    .await?;
59
60        Ok(rows
61            .into_iter()
62            .map(|r| AiRequestInfo {
63                id: AiRequestId::new(r.id),
64                provider: r.provider,
65                model: r.model,
66                max_tokens: r.max_tokens,
67                input_tokens: r.input_tokens,
68                output_tokens: r.output_tokens,
69                cost_microdollars: r.cost_microdollars,
70                latency_ms: r.latency_ms,
71            })
72            .collect())
73    }
74
75    pub async fn fetch_tool_logs(
76        &self,
77        task_id: &TaskId,
78        context_id: &ContextId,
79    ) -> Result<Vec<ToolLogEntry>> {
80        let rows = sqlx::query!(
81        r#"SELECT timestamp, level, module, message
82           FROM logs
83           WHERE (task_id = $1 OR context_id = $2)
84             AND (
85                 (module LIKE '%_tools' OR module LIKE '%_manager' OR module LIKE 'create_%' OR module LIKE 'update_%' OR module LIKE 'research_%')
86                 OR (level = 'ERROR' AND message LIKE '%tool%')
87                 OR message LIKE 'Tool executed%'
88                 OR message LIKE 'Tool failed%'
89                 OR message LIKE 'MCP execution%'
90             )
91           ORDER BY timestamp"#,
92        task_id.as_str(),
93        context_id.as_str()
94    )
95    .fetch_all(&*self.pool)
96    .await?;
97
98        rows.into_iter()
99            .map(|r| -> Result<ToolLogEntry> {
100                Ok(ToolLogEntry {
101                    timestamp: r.timestamp,
102                    level: r.level.parse()?,
103                    module: r.module,
104                    message: r.message,
105                })
106            })
107            .collect()
108    }
109
110    pub async fn fetch_task_artifacts(
111        &self,
112        task_id: &TaskId,
113        context_id: &ContextId,
114    ) -> Result<Vec<TaskArtifact>> {
115        let rows = sqlx::query!(
116        r#"SELECT ta.artifact_id, ta.artifact_type, ta.name, ta.source, ta.tool_name,
117                  ap.part_kind as "part_kind?", ap.text_content as "text_content?",
118                  ap.data_content as "data_content?"
119           FROM task_artifacts ta
120           LEFT JOIN artifact_parts ap ON ta.artifact_id = ap.artifact_id AND ta.context_id = ap.context_id
121           WHERE ta.task_id = $1 OR ta.context_id = $2
122           ORDER BY ta.created_at, ap.sequence_number"#,
123        task_id.as_str(),
124        context_id.as_str()
125    )
126    .fetch_all(&*self.pool)
127    .await?;
128
129        Ok(rows
130            .into_iter()
131            .map(|r| TaskArtifact {
132                artifact_id: ArtifactId::new(r.artifact_id),
133                artifact_type: r.artifact_type,
134                name: r.name,
135                source: r.source,
136                tool_name: r.tool_name.map(McpToolName::new),
137                part_kind: r.part_kind,
138                text_content: r.text_content,
139                data_content: r.data_content,
140            })
141            .collect())
142    }
143}
systemprompt_runtime/trace/repository/mcp_trace.rs

systemprompt_runtime/trace/repository/
mcp_trace.rs
