rtk_mcp/
rtk-mcp.rs

1//! `rtk-mcp` — the plugin's MCP server, spoken over stdio.
2//!
3//! Antigravity launches this binary and talks line-delimited JSON-RPC to it on
4//! stdin/stdout. All the judgement lives in [`rtk_mcp_agy`]; this file only
5//! moves bytes and spawns the two external processes — `rtk` to rewrite a
6//! command, and a shell to run it.
7//!
8//! stdout belongs to the protocol: one JSON response per line, nothing else.
9
10use rtk_mcp_agy::{handle_line, parse_rewrite, resolve_rtk_bin, resolve_shell, Execution};
11use std::env;
12use std::io::{self, BufRead, Write};
13use std::process::Command;
14
15fn main() {
16    let args: Vec<String> = env::args().collect();
17
18    if args.iter().any(|a| a == "--version" || a == "-V") {
19        println!("rtk-mcp {}", rtk_mcp_agy::version());
20        return;
21    }
22
23    if args.iter().any(|a| a == "--help" || a == "-h") {
24        println!(
25            "rtk-mcp {} — MCP server for Antigravity rtk shell command optimization",
26            rtk_mcp_agy::version()
27        );
28        println!();
29        println!("Usage: rtk-mcp");
30        println!("  JSON-RPC MCP server over stdin/stdout. Exposes 4 tools:");
31        println!(
32            "    {}      — Rewrites the given CommandLine via `rtk rewrite` and executes it.",
33            rtk_mcp_agy::TOOL_NAME
34        );
35        println!(
36            "    {}     — Reports rtk token-savings statistics.",
37            rtk_mcp_agy::TOOL_GAIN
38        );
39        println!(
40            "    {}    — Shows how rtk would rewrite a shell command without running it.",
41            rtk_mcp_agy::TOOL_CHECK
42        );
43        println!(
44            "    {} — Analyzes past session history for missed token savings.",
45            rtk_mcp_agy::TOOL_DISCOVER
46        );
47        return;
48    }
49
50    let stdin = io::stdin();
51    let mut stdout = io::stdout();
52
53    let mut lines = stdin.lock().lines();
54    while let Some(Ok(line)) = lines.next() {
55        let Some(response) = handle_line(&line, &rtk_rewrite, &run_command) else {
56            continue;
57        };
58        if let Ok(response_str) = serde_json::to_string(&response) {
59            println!("{response_str}");
60            let _ = stdout.flush();
61        }
62    }
63}
64
65/// Ask `rtk` for a token-optimized form of `command`.
66///
67/// Any failure — rtk absent, non-zero exit, no output — is reported as "no
68/// rewrite available", and the caller runs the original command. rtk only
69/// covers specific known commands, so this is the ordinary path, not an error.
70fn rtk_rewrite(command: &str) -> Option<String> {
71    let output = Command::new(resolve_rtk_bin(env::var("RTK_BIN").ok()))
72        .arg("rewrite")
73        .arg("--")
74        .arg(command)
75        .output()
76        .ok()?;
77    parse_rewrite(&String::from_utf8_lossy(&output.stdout))
78}
79
80/// Run `command` in `cwd` through a shell.
81///
82/// `pwsh` on Windows, `sh` elsewhere, overridable with `RTK_AGY_SHELL`. Both
83/// take the command after `-c`. A shell that cannot be spawned is reported to
84/// the model as a tool error rather than taking the server down.
85fn run_command(command: &str, cwd: &str) -> Execution {
86    let shell = resolve_shell(env::var("RTK_AGY_SHELL").ok(), cfg!(windows));
87    match Command::new(&shell)
88        .arg("-c")
89        .arg(command)
90        .current_dir(cwd)
91        .output()
92    {
93        Ok(out) => Execution::Ran {
94            stdout: String::from_utf8_lossy(&out.stdout).to_string(),
95            stderr: String::from_utf8_lossy(&out.stderr).to_string(),
96            success: out.status.success(),
97        },
98        Err(e) => Execution::NotRun(format!("could not run {shell}: {e}")),
99    }
100}
rtk_mcp/rtk-mcp.rs

rtk_mcp/
rtk-mcp.rs
