ati/proxy/
client.rs

1/// Proxy client — forwards tool calls to an external ATI proxy server.
2///
3/// When ATI_PROXY_URL is set, `ati run <tool>` sends tool_name + args
4/// to the proxy. Authentication is via JWT in the Authorization header
5/// (ATI_SESSION_TOKEN env var).
6use reqwest::Client;
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use std::collections::HashMap;
10use std::time::Duration;
11use thiserror::Error;
12
13#[derive(Error, Debug)]
14pub enum ProxyError {
15    #[error("Proxy request failed: {0}")]
16    Request(#[from] reqwest::Error),
17    #[error("Proxy error ({status}): {body}")]
18    ProxyResponse { status: u16, body: String },
19    #[error("Invalid proxy URL: {0}")]
20    InvalidUrl(String),
21    #[error("Proxy returned invalid response: {0}")]
22    InvalidResponse(String),
23}
24
25/// Request payload sent to the proxy server's /call endpoint.
26#[derive(Debug, Serialize)]
27pub struct ProxyCallRequest {
28    pub tool_name: String,
29    /// Tool arguments — JSON object for HTTP/MCP tools, or JSON array for CLI tools.
30    pub args: Value,
31}
32
33/// Response payload from the proxy server.
34#[derive(Debug, Deserialize)]
35pub struct ProxyCallResponse {
36    pub result: Value,
37    #[serde(default)]
38    pub error: Option<String>,
39}
40
41/// Request payload for the proxy's /help endpoint.
42#[derive(Debug, Serialize)]
43pub struct ProxyHelpRequest {
44    pub query: String,
45    #[serde(skip_serializing_if = "Option::is_none")]
46    pub tool: Option<String>,
47}
48
49/// Response from the proxy's /help endpoint.
50#[derive(Debug, Deserialize)]
51pub struct ProxyHelpResponse {
52    pub content: String,
53    #[serde(default)]
54    pub error: Option<String>,
55}
56
57const PROXY_TIMEOUT_SECS: u64 = 120;
58
59/// Build an HTTP request builder with JWT Bearer auth from ATI_SESSION_TOKEN.
60fn build_proxy_request(
61    client: &Client,
62    method: reqwest::Method,
63    url: &str,
64) -> reqwest::RequestBuilder {
65    let mut req = client.request(method, url);
66    if let Ok(token) = std::env::var("ATI_SESSION_TOKEN") {
67        if !token.is_empty() {
68            req = req.header("Authorization", format!("Bearer {token}"));
69        }
70    }
71    req
72}
73
74/// Execute a tool call via the proxy server.
75///
76/// POST {proxy_url}/call with JSON body: { tool_name, args }
77/// Scopes are carried inside the JWT — not in the request body.
78///
79/// `args` carries key-value pairs for HTTP/MCP tools.
80/// `raw_args`, if provided, is sent as an array in the `args` field for CLI tools.
81pub async fn call_tool(
82    proxy_url: &str,
83    tool_name: &str,
84    args: &HashMap<String, Value>,
85    raw_args: Option<&[String]>,
86) -> Result<Value, ProxyError> {
87    let client = Client::builder()
88        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
89        .build()?;
90
91    let url = format!("{}/call", proxy_url.trim_end_matches('/'));
92
93    // If raw_args are provided (CLI tool), send them as a JSON array in `args`.
94    // Otherwise send the key-value map.
95    let args_value = match raw_args {
96        Some(raw) if !raw.is_empty() => {
97            Value::Array(raw.iter().map(|s| Value::String(s.clone())).collect())
98        }
99        _ => serde_json::to_value(args).unwrap_or(Value::Object(serde_json::Map::new())),
100    };
101
102    let payload = ProxyCallRequest {
103        tool_name: tool_name.to_string(),
104        args: args_value,
105    };
106
107    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
108        .json(&payload)
109        .send()
110        .await?;
111    let status = response.status();
112
113    if !status.is_success() {
114        let body = response.text().await.unwrap_or_else(|_| "empty".into());
115        return Err(ProxyError::ProxyResponse {
116            status: status.as_u16(),
117            body,
118        });
119    }
120
121    let body: ProxyCallResponse = response
122        .json()
123        .await
124        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))?;
125
126    if let Some(err) = body.error {
127        return Err(ProxyError::ProxyResponse {
128            status: 200,
129            body: err,
130        });
131    }
132
133    Ok(body.result)
134}
135
136/// List available tools from the proxy.
137pub async fn list_tools(proxy_url: &str, query_params: &str) -> Result<Value, ProxyError> {
138    let client = Client::builder()
139        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
140        .build()?;
141    let mut url = format!("{}/tools", proxy_url.trim_end_matches('/'));
142    if !query_params.is_empty() {
143        url.push('?');
144        url.push_str(query_params);
145    }
146    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
147        .send()
148        .await?;
149    let status = response.status();
150    if !status.is_success() {
151        let body = response.text().await.unwrap_or_default();
152        return Err(ProxyError::ProxyResponse {
153            status: status.as_u16(),
154            body,
155        });
156    }
157    Ok(response.json().await?)
158}
159
160/// Get detailed info about a specific tool from the proxy.
161pub async fn get_tool_info(proxy_url: &str, name: &str) -> Result<Value, ProxyError> {
162    let client = Client::builder()
163        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
164        .build()?;
165    let url = format!("{}/tools/{}", proxy_url.trim_end_matches('/'), name);
166    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
167        .send()
168        .await?;
169    let status = response.status();
170    if !status.is_success() {
171        let body = response.text().await.unwrap_or_default();
172        return Err(ProxyError::ProxyResponse {
173            status: status.as_u16(),
174            body,
175        });
176    }
177    Ok(response.json().await?)
178}
179
180/// Forward a raw MCP JSON-RPC message via the proxy's /mcp endpoint.
181pub async fn call_mcp(
182    proxy_url: &str,
183    method: &str,
184    params: Option<Value>,
185) -> Result<Value, ProxyError> {
186    use std::sync::atomic::{AtomicU64, Ordering};
187    static MCP_ID: AtomicU64 = AtomicU64::new(1);
188
189    let id = MCP_ID.fetch_add(1, Ordering::SeqCst);
190    let msg = serde_json::json!({
191        "jsonrpc": "2.0",
192        "id": id,
193        "method": method,
194        "params": params,
195    });
196
197    let client = Client::builder()
198        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
199        .build()?;
200
201    let url = format!("{}/mcp", proxy_url.trim_end_matches('/'));
202
203    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
204        .json(&msg)
205        .send()
206        .await?;
207    let status = response.status();
208
209    if status == reqwest::StatusCode::ACCEPTED {
210        return Ok(Value::Null);
211    }
212
213    if !status.is_success() {
214        let body = response.text().await.unwrap_or_else(|_| "empty".into());
215        return Err(ProxyError::ProxyResponse {
216            status: status.as_u16(),
217            body,
218        });
219    }
220
221    let body: Value = response
222        .json()
223        .await
224        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))?;
225
226    if let Some(err) = body.get("error") {
227        let message = err
228            .get("message")
229            .and_then(|m| m.as_str())
230            .unwrap_or("MCP proxy error");
231        return Err(ProxyError::ProxyResponse {
232            status: 200,
233            body: message.to_string(),
234        });
235    }
236
237    Ok(body.get("result").cloned().unwrap_or(Value::Null))
238}
239
240/// Fetch skill list from the proxy server.
241pub async fn list_skills(
242    proxy_url: &str,
243    query_params: &str,
244) -> Result<serde_json::Value, ProxyError> {
245    let client = Client::builder()
246        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
247        .build()?;
248
249    let url = if query_params.is_empty() {
250        format!("{}/skills", proxy_url.trim_end_matches('/'))
251    } else {
252        format!("{}/skills?{query_params}", proxy_url.trim_end_matches('/'))
253    };
254
255    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
256        .send()
257        .await?;
258    let status = response.status();
259
260    if !status.is_success() {
261        let body = response.text().await.unwrap_or_else(|_| "empty".into());
262        return Err(ProxyError::ProxyResponse {
263            status: status.as_u16(),
264            body,
265        });
266    }
267
268    response
269        .json()
270        .await
271        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
272}
273
274/// Fetch a skill's detail from the proxy server.
275pub async fn get_skill(
276    proxy_url: &str,
277    name: &str,
278    query_params: &str,
279) -> Result<serde_json::Value, ProxyError> {
280    let client = Client::builder()
281        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
282        .build()?;
283
284    let url = if query_params.is_empty() {
285        format!("{}/skills/{name}", proxy_url.trim_end_matches('/'))
286    } else {
287        format!(
288            "{}/skills/{name}?{query_params}",
289            proxy_url.trim_end_matches('/')
290        )
291    };
292
293    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
294        .send()
295        .await?;
296    let status = response.status();
297
298    if !status.is_success() {
299        let body = response.text().await.unwrap_or_else(|_| "empty".into());
300        return Err(ProxyError::ProxyResponse {
301            status: status.as_u16(),
302            body,
303        });
304    }
305
306    response
307        .json()
308        .await
309        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
310}
311
312/// Resolve skills for given scopes via the proxy.
313pub async fn resolve_skills(
314    proxy_url: &str,
315    scopes: &serde_json::Value,
316) -> Result<serde_json::Value, ProxyError> {
317    let client = Client::builder()
318        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
319        .build()?;
320
321    let url = format!("{}/skills/resolve", proxy_url.trim_end_matches('/'));
322
323    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
324        .json(scopes)
325        .send()
326        .await?;
327    let status = response.status();
328
329    if !status.is_success() {
330        let body = response.text().await.unwrap_or_else(|_| "empty".into());
331        return Err(ProxyError::ProxyResponse {
332            status: status.as_u16(),
333            body,
334        });
335    }
336
337    response
338        .json()
339        .await
340        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
341}
342
343/// Execute an LLM help query via the proxy server.
344pub async fn call_help(
345    proxy_url: &str,
346    query: &str,
347    tool: Option<&str>,
348) -> Result<String, ProxyError> {
349    let client = Client::builder()
350        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
351        .build()?;
352
353    let url = format!("{}/help", proxy_url.trim_end_matches('/'));
354
355    let payload = ProxyHelpRequest {
356        query: query.to_string(),
357        tool: tool.map(|t| t.to_string()),
358    };
359
360    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
361        .json(&payload)
362        .send()
363        .await?;
364    let status = response.status();
365
366    if !status.is_success() {
367        let body = response.text().await.unwrap_or_else(|_| "empty".into());
368        return Err(ProxyError::ProxyResponse {
369            status: status.as_u16(),
370            body,
371        });
372    }
373
374    let body: ProxyHelpResponse = response
375        .json()
376        .await
377        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))?;
378
379    if let Some(err) = body.error {
380        return Err(ProxyError::ProxyResponse {
381            status: 200,
382            body: err,
383        });
384    }
385
386    Ok(body.content)
387}
ati/proxy/client.rs

ati/proxy/
client.rs
