Skip to main content

aether_cli/
mcp_command.rs

1use aether_core::mcp::{AETHER_MCP_IPC_SOCKET, DeferredServer, DeferredToolsClient, DeferredToolsError};
2use clap::{ArgAction, Args};
3use mcp_utils::client::ToolCallError;
4use mcp_utils::model::{CallToolResult, Tool};
5use serde_json::{Map, Value};
6use std::env::var_os;
7use std::io::{self, IsTerminal, Read};
8use std::path::PathBuf;
9use std::time::Duration;
10
11const DEFAULT_CALL_TIMEOUT_SECONDS: u64 = 600;
12
13#[derive(Debug, Args)]
14#[command(disable_help_flag = true)]
15pub struct McpArgs {
16    #[arg(value_name = "SERVER")]
17    server: Option<String>,
18
19    #[arg(value_name = "TOOL", requires = "server")]
20    tool: Option<String>,
21
22    #[arg(long, action = ArgAction::SetTrue, conflicts_with_all = ["json", "timeout_seconds"])]
23    help: bool,
24
25    #[arg(long, value_name = "OBJECT", requires = "tool")]
26    json: Option<String>,
27
28    #[arg(
29        long = "timeout",
30        value_name = "SECONDS",
31        default_value_t = DEFAULT_CALL_TIMEOUT_SECONDS,
32        value_parser = clap::value_parser!(u64).range(1..),
33        requires = "tool"
34    )]
35    timeout_seconds: u64,
36}
37
38#[derive(Debug, thiserror::Error)]
39pub enum McpCommandError {
40    #[error("{0}")]
41    Usage(String),
42    #[error("`aether mcp` requires the inherited {AETHER_MCP_IPC_SOCKET} from an active Aether session")]
43    SessionUnavailable,
44    #[error("invalid {AETHER_MCP_IPC_SOCKET}: socket path must be absolute")]
45    RelativeSocket,
46    #[error(transparent)]
47    Session(#[from] DeferredToolsError),
48    #[error("MCP request timed out after {0} seconds")]
49    Timeout(u64),
50    #[error("MCP tool call failed: {0}")]
51    Call(#[source] ToolCallError),
52    #[error("deferred tool returned an error: {0}")]
53    Tool(String),
54    #[error("failed to read JSON from stdin: {0}")]
55    Stdin(#[source] io::Error),
56    #[error("failed to print JSON result: {0}")]
57    Output(#[source] serde_json::Error),
58}
59
60#[derive(Debug)]
61enum Request {
62    Help(HelpLevel),
63    Call { server: String, tool: String, json: Map<String, Value>, timeout_seconds: u64 },
64}
65
66#[derive(Debug)]
67enum HelpLevel {
68    Servers,
69    Server(String),
70    Tool { server: String, tool: String },
71}
72
73impl McpCommandError {
74    pub fn exit_code(&self) -> u8 {
75        match self {
76            Self::Usage(_) | Self::Stdin(_) => 2,
77            Self::SessionUnavailable
78            | Self::RelativeSocket
79            | Self::Session(_)
80            | Self::Timeout(_)
81            | Self::Call(_)
82            | Self::Tool(_)
83            | Self::Output(_) => 1,
84        }
85    }
86}
87
88pub async fn run(args: McpArgs) -> Result<(), McpCommandError> {
89    let request = Request::try_from(args)?;
90    let client = DeferredToolsClient::connect(inherited_socket()?).await?;
91    execute_request(&client, request).await
92}
93
94impl TryFrom<McpArgs> for Request {
95    type Error = McpCommandError;
96
97    fn try_from(args: McpArgs) -> Result<Self, Self::Error> {
98        match (args.server, args.tool, args.help) {
99            (None, None, _) => Ok(Request::Help(HelpLevel::Servers)),
100            (Some(server), None, true) => Ok(Request::Help(HelpLevel::Server(server))),
101            (Some(server), Some(tool), true) => Ok(Request::Help(HelpLevel::Tool { server, tool })),
102            (Some(server), Some(tool), false) => {
103                let json = if let Some(input) = args.json.as_deref() {
104                    parse_json_object(Some(input))?
105                } else {
106                    let stdin = read_stdin()?;
107                    parse_json_object(stdin.as_deref())?
108                };
109                Ok(Request::Call { server, tool, json, timeout_seconds: args.timeout_seconds })
110            }
111            (Some(_), None, false) | (None, Some(_), _) => Err(McpCommandError::Usage(
112                "usage: aether mcp <server> <tool> [--json <object>] [--timeout <seconds>]".into(),
113            )),
114        }
115    }
116}
117
118fn inherited_socket() -> Result<PathBuf, McpCommandError> {
119    let path = PathBuf::from(var_os(AETHER_MCP_IPC_SOCKET).ok_or(McpCommandError::SessionUnavailable)?);
120    if path.is_absolute() { Ok(path) } else { Err(McpCommandError::RelativeSocket) }
121}
122
123fn read_stdin() -> Result<Option<String>, McpCommandError> {
124    let mut stdin = io::stdin();
125    if stdin.is_terminal() {
126        return Ok(None);
127    }
128    let mut input = String::new();
129    stdin.read_to_string(&mut input).map_err(McpCommandError::Stdin)?;
130    Ok((!input.trim().is_empty()).then_some(input))
131}
132
133fn parse_json_object(input: Option<&str>) -> Result<Map<String, Value>, McpCommandError> {
134    let Some(input) = input else { return Ok(Map::new()) };
135    serde_json::from_str(input).map_err(|error| match error.classify() {
136        serde_json::error::Category::Data => McpCommandError::Usage("tool input must be a JSON object".into()),
137        _ => McpCommandError::Usage(format!("invalid JSON input: {error}")),
138    })
139}
140
141async fn execute_request(client: &DeferredToolsClient, request: Request) -> Result<(), McpCommandError> {
142    match request {
143        Request::Help(HelpLevel::Servers) => show_servers_help(client).await,
144        Request::Help(HelpLevel::Server(server)) => show_server_help(client, &server).await,
145        Request::Help(HelpLevel::Tool { server, tool }) => show_tool_help(client, &server, &tool).await,
146        Request::Call { server, tool, json, timeout_seconds } => {
147            call_tool(client, &server, &tool, json, timeout_seconds).await
148        }
149    }
150}
151
152async fn show_servers_help(client: &DeferredToolsClient) -> Result<(), McpCommandError> {
153    print_servers_help(&client.servers().await?);
154    Ok(())
155}
156
157async fn show_server_help(client: &DeferredToolsClient, server: &str) -> Result<(), McpCommandError> {
158    let servers = client.servers().await?;
159    let deferred = servers
160        .iter()
161        .find(|deferred| deferred.name == server)
162        .ok_or_else(|| McpCommandError::Usage(format!("unknown deferred server `{server}`")))?;
163    print_server_help(deferred);
164    Ok(())
165}
166
167async fn show_tool_help(client: &DeferredToolsClient, server: &str, tool: &str) -> Result<(), McpCommandError> {
168    let definition = client
169        .tool(server, tool)
170        .await?
171        .ok_or_else(|| McpCommandError::Usage(format!("unknown deferred tool `{server} {tool}`")))?;
172    print_tool_help(server, tool, &definition)
173}
174
175async fn call_tool(
176    client: &DeferredToolsClient,
177    server: &str,
178    tool: &str,
179    json: Map<String, Value>,
180    timeout_seconds: u64,
181) -> Result<(), McpCommandError> {
182    let call = client.call(server, tool, json, Duration::from_secs(timeout_seconds));
183    let result = call.result().await.map_err(call_error)?;
184    if result.is_error.unwrap_or(false) {
185        return Err(McpCommandError::Tool(result_json(&result).to_string()));
186    }
187    println!("{}", serde_json::to_string(&result_json(&result)).map_err(McpCommandError::Output)?);
188    Ok(())
189}
190
191fn call_error(error: ToolCallError) -> McpCommandError {
192    match error {
193        ToolCallError::TimedOut(timeout) => McpCommandError::Timeout(timeout.as_secs()),
194        error => McpCommandError::Call(error),
195    }
196}
197
198fn print_servers_help(servers: &[DeferredServer]) {
199    println!("Discover and call deferred MCP tools through the active Aether session.\n");
200    println!("Usage: aether mcp <server> --help\n");
201    println!("Deferred MCP servers:");
202    for server in servers {
203        println!("  {:<20} {}", server.name, server.description);
204    }
205}
206
207fn print_server_help(server: &DeferredServer) {
208    let name = &server.name;
209    println!("Deferred tools from `{name}`.\n");
210    println!("Usage: aether mcp {name} <tool> --help\n");
211    println!("Tools:");
212    for tool in &server.tools {
213        println!("  {:<20} {}", tool.name, tool.description.as_deref().unwrap_or_default());
214    }
215}
216
217fn print_tool_help(server: &str, tool: &str, definition: &Tool) -> Result<(), McpCommandError> {
218    println!("{}\n", definition.description.as_deref().unwrap_or_default());
219    println!("Usage:");
220    println!("  aether mcp {server} {tool} --json '{{...}}'");
221    println!("  printf '%s' '{{...}}' | aether mcp {server} {tool}\n");
222    println!("Input schema:");
223    println!("{}", serde_json::to_string_pretty(definition.input_schema.as_ref()).map_err(McpCommandError::Output)?);
224    Ok(())
225}
226
227fn result_json(result: &CallToolResult) -> Value {
228    result.structured_content.clone().unwrap_or_else(|| serde_json::to_value(&result.content).unwrap_or(Value::Null))
229}