rtk_cc_mcp/
mcp.rs

1//! `rtk-cc-mcp` — the plugin's MCP server, spoken over stdio.
2//!
3//! Claude Code launches this binary via `.mcp.json` and talks JSON-RPC to it on
4//! stdin/stdout. Two rules follow from that, and both are easy to get wrong:
5//!
6//! 1. **stdout belongs to the protocol.** Every log line goes to stderr. A
7//!    stray `println!` corrupts the stream and the server appears to hang.
8//! 2. **The process is long-lived.** It stays up until Claude Code closes the
9//!    pipe, so `waiting()` is what keeps it alive.
10//!
11//! The tools here are rtk's *analytics* surface. Three of them are read-only.
12//! The fourth, `rtk_proxy`, executes — see its doc comment.
13
14use rmcp::{
15    handler::server::{router::tool::ToolRouter, wrapper::Parameters},
16    model::{Implementation, ServerCapabilities, ServerConfig},
17    tool, tool_handler, tool_router,
18    transport::stdio,
19    ErrorData, ServerHandler, ServiceExt,
20};
21use rtk_mcp_cc::{
22    check_args, discover_args, gain_args, proxy_args, Config, DiscoverOptions, GainOptions,
23};
24use schemars::JsonSchema;
25use serde::Deserialize;
26use std::path::PathBuf;
27use std::process::Command;
28use tracing::info;
29
30// ── Tool parameter types ───────────────────────────────────────────────
31//
32// Each `#[tool]` takes one `Parameters<T>`; `T`'s JsonSchema derive becomes the
33// tool's input schema, and the doc comments become the field descriptions the
34// model reads. Describe fields properly — this is the model's only documentation.
35
36/// Arguments for the `rtk_gain` tool.
37#[derive(Debug, Deserialize, JsonSchema)]
38pub struct GainRequest {
39    /// Include the recent per-command history rather than the summary alone.
40    pub history: Option<bool>,
41    /// Restrict the statistics to the current working project.
42    pub project_only: Option<bool>,
43    /// Output format: `text` (default), `json`, or `csv`.
44    pub format: Option<String>,
45}
46
47/// Arguments for the `rtk_discover` tool.
48#[derive(Debug, Deserialize, JsonSchema)]
49pub struct DiscoverRequest {
50    /// Filter to projects whose path contains this substring.
51    pub project: Option<String>,
52    /// Maximum commands reported per section. rtk's default is 15.
53    pub limit: Option<u32>,
54    /// Only consider sessions from the last N days. rtk's default is 30.
55    pub since: Option<u32>,
56    /// Scan every project rather than only the current one.
57    pub all_projects: Option<bool>,
58    /// Output format: `text` (default) or `json`.
59    pub format: Option<String>,
60}
61
62/// Arguments for the `rtk_check` tool.
63#[derive(Debug, Deserialize, JsonSchema)]
64pub struct CheckRequest {
65    /// The shell command to analyze, e.g. `grep -rn foo src/`. It is never
66    /// executed — this reports only how rtk would rewrite it.
67    pub command: String,
68}
69
70/// Arguments for the `rtk_proxy` tool.
71#[derive(Debug, Deserialize, JsonSchema)]
72pub struct ProxyRequest {
73    /// The command as an already-split argument vector, program first —
74    /// e.g. `["git", "status", "--short"]`. This is not a shell command line:
75    /// pipes, redirects, globs and `;` are not interpreted, and each element is
76    /// passed through verbatim as one argument.
77    pub argv: Vec<String>,
78}
79
80// ── Server ─────────────────────────────────────────────────────────────
81
82/// The MCP server. Holds the generated tool router plus the resolved settings.
83#[derive(Debug, Clone)]
84pub struct RtkServer {
85    // Read by the code `#[tool_handler]` generates, which dead-code analysis
86    // does not see through. rmcp's own test suite annotates this same field.
87    #[allow(dead_code)]
88    tool_router: ToolRouter<Self>,
89    config: Config,
90}
91
92impl RtkServer {
93    /// Build a server, loading settings relative to `project_dir`.
94    pub fn new(project_dir: &std::path::Path) -> Self {
95        Self {
96            tool_router: Self::tool_router(),
97            config: Config::load(project_dir),
98        }
99    }
100
101    /// Run rtk with `args` and return its output as the tool result.
102    ///
103    /// stderr is folded into the result rather than dropped: rtk puts its
104    /// advisory notices there, and a tool that silently discards the only
105    /// explanation of an empty result is worse than a noisy one.
106    fn run(&self, args: &[String]) -> Result<String, ErrorData> {
107        let output = Command::new(&self.config.rtk_bin)
108            .args(args)
109            .output()
110            .map_err(|e| {
111                ErrorData::internal_error(
112                    format!(
113                        "could not run {} ({e}). Install rtk, or set RTK_BIN to its path.",
114                        self.config.rtk_bin
115                    ),
116                    None,
117                )
118            })?;
119
120        let stdout = String::from_utf8_lossy(&output.stdout)
121            .trim_end()
122            .to_string();
123        let stderr = String::from_utf8_lossy(&output.stderr)
124            .trim_end()
125            .to_string();
126
127        if !output.status.success() {
128            return Err(ErrorData::internal_error(
129                format!(
130                    "rtk exited with {}: {}",
131                    output.status,
132                    if stderr.is_empty() { &stdout } else { &stderr }
133                ),
134                None,
135            ));
136        }
137
138        Ok(match (stdout.is_empty(), stderr.is_empty()) {
139            (true, true) => "(rtk produced no output)".to_string(),
140            (false, true) => stdout,
141            (true, false) => stderr,
142            (false, false) => format!("{stdout}\n\n[rtk stderr]\n{stderr}"),
143        })
144    }
145}
146
147/// Turn an argument-building error into an MCP invalid-params error.
148fn bad_params(message: String) -> ErrorData {
149    ErrorData::invalid_params(message, None)
150}
151
152#[tool_router]
153impl RtkServer {
154    /// Token-savings analytics.
155    #[tool(
156        name = "rtk_gain",
157        description = "Report rtk's token-savings statistics: how many tokens rtk has saved, and \
158                       optionally the recent per-command history. Read-only."
159    )]
160    async fn rtk_gain(
161        &self,
162        Parameters(GainRequest {
163            history,
164            project_only,
165            format,
166        }): Parameters<GainRequest>,
167    ) -> Result<String, ErrorData> {
168        let opts = GainOptions {
169            history: history.unwrap_or(false),
170            project_only: project_only.unwrap_or(false),
171            format,
172        };
173        let args = gain_args(&opts, &self.config).map_err(bad_params)?;
174        info!(?args, "rtk_gain");
175        self.run(&args)
176    }
177
178    /// Missed-optimization analysis over past sessions.
179    #[tool(
180        name = "rtk_discover",
181        description = "Analyze past Claude Code session history for commands that rtk could have \
182                       optimized but did not, i.e. missed token savings. Read-only."
183    )]
184    async fn rtk_discover(
185        &self,
186        Parameters(DiscoverRequest {
187            project,
188            limit,
189            since,
190            all_projects,
191            format,
192        }): Parameters<DiscoverRequest>,
193    ) -> Result<String, ErrorData> {
194        let opts = DiscoverOptions {
195            project,
196            limit,
197            since,
198            all_projects: all_projects.unwrap_or(false),
199            format,
200        };
201        let args = discover_args(&opts, &self.config).map_err(bad_params)?;
202        info!(?args, "rtk_discover");
203        self.run(&args)
204    }
205
206    /// Dry-run the rewrite engine.
207    #[tool(
208        name = "rtk_check",
209        description = "Show how rtk would rewrite a given shell command, without running it. Use \
210                       this to explain why a command was or was not optimized. Read-only: the \
211                       command is analyzed, never executed."
212    )]
213    async fn rtk_check(
214        &self,
215        Parameters(CheckRequest { command }): Parameters<CheckRequest>,
216    ) -> Result<String, ErrorData> {
217        let args = check_args(&command, &self.config).map_err(bad_params)?;
218        info!(?args, "rtk_check");
219        self.run(&args)
220    }
221
222    /// Execute a command unoptimized, but still counted.
223    ///
224    /// This is the one tool here that runs something. It exists for the case
225    /// rtk's own docs call out — debugging rtk itself, when you need a command
226    /// to bypass the rewrite rules while still being tracked in the statistics.
227    /// It takes an argv vector and spawns it directly, so no shell metacharacter
228    /// is ever interpreted.
229    #[tool(
230        name = "rtk_proxy",
231        description = "Execute a command through `rtk proxy`: it runs UNFILTERED (no rewriting) \
232                       but is still counted in rtk's statistics. Takes an argument vector, not a \
233                       shell line — no pipes, redirects or globbing. This tool EXECUTES the \
234                       command; prefer the ordinary Bash tool unless you specifically need to \
235                       bypass rtk's rewriting while keeping it tracked."
236    )]
237    async fn rtk_proxy(
238        &self,
239        Parameters(ProxyRequest { argv }): Parameters<ProxyRequest>,
240    ) -> Result<String, ErrorData> {
241        let args = proxy_args(&argv, &self.config).map_err(bad_params)?;
242        info!(?args, "rtk_proxy");
243        self.run(&args)
244    }
245}
246
247#[tool_handler]
248impl ServerHandler for RtkServer {
249    fn get_info(&self) -> ServerConfig {
250        ServerConfig::new(ServerCapabilities::builder().enable_tools().build())
251            // Without this the server introduces itself as "rmcp" and is
252            // indistinguishable from every other rmcp server in `/mcp`.
253            //
254            // Do not reach for `Implementation::from_build_env()` here: its
255            // `env!` calls expand inside rmcp's own crate, so it reports
256            // "rmcp" no matter who calls it. The `env!`s below expand here,
257            // which is the whole point.
258            .with_server_info(Implementation::new(
259                env!("CARGO_BIN_NAME"),
260                env!("CARGO_PKG_VERSION"),
261            ))
262            .with_instructions(
263                "Analytics for rtk (Rust Token Killer), which rewrites shell commands into \
264                 token-cheaper equivalents. Use rtk_gain for how much has been saved, \
265                 rtk_discover for savings being missed, and rtk_check to explain how one \
266                 particular command would be rewritten. rtk_proxy executes a command \
267                 unfiltered and is rarely what you want.",
268            )
269    }
270}
271
272// ── Entry point ────────────────────────────────────────────────────────
273
274#[tokio::main]
275async fn main() -> anyhow::Result<()> {
276    let args: Vec<String> = std::env::args().collect();
277    if args.iter().any(|a| a == "--version" || a == "-V") {
278        println!("rtk-cc-mcp {}", env!("CARGO_PKG_VERSION"));
279        return Ok(());
280    }
281    if args.iter().any(|a| a == "--help" || a == "-h") {
282        println!(
283            "rtk-cc-mcp {} — MCP server exposing rtk analytics",
284            env!("CARGO_PKG_VERSION")
285        );
286        println!();
287        println!("Usage: rtk-cc-mcp");
288        println!("  Speaks MCP over stdio. Launched by Claude Code via .mcp.json;");
289        println!("  running it by hand is only useful for debugging.");
290        println!();
291        println!("Tools: rtk_gain, rtk_discover, rtk_check, rtk_proxy");
292        println!();
293        println!("Environment:");
294        println!("  RTK_BIN  Path to the rtk executable (default: `rtk` on PATH)");
295        return Ok(());
296    }
297
298    // stderr, never stdout: stdout is the JSON-RPC channel.
299    tracing_subscriber::fmt()
300        .with_writer(std::io::stderr)
301        .with_env_filter(
302            tracing_subscriber::EnvFilter::try_from_default_env()
303                .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
304        )
305        .init();
306
307    let project_dir = std::env::var("CLAUDE_PROJECT_DIR")
308        .map(PathBuf::from)
309        .or_else(|_| std::env::current_dir())
310        .unwrap_or_else(|_| PathBuf::from("."));
311
312    info!(project_dir = %project_dir.display(), "starting rtk-cc-mcp");
313
314    let service = RtkServer::new(&project_dir).serve(stdio()).await?;
315    service.waiting().await?;
316    Ok(())
317}
rtk_cc_mcp/mcp.rs

rtk_cc_mcp/
mcp.rs
