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    /// Raw positional args for CLI tools. When present, the proxy's
32    /// `args_as_positional()` uses these instead of parsing `args`.
33    /// This preserves bare positional words like `browse status` that
34    /// don't survive the `--key value` parse into the args map.
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub raw_args: Option<Vec<String>>,
37}
38
39/// Response payload from the proxy server.
40#[derive(Debug, Deserialize)]
41pub struct ProxyCallResponse {
42    pub result: Value,
43    #[serde(default)]
44    pub error: Option<String>,
45}
46
47/// Request payload for the proxy's /help endpoint.
48#[derive(Debug, Serialize)]
49pub struct ProxyHelpRequest {
50    pub query: String,
51    #[serde(skip_serializing_if = "Option::is_none")]
52    pub tool: Option<String>,
53}
54
55/// Response from the proxy's /help endpoint.
56#[derive(Debug, Deserialize)]
57pub struct ProxyHelpResponse {
58    pub content: String,
59    #[serde(default)]
60    pub error: Option<String>,
61}
62
63const PROXY_TIMEOUT_SECS: u64 = 120;
64
65/// Build an HTTP request builder with JWT Bearer auth from ATI_SESSION_TOKEN.
66fn build_proxy_request(
67    client: &Client,
68    method: reqwest::Method,
69    url: &str,
70) -> reqwest::RequestBuilder {
71    let mut req = client.request(method, url);
72    if let Ok(token) = std::env::var("ATI_SESSION_TOKEN") {
73        if !token.is_empty() {
74            req = req.header("Authorization", format!("Bearer {token}"));
75        }
76    }
77    req
78}
79
80/// Execute a tool call via the proxy server.
81///
82/// POST {proxy_url}/call with JSON body: { tool_name, args }
83/// Scopes are carried inside the JWT — not in the request body.
84///
85/// `args` carries key-value pairs for HTTP/MCP tools.
86/// `raw_args`, if provided, is sent as an array in the `args` field for CLI tools.
87pub async fn call_tool(
88    proxy_url: &str,
89    tool_name: &str,
90    args: &HashMap<String, Value>,
91    raw_args: Option<&[String]>,
92) -> Result<Value, ProxyError> {
93    let client = Client::builder()
94        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
95        .build()?;
96
97    let url = format!("{}/call", proxy_url.trim_end_matches('/'));
98
99    // Send both the parsed args map (for HTTP/MCP/OpenAPI tools) AND the raw
100    // positional args (for CLI tools). The proxy's CallRequest handler uses
101    // args_as_map() for HTTP tools and args_as_positional() for CLI tools.
102    // args_as_positional() checks `raw_args` first, so CLI tools always get
103    // their original positional args even when the map is empty.
104    let args_value = serde_json::to_value(args).unwrap_or(Value::Object(serde_json::Map::new()));
105    let raw_args_vec = raw_args.filter(|r| !r.is_empty()).map(|r| r.to_vec());
106
107    let payload = ProxyCallRequest {
108        tool_name: tool_name.to_string(),
109        args: args_value,
110        raw_args: raw_args_vec,
111    };
112
113    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
114        .json(&payload)
115        .send()
116        .await?;
117    let status = response.status();
118
119    if !status.is_success() {
120        let body = response.text().await.unwrap_or_else(|_| "empty".into());
121        return Err(ProxyError::ProxyResponse {
122            status: status.as_u16(),
123            body,
124        });
125    }
126
127    let body: ProxyCallResponse = response
128        .json()
129        .await
130        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))?;
131
132    if let Some(err) = body.error {
133        return Err(ProxyError::ProxyResponse {
134            status: 200,
135            body: err,
136        });
137    }
138
139    Ok(body.result)
140}
141
142/// List available tools from the proxy.
143pub async fn list_tools(proxy_url: &str, query_params: &str) -> Result<Value, ProxyError> {
144    let client = Client::builder()
145        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
146        .build()?;
147    let mut url = format!("{}/tools", proxy_url.trim_end_matches('/'));
148    if !query_params.is_empty() {
149        url.push('?');
150        url.push_str(query_params);
151    }
152    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
153        .send()
154        .await?;
155    let status = response.status();
156    if !status.is_success() {
157        let body = response.text().await.unwrap_or_default();
158        return Err(ProxyError::ProxyResponse {
159            status: status.as_u16(),
160            body,
161        });
162    }
163    Ok(response.json().await?)
164}
165
166/// Get detailed info about a specific tool from the proxy.
167pub async fn get_tool_info(proxy_url: &str, name: &str) -> Result<Value, ProxyError> {
168    let client = Client::builder()
169        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
170        .build()?;
171    let url = format!("{}/tools/{}", proxy_url.trim_end_matches('/'), name);
172    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
173        .send()
174        .await?;
175    let status = response.status();
176    if !status.is_success() {
177        let body = response.text().await.unwrap_or_default();
178        return Err(ProxyError::ProxyResponse {
179            status: status.as_u16(),
180            body,
181        });
182    }
183    Ok(response.json().await?)
184}
185
186/// Forward a raw MCP JSON-RPC message via the proxy's /mcp endpoint.
187pub async fn call_mcp(
188    proxy_url: &str,
189    method: &str,
190    params: Option<Value>,
191) -> Result<Value, ProxyError> {
192    use std::sync::atomic::{AtomicU64, Ordering};
193    static MCP_ID: AtomicU64 = AtomicU64::new(1);
194
195    let id = MCP_ID.fetch_add(1, Ordering::SeqCst);
196    let msg = serde_json::json!({
197        "jsonrpc": "2.0",
198        "id": id,
199        "method": method,
200        "params": params,
201    });
202
203    let client = Client::builder()
204        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
205        .build()?;
206
207    let url = format!("{}/mcp", proxy_url.trim_end_matches('/'));
208
209    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
210        .json(&msg)
211        .send()
212        .await?;
213    let status = response.status();
214
215    if status == reqwest::StatusCode::ACCEPTED {
216        return Ok(Value::Null);
217    }
218
219    if !status.is_success() {
220        let body = response.text().await.unwrap_or_else(|_| "empty".into());
221        return Err(ProxyError::ProxyResponse {
222            status: status.as_u16(),
223            body,
224        });
225    }
226
227    let body: Value = response
228        .json()
229        .await
230        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))?;
231
232    if let Some(err) = body.get("error") {
233        let message = err
234            .get("message")
235            .and_then(|m| m.as_str())
236            .unwrap_or("MCP proxy error");
237        return Err(ProxyError::ProxyResponse {
238            status: 200,
239            body: message.to_string(),
240        });
241    }
242
243    Ok(body.get("result").cloned().unwrap_or(Value::Null))
244}
245
246/// Fetch skill list from the proxy server.
247pub async fn list_skills(
248    proxy_url: &str,
249    query_params: &str,
250) -> Result<serde_json::Value, ProxyError> {
251    let client = Client::builder()
252        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
253        .build()?;
254
255    let url = if query_params.is_empty() {
256        format!("{}/skills", proxy_url.trim_end_matches('/'))
257    } else {
258        format!("{}/skills?{query_params}", proxy_url.trim_end_matches('/'))
259    };
260
261    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
262        .send()
263        .await?;
264    let status = response.status();
265
266    if !status.is_success() {
267        let body = response.text().await.unwrap_or_else(|_| "empty".into());
268        return Err(ProxyError::ProxyResponse {
269            status: status.as_u16(),
270            body,
271        });
272    }
273
274    response
275        .json()
276        .await
277        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
278}
279
280/// Fetch a skill's detail from the proxy server.
281pub async fn get_skill(
282    proxy_url: &str,
283    name: &str,
284    query_params: &str,
285) -> Result<serde_json::Value, ProxyError> {
286    let client = Client::builder()
287        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
288        .build()?;
289
290    let url = if query_params.is_empty() {
291        format!("{}/skills/{name}", proxy_url.trim_end_matches('/'))
292    } else {
293        format!(
294            "{}/skills/{name}?{query_params}",
295            proxy_url.trim_end_matches('/')
296        )
297    };
298
299    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
300        .send()
301        .await?;
302    let status = response.status();
303
304    if !status.is_success() {
305        let body = response.text().await.unwrap_or_else(|_| "empty".into());
306        return Err(ProxyError::ProxyResponse {
307            status: status.as_u16(),
308            body,
309        });
310    }
311
312    response
313        .json()
314        .await
315        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
316}
317
318async fn get_proxy_json(proxy_url: &str, path: &str) -> Result<serde_json::Value, ProxyError> {
319    let client = Client::builder()
320        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
321        .build()?;
322
323    let url = format!(
324        "{}/{}",
325        proxy_url.trim_end_matches('/'),
326        path.trim_start_matches('/')
327    );
328
329    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
330        .send()
331        .await?;
332    let status = response.status();
333
334    if !status.is_success() {
335        let body = response.text().await.unwrap_or_else(|_| "empty".into());
336        return Err(ProxyError::ProxyResponse {
337            status: status.as_u16(),
338            body,
339        });
340    }
341
342    response
343        .json()
344        .await
345        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
346}
347
348async fn get_proxy_json_with_query(
349    proxy_url: &str,
350    path: &str,
351    query: &[(&str, String)],
352) -> Result<serde_json::Value, ProxyError> {
353    let client = Client::builder()
354        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
355        .build()?;
356
357    let mut url = format!(
358        "{}/{}",
359        proxy_url.trim_end_matches('/'),
360        path.trim_start_matches('/')
361    );
362
363    if !query.is_empty() {
364        let params = query
365            .iter()
366            .map(|(key, value)| format!("{key}={}", urlencoding(value)))
367            .collect::<Vec<_>>()
368            .join("&");
369        url.push('?');
370        url.push_str(&params);
371    }
372
373    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
374        .send()
375        .await?;
376    let status = response.status();
377
378    if !status.is_success() {
379        let body = response.text().await.unwrap_or_else(|_| "empty".into());
380        return Err(ProxyError::ProxyResponse {
381            status: status.as_u16(),
382            body,
383        });
384    }
385
386    response
387        .json()
388        .await
389        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
390}
391
392/// List remote SkillATI skills from the proxy server.
393pub async fn get_skillati_catalog(
394    proxy_url: &str,
395    search: Option<&str>,
396) -> Result<serde_json::Value, ProxyError> {
397    let query = search
398        .map(|value| vec![("search", value.to_string())])
399        .unwrap_or_default();
400    get_proxy_json_with_query(proxy_url, "skillati/catalog", &query).await
401}
402
403/// Read a remote SkillATI skill from the proxy server.
404pub async fn get_skillati_read(
405    proxy_url: &str,
406    name: &str,
407) -> Result<serde_json::Value, ProxyError> {
408    get_proxy_json(proxy_url, &format!("skillati/{}", urlencoding(name))).await
409}
410
411/// List bundled resources for a remote SkillATI skill via the proxy server.
412pub async fn get_skillati_resources(
413    proxy_url: &str,
414    name: &str,
415    prefix: Option<&str>,
416) -> Result<serde_json::Value, ProxyError> {
417    let query = prefix
418        .map(|value| vec![("prefix", value.to_string())])
419        .unwrap_or_default();
420    get_proxy_json_with_query(
421        proxy_url,
422        &format!("skillati/{}/resources", urlencoding(name)),
423        &query,
424    )
425    .await
426}
427
428/// Read one arbitrary skill-relative path from a remote SkillATI skill via the proxy server.
429pub async fn get_skillati_file(
430    proxy_url: &str,
431    name: &str,
432    path: &str,
433) -> Result<serde_json::Value, ProxyError> {
434    get_proxy_json_with_query(
435        proxy_url,
436        &format!("skillati/{}/file", urlencoding(name)),
437        &[("path", path.to_string())],
438    )
439    .await
440}
441
442/// List on-demand references for a remote SkillATI skill via the proxy server.
443pub async fn get_skillati_refs(
444    proxy_url: &str,
445    name: &str,
446) -> Result<serde_json::Value, ProxyError> {
447    get_proxy_json(proxy_url, &format!("skillati/{}/refs", urlencoding(name))).await
448}
449
450/// Read one reference file from a remote SkillATI skill via the proxy server.
451pub async fn get_skillati_ref(
452    proxy_url: &str,
453    name: &str,
454    reference: &str,
455) -> Result<serde_json::Value, ProxyError> {
456    get_proxy_json(
457        proxy_url,
458        &format!(
459            "skillati/{}/ref/{}",
460            urlencoding(name),
461            urlencoding(reference)
462        ),
463    )
464    .await
465}
466
467fn urlencoding(s: &str) -> String {
468    s.replace('%', "%25")
469        .replace(' ', "%20")
470        .replace('#', "%23")
471        .replace('&', "%26")
472        .replace('?', "%3F")
473        .replace('/', "%2F")
474        .replace('=', "%3D")
475}
476
477/// Resolve skills for given scopes via the proxy.
478pub async fn resolve_skills(
479    proxy_url: &str,
480    scopes: &serde_json::Value,
481) -> Result<serde_json::Value, ProxyError> {
482    let client = Client::builder()
483        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
484        .build()?;
485
486    let url = format!("{}/skills/resolve", proxy_url.trim_end_matches('/'));
487
488    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
489        .json(scopes)
490        .send()
491        .await?;
492    let status = response.status();
493
494    if !status.is_success() {
495        let body = response.text().await.unwrap_or_else(|_| "empty".into());
496        return Err(ProxyError::ProxyResponse {
497            status: status.as_u16(),
498            body,
499        });
500    }
501
502    response
503        .json()
504        .await
505        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
506}
507
508/// Execute an LLM help query via the proxy server.
509pub async fn call_help(
510    proxy_url: &str,
511    query: &str,
512    tool: Option<&str>,
513) -> Result<String, ProxyError> {
514    let client = Client::builder()
515        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
516        .build()?;
517
518    let url = format!("{}/help", proxy_url.trim_end_matches('/'));
519
520    let payload = ProxyHelpRequest {
521        query: query.to_string(),
522        tool: tool.map(|t| t.to_string()),
523    };
524
525    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
526        .json(&payload)
527        .send()
528        .await?;
529    let status = response.status();
530
531    if !status.is_success() {
532        let body = response.text().await.unwrap_or_else(|_| "empty".into());
533        return Err(ProxyError::ProxyResponse {
534            status: status.as_u16(),
535            body,
536        });
537    }
538
539    let body: ProxyHelpResponse = response
540        .json()
541        .await
542        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))?;
543
544    if let Some(err) = body.error {
545        return Err(ProxyError::ProxyResponse {
546            status: 200,
547            body: err,
548        });
549    }
550
551    Ok(body.content)
552}
ati/proxy/client.rs

ati/proxy/
client.rs
