agtrace_sdk/mcp/
server.rs

1//! MCP JSON-RPC server.
2
3use schemars::schema_for;
4use serde::{Deserialize, Serialize};
5use serde_json::{Value, json};
6use std::io::{BufRead, BufReader, Write};
7
8use crate::Client;
9use crate::query::{
10    AnalyzeSessionArgs, GetTurnsArgs, ListSessionsArgs, ListTurnsArgs, SearchEventsArgs,
11};
12
13use super::tools::{
14    handle_analyze_session, handle_get_project_info, handle_get_turns, handle_list_sessions,
15    handle_list_turns, handle_search_events,
16};
17
18#[derive(Debug, Deserialize)]
19struct JsonRpcRequest {
20    #[allow(dead_code)]
21    jsonrpc: String,
22    id: Option<Value>,
23    method: String,
24    params: Option<Value>,
25}
26
27#[derive(Debug, Serialize)]
28struct JsonRpcResponse {
29    jsonrpc: String,
30    id: Value,
31    #[serde(skip_serializing_if = "Option::is_none")]
32    result: Option<Value>,
33    #[serde(skip_serializing_if = "Option::is_none")]
34    error: Option<JsonRpcError>,
35}
36
37#[derive(Debug, Serialize)]
38struct JsonRpcError {
39    code: i32,
40    message: String,
41    #[serde(skip_serializing_if = "Option::is_none")]
42    data: Option<Value>,
43}
44
45pub struct AgTraceServer {
46    client: Client,
47}
48
49impl AgTraceServer {
50    pub fn new(client: Client) -> Self {
51        Self { client }
52    }
53
54    /// Convert serde deserialization error to MCP-compliant JSON-RPC error
55    fn parse_validation_error(tool_name: &str, error: serde_json::Error) -> JsonRpcError {
56        let error_msg = error.to_string();
57
58        // Check if it's a "missing field" error
59        if error_msg.contains("missing field")
60            && let Some(field_start) = error_msg.find('`')
61            && let Some(field_end) = error_msg[field_start + 1..].find('`')
62        {
63            let field_name = &error_msg[field_start + 1..field_start + 1 + field_end];
64            return JsonRpcError {
65                code: -32602,
66                message: format!("Invalid params: missing required field \"{}\"", field_name),
67                data: Some(json!({
68                    "missing": [field_name],
69                    "tool": tool_name,
70                })),
71            };
72        }
73
74        // Fallback for other validation errors
75        JsonRpcError {
76            code: -32602,
77            message: format!("Invalid params: {}", error),
78            data: Some(json!({
79                "tool": tool_name,
80                "detail": error_msg,
81            })),
82        }
83    }
84
85    async fn handle_request(&self, request: JsonRpcRequest) -> JsonRpcResponse {
86        // MCP requires all requests to have an id, use a default if missing
87        let id = request
88            .id
89            .clone()
90            .unwrap_or_else(|| Value::Number(serde_json::Number::from(0)));
91
92        match request.method.as_str() {
93            "initialize" => self.handle_initialize(id, request.params).await,
94            "tools/list" => self.handle_list_tools(id).await,
95            "tools/call" => self.handle_call_tool(id, request.params).await,
96            _ => JsonRpcResponse {
97                jsonrpc: "2.0".to_string(),
98                id,
99                result: None,
100                error: Some(JsonRpcError {
101                    code: -32601,
102                    message: format!("Method not found: {}", request.method),
103                    data: None,
104                }),
105            },
106        }
107    }
108
109    async fn handle_initialize(&self, id: Value, _params: Option<Value>) -> JsonRpcResponse {
110        JsonRpcResponse {
111            jsonrpc: "2.0".to_string(),
112            id,
113            result: Some(json!({
114                "protocolVersion": "2024-11-05",
115                "capabilities": {
116                    "tools": {}
117                },
118                "serverInfo": {
119                    "name": "agtrace",
120                    "version": env!("CARGO_PKG_VERSION")
121                },
122                "instructions": "AgTrace MCP Server - AI agent execution observability. Use these tools to query historical sessions, analyze failures, search event payloads, and debug agent behavior."
123            })),
124            error: None,
125        }
126    }
127
128    async fn handle_list_tools(&self, id: Value) -> JsonRpcResponse {
129        // Generate JSON Schemas from Rust types - single source of truth!
130        let list_sessions_schema = schema_for!(ListSessionsArgs);
131        let analyze_session_schema = schema_for!(AnalyzeSessionArgs);
132        let search_events_schema = schema_for!(SearchEventsArgs);
133        let list_turns_schema = schema_for!(ListTurnsArgs);
134        let get_turns_schema = schema_for!(GetTurnsArgs);
135
136        JsonRpcResponse {
137            jsonrpc: "2.0".to_string(),
138            id,
139            result: Some(json!({
140                "tools": [
141                    {
142                        "name": "list_sessions",
143                        "description": "List recent AI agent sessions with cursor-based pagination. WORKFLOW: Call this first to discover available sessions, then use session IDs with other tools. Safe to call multiple times with different filters.",
144                        "inputSchema": serde_json::to_value(&list_sessions_schema).unwrap(),
145                    },
146                    {
147                        "name": "get_project_info",
148                        "description": "List all projects that have been indexed by agtrace with their metadata. WORKFLOW: Use this to discover available projects and their hashes. Safe to call anytime.",
149                        "inputSchema": {
150                            "type": "object",
151                            "properties": {}
152                        }
153                    },
154                    {
155                        "name": "analyze_session",
156                        "description": "Run diagnostic analysis on a session to identify failures, loops, and issues. WORKFLOW: First call list_sessions to obtain session IDs, then use those IDs with this tool. Safe to call in parallel for multiple known session IDs.",
157                        "inputSchema": serde_json::to_value(&analyze_session_schema).unwrap(),
158                    },
159                    {
160                        "name": "search_events",
161                        "description": "Search for events and return navigation coordinates (session_id, event_index, turn_index, step_index). Use this to find specific events, then use turn_index with list_turns or get_turns for detailed analysis.",
162                        "inputSchema": serde_json::to_value(&search_events_schema).unwrap(),
163                    },
164                    {
165                        "name": "list_turns",
166                        "description": "List turns with metadata only (no payload content). Returns turn statistics including step_count, duration_ms, total_tokens, and tools_used. Use this to get an overview before drilling down with get_turns.",
167                        "inputSchema": serde_json::to_value(&list_turns_schema).unwrap(),
168                    },
169                    {
170                        "name": "get_turns",
171                        "description": "Get details for specific turns. Defaults are tuned for safety based on data distribution (max 30 steps/turn, 3000 chars/field). WORKFLOW: Fetch 1-2 turns at a time to avoid token limits. If data is marked '[TRUNCATED]' and critical, retry with higher limits.",
172                        "inputSchema": serde_json::to_value(&get_turns_schema).unwrap(),
173                    }
174                ]
175            })),
176            error: None,
177        }
178    }
179
180    async fn handle_call_tool(&self, id: Value, params: Option<Value>) -> JsonRpcResponse {
181        let params = match params {
182            Some(p) => p,
183            None => {
184                return JsonRpcResponse {
185                    jsonrpc: "2.0".to_string(),
186                    id,
187                    result: None,
188                    error: Some(JsonRpcError {
189                        code: -32602,
190                        message: "Missing params".to_string(),
191                        data: None,
192                    }),
193                };
194            }
195        };
196
197        let tool_name = match params.get("name").and_then(|v| v.as_str()) {
198            Some(name) => name,
199            None => {
200                return JsonRpcResponse {
201                    jsonrpc: "2.0".to_string(),
202                    id,
203                    result: None,
204                    error: Some(JsonRpcError {
205                        code: -32602,
206                        message: "Missing tool name".to_string(),
207                        data: None,
208                    }),
209                };
210            }
211        };
212
213        let arguments = params.get("arguments").cloned().unwrap_or(json!({}));
214
215        let result = match tool_name {
216            "list_sessions" => {
217                let args: ListSessionsArgs = match serde_json::from_value(arguments) {
218                    Ok(args) => args,
219                    Err(e) => {
220                        return JsonRpcResponse {
221                            jsonrpc: "2.0".to_string(),
222                            id,
223                            result: None,
224                            error: Some(Self::parse_validation_error("list_sessions", e)),
225                        };
226                    }
227                };
228                handle_list_sessions(&self.client, args).await
229            }
230            "get_project_info" => handle_get_project_info(&self.client).await,
231            "analyze_session" => {
232                let args: AnalyzeSessionArgs = match serde_json::from_value(arguments) {
233                    Ok(args) => args,
234                    Err(e) => {
235                        return JsonRpcResponse {
236                            jsonrpc: "2.0".to_string(),
237                            id,
238                            result: None,
239                            error: Some(Self::parse_validation_error("analyze_session", e)),
240                        };
241                    }
242                };
243                handle_analyze_session(&self.client, args).await
244            }
245            "search_events" => {
246                let args: SearchEventsArgs = match serde_json::from_value(arguments) {
247                    Ok(args) => args,
248                    Err(e) => {
249                        return JsonRpcResponse {
250                            jsonrpc: "2.0".to_string(),
251                            id,
252                            result: None,
253                            error: Some(Self::parse_validation_error("search_events", e)),
254                        };
255                    }
256                };
257                handle_search_events(&self.client, args).await
258            }
259            "list_turns" => {
260                let args: ListTurnsArgs = match serde_json::from_value(arguments) {
261                    Ok(args) => args,
262                    Err(e) => {
263                        return JsonRpcResponse {
264                            jsonrpc: "2.0".to_string(),
265                            id,
266                            result: None,
267                            error: Some(Self::parse_validation_error("list_turns", e)),
268                        };
269                    }
270                };
271                handle_list_turns(&self.client, args).await
272            }
273            "get_turns" => {
274                let args: GetTurnsArgs = match serde_json::from_value(arguments) {
275                    Ok(args) => args,
276                    Err(e) => {
277                        return JsonRpcResponse {
278                            jsonrpc: "2.0".to_string(),
279                            id,
280                            result: None,
281                            error: Some(Self::parse_validation_error("get_turns", e)),
282                        };
283                    }
284                };
285                handle_get_turns(&self.client, args).await
286            }
287            _ => Err(format!("Unknown tool: {}", tool_name)),
288        };
289
290        match result {
291            Ok(content) => JsonRpcResponse {
292                jsonrpc: "2.0".to_string(),
293                id,
294                result: Some(json!({
295                    "content": [
296                        {
297                            "type": "text",
298                            "text": serde_json::to_string(&content).unwrap_or_else(|_| content.to_string())
299                        }
300                    ]
301                })),
302                error: None,
303            },
304            Err(e) => JsonRpcResponse {
305                jsonrpc: "2.0".to_string(),
306                id,
307                result: None,
308                error: Some(JsonRpcError {
309                    code: -32603,
310                    message: e,
311                    data: None,
312                }),
313            },
314        }
315    }
316}
317
318/// Run the MCP server over stdio.
319pub async fn run_server(client: Client) -> anyhow::Result<()> {
320    let server = AgTraceServer::new(client);
321    let stdin = std::io::stdin();
322    let mut stdout = std::io::stdout();
323    let reader = BufReader::new(stdin);
324
325    for line in reader.lines() {
326        let line = line?;
327        let trimmed = line.trim();
328
329        if trimmed.is_empty() {
330            continue;
331        }
332
333        let request: JsonRpcRequest = match serde_json::from_str(trimmed) {
334            Ok(req) => req,
335            Err(e) => {
336                // For parse errors, we can't get a valid id, so we use a sentinel value
337                let error_response = JsonRpcResponse {
338                    jsonrpc: "2.0".to_string(),
339                    id: Value::Number(serde_json::Number::from(-1)),
340                    result: None,
341                    error: Some(JsonRpcError {
342                        code: -32700,
343                        message: format!("Parse error: {}", e),
344                        data: None,
345                    }),
346                };
347                let response_json = serde_json::to_string(&error_response)?;
348                writeln!(stdout, "{}", response_json)?;
349                stdout.flush()?;
350                continue;
351            }
352        };
353
354        let response = server.handle_request(request).await;
355        let response_json = serde_json::to_string(&response)?;
356        writeln!(stdout, "{}", response_json)?;
357        stdout.flush()?;
358    }
359
360    Ok(())
361}
362
363#[cfg(test)]
364mod tests {}
agtrace_sdk/mcp/server.rs

agtrace_sdk/mcp/
server.rs
