mcp_rtk/
tracking.rs

1//! SQLite-backed token savings metrics.
2//!
3//! The [`Tracker`] records every tool call's raw and filtered byte sizes in a
4//! local SQLite database. The `mcp-rtk gain` subcommand reads these metrics to
5//! display a colorful summary with per-tool breakdowns and an efficiency meter.
6//!
7//! # Database Schema
8//!
9//! ```sql
10//! CREATE TABLE tool_calls (
11//!     id INTEGER PRIMARY KEY AUTOINCREMENT,
12//!     timestamp TEXT DEFAULT (datetime('now')),
13//!     tool_name TEXT NOT NULL,
14//!     input_bytes INTEGER NOT NULL,
15//!     output_bytes INTEGER NOT NULL,
16//!     saved_bytes INTEGER NOT NULL,
17//!     savings_pct REAL NOT NULL
18//! );
19//! ```
20
21use anyhow::{Context, Result};
22use rusqlite::Connection;
23use std::path::PathBuf;
24use std::sync::Mutex;
25use std::time::Duration;
26
27use crate::display::*;
28
29/// SQLite-backed tracker for recording and displaying token savings metrics.
30///
31/// Thread-safe via an internal `Mutex<Connection>`, satisfying the `Sync`
32/// requirement of [`ServerHandler`](rmcp::handler::server::ServerHandler).
33///
34/// # Examples
35///
36/// ```no_run
37/// # use mcp_rtk::tracking::Tracker;
38/// let tracker = Tracker::new("~/.local/share/mcp-rtk/metrics.db").unwrap();
39/// tracker.track("list_merge_requests", "{...raw...}", "{...filtered...}", "gitlab").unwrap();
40/// tracker.print_stats().unwrap();
41/// ```
42pub struct Tracker {
43    conn: Mutex<Connection>,
44}
45
46impl Tracker {
47    /// Open or create the tracking database at the given path.
48    ///
49    /// Supports `~/` expansion. Creates parent directories if needed.
50    ///
51    /// # Errors
52    ///
53    /// Returns an error if the database directory cannot be created or the
54    /// SQLite connection fails to open.
55    pub fn new(db_path: &str) -> Result<Self> {
56        let expanded = expand_path(db_path);
57        if let Some(parent) = expanded.parent() {
58            std::fs::create_dir_all(parent)
59                .context("Failed to create tracking database directory")?;
60        }
61        let conn = Connection::open(&expanded).context("Failed to open tracking database")?;
62        conn.busy_timeout(Duration::from_secs(5))?;
63        conn.execute_batch("PRAGMA journal_mode=WAL;")?;
64        conn.execute_batch(
65            "CREATE TABLE IF NOT EXISTS tool_calls (
66                id INTEGER PRIMARY KEY AUTOINCREMENT,
67                timestamp TEXT DEFAULT (datetime('now')),
68                tool_name TEXT NOT NULL,
69                input_bytes INTEGER NOT NULL,
70                output_bytes INTEGER NOT NULL,
71                saved_bytes INTEGER NOT NULL,
72                savings_pct REAL NOT NULL
73            );
74            CREATE INDEX IF NOT EXISTS idx_tool_calls_timestamp ON tool_calls(timestamp);
75            CREATE INDEX IF NOT EXISTS idx_tool_calls_tool ON tool_calls(tool_name);",
76        )
77        .context("Failed to initialize tracking tables")?;
78
79        // Migration: add preset column if missing
80        let has_preset: bool = conn
81            .prepare("SELECT preset FROM tool_calls LIMIT 0")
82            .is_ok();
83        if !has_preset {
84            conn.execute_batch(
85                "ALTER TABLE tool_calls ADD COLUMN preset TEXT NOT NULL DEFAULT 'unknown';",
86            )
87            .context("Failed to add preset column")?;
88        }
89
90        Ok(Self {
91            conn: Mutex::new(conn),
92        })
93    }
94
95    /// Record a single tool call's raw and filtered output sizes.
96    ///
97    /// Token count is estimated as `bytes / 4`.
98    ///
99    /// # Errors
100    ///
101    /// Returns an error if the database lock is poisoned or the insert fails.
102    pub fn track(
103        &self,
104        tool_name: &str,
105        raw_output: &str,
106        filtered_output: &str,
107        preset: &str,
108    ) -> Result<()> {
109        let input_bytes = raw_output.len() as i64;
110        let output_bytes = filtered_output.len() as i64;
111        // Clamp to zero: filtered output can rarely exceed raw when
112        // JSON re-serialization or custom transforms add characters.
113        let saved_bytes = (input_bytes - output_bytes).max(0);
114        let savings_pct = if input_bytes > 0 {
115            (saved_bytes as f64 / input_bytes as f64) * 100.0
116        } else {
117            0.0
118        };
119
120        let conn = self
121            .conn
122            .lock()
123            .map_err(|e| anyhow::anyhow!("lock poisoned: {e}"))?;
124        conn.execute(
125            "INSERT INTO tool_calls (tool_name, input_bytes, output_bytes, saved_bytes, savings_pct, preset)
126             VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
127            rusqlite::params![tool_name, input_bytes, output_bytes, saved_bytes, savings_pct, preset],
128        )?;
129
130        Ok(())
131    }
132
133    /// Print a colorful summary of all-time token savings to stdout.
134    ///
135    /// Includes an efficiency meter bar and a per-tool breakdown table with
136    /// impact bars.
137    ///
138    /// # Errors
139    ///
140    /// Returns an error if the database lock is poisoned or query fails.
141    pub fn print_stats(&self) -> Result<()> {
142        let conn = self
143            .conn
144            .lock()
145            .map_err(|e| anyhow::anyhow!("lock poisoned: {e}"))?;
146
147        // Fetch per-tool stats grouped by preset
148        let mut stmt = conn.prepare(
149            "SELECT
150                preset,
151                tool_name,
152                COUNT(*) as calls,
153                SUM(input_bytes) as total_input,
154                SUM(output_bytes) as total_output,
155                SUM(saved_bytes) as total_saved,
156                AVG(savings_pct) as avg_pct
157             FROM tool_calls
158             GROUP BY preset, tool_name
159             ORDER BY preset, total_saved DESC",
160        )?;
161
162        struct ToolRow {
163            preset: String,
164            name: String,
165            calls: i64,
166            saved: i64,
167            avg_pct: f64,
168        }
169
170        let rows: Vec<ToolRow> = stmt
171            .query_map([], |row| {
172                Ok(ToolRow {
173                    preset: row.get(0)?,
174                    name: row.get(1)?,
175                    calls: row.get(2)?,
176                    saved: row.get(5)?,
177                    avg_pct: row.get(6)?,
178                })
179            })?
180            .filter_map(|r| r.ok())
181            .collect();
182
183        let grand_calls: i64 = rows.iter().map(|r| r.calls).sum();
184        let grand_input: i64 = conn.query_row(
185            "SELECT COALESCE(SUM(input_bytes), 0) FROM tool_calls",
186            [],
187            |row| row.get(0),
188        )?;
189        let grand_saved: i64 = rows.iter().map(|r| r.saved).sum();
190        let grand_output = grand_input - grand_saved;
191        let grand_pct = if grand_input > 0 {
192            (grand_saved as f64 / grand_input as f64) * 100.0
193        } else {
194            0.0
195        };
196
197        let saved_tokens = grand_saved / 4;
198
199        // ── Header ──────────────────────────────────────────
200        println!();
201        println!("  {BOLD}{GREEN}MCP-RTK{RESET}{DIM} - Token Savings{RESET}");
202        println!("  {DIM}{}{RESET}", "─".repeat(56));
203        println!();
204
205        // ── Summary (two columns) ───────────────────────────
206        println!(
207            "  {DIM}Calls{RESET}  {BOLD}{WHITE}{:<12}{RESET}  {DIM}Input{RESET}   {WHITE}{} tokens{RESET}",
208            grand_calls,
209            format_number(grand_input / 4),
210        );
211        println!(
212            "  {DIM}Saved{RESET}  {BOLD}{GREEN}{:<12}{RESET}  {DIM}Output{RESET}  {WHITE}{} tokens{RESET}",
213            format!("{} ({:.0}%)", format_number(saved_tokens), grand_pct),
214            format_number(grand_output / 4),
215        );
216        println!();
217
218        // ── Efficiency bar ──────────────────────────────────
219        let bar_width: usize = 40;
220        let bar = render_block_bar(grand_pct / 100.0, bar_width);
221        let pct_color = pct_to_color(grand_pct);
222        println!("  {bar}  {pct_color}{BOLD}{:.1}%{RESET}", grand_pct);
223        println!();
224
225        // ── Per-tool table ──────────────────────────────────
226        if rows.is_empty() {
227            println!("  {DIM}No tool calls recorded yet.{RESET}");
228            println!();
229            return Ok(());
230        }
231
232        // Collect unique presets in insertion order
233        let mut seen = std::collections::HashSet::new();
234        let mut presets: Vec<String> = Vec::new();
235        for row in &rows {
236            if seen.insert(row.preset.clone()) {
237                presets.push(row.preset.clone());
238            }
239        }
240
241        let max_saved = rows.iter().map(|r| r.saved).max().unwrap_or(1).max(1);
242
243        for preset in &presets {
244            let preset_rows: Vec<&ToolRow> = rows.iter().filter(|r| &r.preset == preset).collect();
245            let preset_saved: i64 = preset_rows.iter().map(|r| r.saved).sum();
246            let preset_calls: i64 = preset_rows.iter().map(|r| r.calls).sum();
247
248            println!(
249                "  {DIM}─── {RESET}{BOLD}{}{RESET}{DIM} ({} calls, {} saved) {}─{RESET}",
250                preset,
251                preset_calls,
252                format_tokens(preset_saved),
253                "─".repeat(30usize.saturating_sub(preset.len())),
254            );
255            println!();
256            println!(
257                "  {DIM}{:<28} {:>5}  {:>8}  {:>5}{RESET}",
258                "Tool", "Count", "Saved", "Avg%"
259            );
260            println!();
261
262            for row in &preset_rows {
263                let pct_color = pct_to_color(row.avg_pct);
264                let bar_ratio = row.saved as f64 / max_saved as f64;
265                let inline_bar = render_block_bar(bar_ratio, 16);
266
267                println!(
268                    "  {BOLD}{WHITE}{:<28}{RESET} {:>5}  {:>8}  {pct_color}{:>4.0}%{RESET}  {inline_bar}",
269                    truncate_name(&row.name, 28),
270                    row.calls,
271                    format_tokens(row.saved),
272                    row.avg_pct,
273                );
274            }
275
276            println!();
277        }
278
279        println!();
280        Ok(())
281    }
282
283    /// Print the last 50 tool calls with timestamps and savings percentages.
284    ///
285    /// # Errors
286    ///
287    /// Returns an error if the database lock is poisoned or query fails.
288    pub fn print_history(&self) -> Result<()> {
289        let conn = self
290            .conn
291            .lock()
292            .map_err(|e| anyhow::anyhow!("lock poisoned: {e}"))?;
293        let mut stmt = conn.prepare(
294            "SELECT timestamp, tool_name, input_bytes, output_bytes, savings_pct, preset
295             FROM tool_calls
296             ORDER BY timestamp DESC
297             LIMIT 50",
298        )?;
299
300        let rows: Vec<(String, String, i64, i64, f64, String)> = stmt
301            .query_map([], |row| {
302                Ok((
303                    row.get::<_, String>(0)?,
304                    row.get::<_, String>(1)?,
305                    row.get::<_, i64>(2)?,
306                    row.get::<_, i64>(3)?,
307                    row.get::<_, f64>(4)?,
308                    row.get::<_, String>(5)?,
309                ))
310            })?
311            .filter_map(|r| r.ok())
312            .collect();
313
314        println!();
315        println!("  {BOLD}{GREEN}MCP-RTK{RESET}{DIM} ── Recent Calls{RESET}");
316        println!("  {DIM}{}{RESET}", "─".repeat(76));
317        println!();
318
319        if rows.is_empty() {
320            println!("  {DIM}No tool calls recorded yet.{RESET}");
321            println!();
322            return Ok(());
323        }
324
325        println!(
326            "  {DIM}{:<19} {:<8} {:<22} {:>7} {:>7} {:>6}{RESET}",
327            "Timestamp", "Preset", "Tool", "In", "Out", "Saved"
328        );
329        println!();
330
331        for (ts, name, input, output, pct, preset) in &rows {
332            let pct_color = pct_to_color(*pct);
333            let saved_bytes = input - output;
334
335            println!(
336                "  {DIM}{:<19}{RESET} {YELLOW}{:<8}{RESET} {WHITE}{:<22}{RESET} {:>7} {:>7} {pct_color}{BOLD}{:>5.0}%{RESET}  {DIM}{}{RESET}",
337                ts.get(..19).unwrap_or(ts),
338                truncate_name(preset, 8),
339                truncate_name(name, 22),
340                format_tokens(*input),
341                format_tokens(*output),
342                pct,
343                if saved_bytes > 0 {
344                    format!("-{} tk", format_tokens(saved_bytes))
345                } else {
346                    String::new()
347                },
348            );
349        }
350
351        println!();
352        Ok(())
353    }
354
355    /// Return all tracking stats as a [`serde_json::Value`] for programmatic use.
356    ///
357    /// # Errors
358    ///
359    /// Returns an error if the database lock is poisoned or query fails.
360    pub fn stats_as_json(&self) -> Result<serde_json::Value> {
361        let conn = self
362            .conn
363            .lock()
364            .map_err(|e| anyhow::anyhow!("lock poisoned: {e}"))?;
365
366        // Grand totals
367        let (total_calls, total_input, total_output, total_saved): (i64, i64, i64, i64) =
368            conn.query_row(
369                "SELECT COUNT(*), COALESCE(SUM(input_bytes),0), COALESCE(SUM(output_bytes),0), COALESCE(SUM(saved_bytes),0) FROM tool_calls",
370                [],
371                |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?, row.get(3)?)),
372            )?;
373
374        let grand_pct = if total_input > 0 {
375            (total_saved as f64 / total_input as f64) * 100.0
376        } else {
377            0.0
378        };
379
380        // Per-preset, per-tool breakdown
381        let mut stmt = conn.prepare(
382            "SELECT preset, tool_name, COUNT(*) as calls, SUM(input_bytes), SUM(output_bytes), SUM(saved_bytes), AVG(savings_pct)
383             FROM tool_calls GROUP BY preset, tool_name ORDER BY preset, SUM(saved_bytes) DESC",
384        )?;
385
386        let rows: Vec<(String, String, i64, i64, i64, i64, f64)> = stmt
387            .query_map([], |row| {
388                Ok((
389                    row.get(0)?,
390                    row.get(1)?,
391                    row.get(2)?,
392                    row.get(3)?,
393                    row.get(4)?,
394                    row.get(5)?,
395                    row.get(6)?,
396                ))
397            })?
398            .filter_map(|r| r.ok())
399            .collect();
400
401        // Build JSON
402        let mut presets_map: serde_json::Map<String, serde_json::Value> = serde_json::Map::new();
403        for (preset, tool, calls, input, output, saved, avg_pct) in &rows {
404            let preset_entry = presets_map
405                .entry(preset.clone())
406                .or_insert_with(|| {
407                    serde_json::json!({"calls": 0, "input_bytes": 0, "output_bytes": 0, "saved_bytes": 0, "tools": {}})
408                });
409            let preset_obj = preset_entry.as_object_mut().unwrap();
410            *preset_obj.get_mut("calls").unwrap() =
411                serde_json::json!(preset_obj["calls"].as_i64().unwrap() + calls);
412            *preset_obj.get_mut("input_bytes").unwrap() =
413                serde_json::json!(preset_obj["input_bytes"].as_i64().unwrap() + input);
414            *preset_obj.get_mut("output_bytes").unwrap() =
415                serde_json::json!(preset_obj["output_bytes"].as_i64().unwrap() + output);
416            *preset_obj.get_mut("saved_bytes").unwrap() =
417                serde_json::json!(preset_obj["saved_bytes"].as_i64().unwrap() + saved);
418
419            let tools = preset_obj
420                .get_mut("tools")
421                .unwrap()
422                .as_object_mut()
423                .unwrap();
424            tools.insert(
425                tool.clone(),
426                serde_json::json!({
427                    "calls": calls,
428                    "input_bytes": input,
429                    "output_bytes": output,
430                    "saved_bytes": saved,
431                    "avg_savings_pct": (avg_pct * 10.0).round() / 10.0,
432                }),
433            );
434        }
435
436        let output = serde_json::json!({
437            "total_calls": total_calls,
438            "total_input_bytes": total_input,
439            "total_output_bytes": total_output,
440            "total_saved_bytes": total_saved,
441            "total_input_tokens": total_input / 4,
442            "total_output_tokens": total_output / 4,
443            "total_saved_tokens": total_saved / 4,
444            "savings_pct": (grand_pct * 10.0).round() / 10.0,
445            "presets": presets_map,
446        });
447
448        Ok(output)
449    }
450
451    /// Export all tracking stats as pretty-printed JSON to stdout.
452    ///
453    /// # Errors
454    ///
455    /// Returns an error if the database lock is poisoned or query fails.
456    pub fn export_json(&self) -> Result<()> {
457        let output = self.stats_as_json()?;
458        println!("{}", serde_json::to_string_pretty(&output).unwrap());
459        Ok(())
460    }
461
462    /// Return the set of preset names that have tracking data.
463    ///
464    /// Used by `discover` to detect which servers are already proxied.
465    pub fn tracked_presets(&self) -> Result<std::collections::HashSet<String>> {
466        let conn = self
467            .conn
468            .lock()
469            .map_err(|e| anyhow::anyhow!("lock poisoned: {e}"))?;
470        let mut stmt =
471            conn.prepare("SELECT DISTINCT preset FROM tool_calls WHERE preset != 'unknown'")?;
472        let presets: std::collections::HashSet<String> = stmt
473            .query_map([], |row| row.get::<_, String>(0))?
474            .filter_map(|r| r.ok())
475            .collect();
476        Ok(presets)
477    }
478}
479
480/// Expand `~/` prefix to the user's home directory.
481fn expand_path(path: &str) -> PathBuf {
482    if let Some(rest) = path.strip_prefix("~/") {
483        if let Ok(home) = std::env::var("HOME").or_else(|_| std::env::var("USERPROFILE")) {
484            return PathBuf::from(home).join(rest);
485        }
486    }
487    PathBuf::from(path)
488}
mcp_rtk/tracking.rs

mcp_rtk/
tracking.rs
