agent_tools_interface/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    pub args: HashMap<String, Value>,
30    #[serde(skip_serializing_if = "Option::is_none")]
31    pub raw_args: Option<Vec<String>>,
32}
33
34/// Response payload from the proxy server.
35#[derive(Debug, Deserialize)]
36pub struct ProxyCallResponse {
37    pub result: Value,
38    #[serde(default)]
39    pub error: Option<String>,
40}
41
42/// Request payload for the proxy's /help endpoint.
43#[derive(Debug, Serialize)]
44pub struct ProxyHelpRequest {
45    pub query: String,
46    #[serde(skip_serializing_if = "Option::is_none")]
47    pub tool: Option<String>,
48}
49
50/// Response from the proxy's /help endpoint.
51#[derive(Debug, Deserialize)]
52pub struct ProxyHelpResponse {
53    pub content: String,
54    #[serde(default)]
55    pub error: Option<String>,
56}
57
58const PROXY_TIMEOUT_SECS: u64 = 120;
59
60/// Build an HTTP request builder with JWT Bearer auth from ATI_SESSION_TOKEN.
61fn build_proxy_request(
62    client: &Client,
63    method: reqwest::Method,
64    url: &str,
65) -> reqwest::RequestBuilder {
66    let mut req = client.request(method, url);
67    if let Ok(token) = std::env::var("ATI_SESSION_TOKEN") {
68        if !token.is_empty() {
69            req = req.header("Authorization", format!("Bearer {token}"));
70        }
71    }
72    req
73}
74
75/// Execute a tool call via the proxy server.
76///
77/// POST {proxy_url}/call with JSON body: { tool_name, args }
78/// Scopes are carried inside the JWT — not in the request body.
79pub async fn call_tool(
80    proxy_url: &str,
81    tool_name: &str,
82    args: &HashMap<String, Value>,
83    raw_args: Option<&[String]>,
84) -> Result<Value, ProxyError> {
85    let client = Client::builder()
86        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
87        .build()?;
88
89    let url = format!("{}/call", proxy_url.trim_end_matches('/'));
90
91    let payload = ProxyCallRequest {
92        tool_name: tool_name.to_string(),
93        args: args.clone(),
94        raw_args: raw_args.map(|r| r.to_vec()),
95    };
96
97    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
98        .json(&payload)
99        .send()
100        .await?;
101    let status = response.status();
102
103    if !status.is_success() {
104        let body = response.text().await.unwrap_or_else(|_| "empty".into());
105        return Err(ProxyError::ProxyResponse {
106            status: status.as_u16(),
107            body,
108        });
109    }
110
111    let body: ProxyCallResponse = response
112        .json()
113        .await
114        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))?;
115
116    if let Some(err) = body.error {
117        return Err(ProxyError::ProxyResponse {
118            status: 200,
119            body: err,
120        });
121    }
122
123    Ok(body.result)
124}
125
126/// Forward a raw MCP JSON-RPC message via the proxy's /mcp endpoint.
127pub async fn call_mcp(
128    proxy_url: &str,
129    method: &str,
130    params: Option<Value>,
131) -> Result<Value, ProxyError> {
132    use std::sync::atomic::{AtomicU64, Ordering};
133    static MCP_ID: AtomicU64 = AtomicU64::new(1);
134
135    let id = MCP_ID.fetch_add(1, Ordering::SeqCst);
136    let msg = serde_json::json!({
137        "jsonrpc": "2.0",
138        "id": id,
139        "method": method,
140        "params": params,
141    });
142
143    let client = Client::builder()
144        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
145        .build()?;
146
147    let url = format!("{}/mcp", proxy_url.trim_end_matches('/'));
148
149    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
150        .json(&msg)
151        .send()
152        .await?;
153    let status = response.status();
154
155    if status == reqwest::StatusCode::ACCEPTED {
156        return Ok(Value::Null);
157    }
158
159    if !status.is_success() {
160        let body = response.text().await.unwrap_or_else(|_| "empty".into());
161        return Err(ProxyError::ProxyResponse {
162            status: status.as_u16(),
163            body,
164        });
165    }
166
167    let body: Value = response
168        .json()
169        .await
170        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))?;
171
172    if let Some(err) = body.get("error") {
173        let message = err
174            .get("message")
175            .and_then(|m| m.as_str())
176            .unwrap_or("MCP proxy error");
177        return Err(ProxyError::ProxyResponse {
178            status: 200,
179            body: message.to_string(),
180        });
181    }
182
183    Ok(body.get("result").cloned().unwrap_or(Value::Null))
184}
185
186/// Fetch skill list from the proxy server.
187pub async fn list_skills(
188    proxy_url: &str,
189    query_params: &str,
190) -> Result<serde_json::Value, ProxyError> {
191    let client = Client::builder()
192        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
193        .build()?;
194
195    let url = if query_params.is_empty() {
196        format!("{}/skills", proxy_url.trim_end_matches('/'))
197    } else {
198        format!("{}/skills?{query_params}", proxy_url.trim_end_matches('/'))
199    };
200
201    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
202        .send()
203        .await?;
204    let status = response.status();
205
206    if !status.is_success() {
207        let body = response.text().await.unwrap_or_else(|_| "empty".into());
208        return Err(ProxyError::ProxyResponse {
209            status: status.as_u16(),
210            body,
211        });
212    }
213
214    response
215        .json()
216        .await
217        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
218}
219
220/// Fetch a skill's detail from the proxy server.
221pub async fn get_skill(
222    proxy_url: &str,
223    name: &str,
224    query_params: &str,
225) -> Result<serde_json::Value, ProxyError> {
226    let client = Client::builder()
227        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
228        .build()?;
229
230    let url = if query_params.is_empty() {
231        format!("{}/skills/{name}", proxy_url.trim_end_matches('/'))
232    } else {
233        format!(
234            "{}/skills/{name}?{query_params}",
235            proxy_url.trim_end_matches('/')
236        )
237    };
238
239    let response = build_proxy_request(&client, reqwest::Method::GET, &url)
240        .send()
241        .await?;
242    let status = response.status();
243
244    if !status.is_success() {
245        let body = response.text().await.unwrap_or_else(|_| "empty".into());
246        return Err(ProxyError::ProxyResponse {
247            status: status.as_u16(),
248            body,
249        });
250    }
251
252    response
253        .json()
254        .await
255        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
256}
257
258/// Resolve skills for given scopes via the proxy.
259pub async fn resolve_skills(
260    proxy_url: &str,
261    scopes: &serde_json::Value,
262) -> Result<serde_json::Value, ProxyError> {
263    let client = Client::builder()
264        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
265        .build()?;
266
267    let url = format!("{}/skills/resolve", proxy_url.trim_end_matches('/'));
268
269    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
270        .json(scopes)
271        .send()
272        .await?;
273    let status = response.status();
274
275    if !status.is_success() {
276        let body = response.text().await.unwrap_or_else(|_| "empty".into());
277        return Err(ProxyError::ProxyResponse {
278            status: status.as_u16(),
279            body,
280        });
281    }
282
283    response
284        .json()
285        .await
286        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))
287}
288
289/// Execute an LLM help query via the proxy server.
290pub async fn call_help(
291    proxy_url: &str,
292    query: &str,
293    tool: Option<&str>,
294) -> Result<String, ProxyError> {
295    let client = Client::builder()
296        .timeout(Duration::from_secs(PROXY_TIMEOUT_SECS))
297        .build()?;
298
299    let url = format!("{}/help", proxy_url.trim_end_matches('/'));
300
301    let payload = ProxyHelpRequest {
302        query: query.to_string(),
303        tool: tool.map(|t| t.to_string()),
304    };
305
306    let response = build_proxy_request(&client, reqwest::Method::POST, &url)
307        .json(&payload)
308        .send()
309        .await?;
310    let status = response.status();
311
312    if !status.is_success() {
313        let body = response.text().await.unwrap_or_else(|_| "empty".into());
314        return Err(ProxyError::ProxyResponse {
315            status: status.as_u16(),
316            body,
317        });
318    }
319
320    let body: ProxyHelpResponse = response
321        .json()
322        .await
323        .map_err(|e| ProxyError::InvalidResponse(e.to_string()))?;
324
325    if let Some(err) = body.error {
326        return Err(ProxyError::ProxyResponse {
327            status: 200,
328            body: err,
329        });
330    }
331
332    Ok(body.content)
333}
agent_tools_interface/proxy/client.rs

agent_tools_interface/proxy/
client.rs
