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
312async fn get_proxy_json(proxy_url: &str, path: &str) -> Result<serde_json::Value, ProxyError> {
313    let client = Client::builder()
314        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
315        .build()?;
316
317    let url = format!(
318        "{}/{}",
319        proxy_url.trim_end_matches('/'),
320        path.trim_start_matches('/')
321    );
322
323    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
324        .send()
325        .await?;
326    let status = response.status();
327
328    if !status.is_success() {
329        let body = response.text().await.unwrap_or_else(|_| "empty".into());
330        return Err(ProxyError::ProxyResponse {
331            status: status.as_u16(),
332            body,
333        });
334    }
335
336    response
337        .json()
338        .await
339        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
340}
341
342async fn get_proxy_json_with_query(
343    proxy_url: &str,
344    path: &str,
345    query: &[(&str, String)],
346) -> Result<serde_json::Value, ProxyError> {
347    let client = Client::builder()
348        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
349        .build()?;
350
351    let mut url = format!(
352        "{}/{}",
353        proxy_url.trim_end_matches('/'),
354        path.trim_start_matches('/')
355    );
356
357    if !query.is_empty() {
358        let params = query
359            .iter()
360            .map(|(key, value)| format!("{key}={}", urlencoding(value)))
361            .collect::<Vec<_>>()
362            .join("&");
363        url.push('?');
364        url.push_str(&params);
365    }
366
367    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
368        .send()
369        .await?;
370    let status = response.status();
371
372    if !status.is_success() {
373        let body = response.text().await.unwrap_or_else(|_| "empty".into());
374        return Err(ProxyError::ProxyResponse {
375            status: status.as_u16(),
376            body,
377        });
378    }
379
380    response
381        .json()
382        .await
383        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
384}
385
386/// List remote SkillATI skills from the proxy server.
387pub async fn get_skillati_catalog(
388    proxy_url: &str,
389    search: Option<&str>,
390) -> Result<serde_json::Value, ProxyError> {
391    let query = search
392        .map(|value| vec![("search", value.to_string())])
393        .unwrap_or_default();
394    get_proxy_json_with_query(proxy_url, "skillati/catalog", &query).await
395}
396
397/// Read a remote SkillATI skill from the proxy server.
398pub async fn get_skillati_read(
399    proxy_url: &str,
400    name: &str,
401) -> Result<serde_json::Value, ProxyError> {
402    get_proxy_json(proxy_url, &format!("skillati/{}", urlencoding(name))).await
403}
404
405/// List bundled resources for a remote SkillATI skill via the proxy server.
406pub async fn get_skillati_resources(
407    proxy_url: &str,
408    name: &str,
409    prefix: Option<&str>,
410) -> Result<serde_json::Value, ProxyError> {
411    let query = prefix
412        .map(|value| vec![("prefix", value.to_string())])
413        .unwrap_or_default();
414    get_proxy_json_with_query(
415        proxy_url,
416        &format!("skillati/{}/resources", urlencoding(name)),
417        &query,
418    )
419    .await
420}
421
422/// Read one arbitrary skill-relative path from a remote SkillATI skill via the proxy server.
423pub async fn get_skillati_file(
424    proxy_url: &str,
425    name: &str,
426    path: &str,
427) -> Result<serde_json::Value, ProxyError> {
428    get_proxy_json_with_query(
429        proxy_url,
430        &format!("skillati/{}/file", urlencoding(name)),
431        &[("path", path.to_string())],
432    )
433    .await
434}
435
436/// List on-demand references for a remote SkillATI skill via the proxy server.
437pub async fn get_skillati_refs(
438    proxy_url: &str,
439    name: &str,
440) -> Result<serde_json::Value, ProxyError> {
441    get_proxy_json(proxy_url, &format!("skillati/{}/refs", urlencoding(name))).await
442}
443
444/// Read one reference file from a remote SkillATI skill via the proxy server.
445pub async fn get_skillati_ref(
446    proxy_url: &str,
447    name: &str,
448    reference: &str,
449) -> Result<serde_json::Value, ProxyError> {
450    get_proxy_json(
451        proxy_url,
452        &format!(
453            "skillati/{}/ref/{}",
454            urlencoding(name),
455            urlencoding(reference)
456        ),
457    )
458    .await
459}
460
461fn urlencoding(s: &str) -> String {
462    s.replace('%', "%25")
463        .replace(' ', "%20")
464        .replace('#', "%23")
465        .replace('&', "%26")
466        .replace('?', "%3F")
467        .replace('/', "%2F")
468        .replace('=', "%3D")
469}
470
471/// Resolve skills for given scopes via the proxy.
472pub async fn resolve_skills(
473    proxy_url: &str,
474    scopes: &serde_json::Value,
475) -> Result<serde_json::Value, ProxyError> {
476    let client = Client::builder()
477        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
478        .build()?;
479
480    let url = format!("{}/skills/resolve", proxy_url.trim_end_matches('/'));
481
482    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
483        .json(scopes)
484        .send()
485        .await?;
486    let status = response.status();
487
488    if !status.is_success() {
489        let body = response.text().await.unwrap_or_else(|_| "empty".into());
490        return Err(ProxyError::ProxyResponse {
491            status: status.as_u16(),
492            body,
493        });
494    }
495
496    response
497        .json()
498        .await
499        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
500}
501
502/// Execute an LLM help query via the proxy server.
503pub async fn call_help(
504    proxy_url: &str,
505    query: &str,
506    tool: Option<&str>,
507) -> Result<String, ProxyError> {
508    let client = Client::builder()
509        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
510        .build()?;
511
512    let url = format!("{}/help", proxy_url.trim_end_matches('/'));
513
514    let payload = ProxyHelpRequest {
515        query: query.to_string(),
516        tool: tool.map(|t| t.to_string()),
517    };
518
519    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
520        .json(&payload)
521        .send()
522        .await?;
523    let status = response.status();
524
525    if !status.is_success() {
526        let body = response.text().await.unwrap_or_else(|_| "empty".into());
527        return Err(ProxyError::ProxyResponse {
528            status: status.as_u16(),
529            body,
530        });
531    }
532
533    let body: ProxyHelpResponse = response
534        .json()
535        .await
536        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))?;
537
538    if let Some(err) = body.error {
539        return Err(ProxyError::ProxyResponse {
540            status: 200,
541            body: err,
542        });
543    }
544
545    Ok(body.content)
546}
ati/proxy/client.rs

ati/proxy/
client.rs
