exeora_cli/
connection.rs

1use crate::{
2    CLI_VERSION,
3    api::ApiClient,
4    auth::AuthManager,
5    config::{ConfigStore, ProjectEntry},
6    error::{ErrorCode, ExeoraError},
7    policy::{CommandPolicy, effective_policy, policy_allows},
8    protocol::{
9        HEARTBEAT_INTERVAL_MS, HEARTBEAT_REQUEST, HEARTBEAT_TIMEOUT_MS, MAX_RESULT_BYTES,
10        PRESENCE_SIGNAL_INTERVAL_MS, PROTOCOL_VERSION, ToolName, now_ms,
11    },
12    tools::ToolEngine,
13};
14use anyhow::{Context, Result, anyhow};
15use futures_util::{SinkExt, StreamExt};
16use serde_json::{Value, json};
17use std::{collections::HashMap, sync::Arc, time::Duration};
18use tokio::sync::{Mutex, mpsc};
19use tokio_tungstenite::{
20    connect_async,
21    tungstenite::{Message, client::IntoClientRequest, http::HeaderValue},
22};
23use tokio_util::sync::CancellationToken;
24use url::Url;
25
26type InFlight = Arc<Mutex<HashMap<String, CancellationToken>>>;
27
28pub async fn connect_forever(
29    config: &ConfigStore,
30    _api: &ApiClient,
31    auth: Arc<AuthManager>,
32    device_id: String,
33    projects: Vec<ProjectEntry>,
34    json_output: bool,
35) -> Result<()> {
36    let engine = Arc::new(ToolEngine::new()?);
37    let project_map: Arc<HashMap<String, ProjectEntry>> = Arc::new(
38        projects
39            .iter()
40            .cloned()
41            .map(|project| (project.id.clone(), project))
42            .collect(),
43    );
44    let gateway = config.gateway_url();
45    let mut delay = Duration::from_secs(1);
46    let stop = CancellationToken::new();
47    let signal_stop = stop.clone();
48    tokio::spawn(async move {
49        let _ = tokio::signal::ctrl_c().await;
50        signal_stop.cancel();
51    });
52
53    loop {
54        if stop.is_cancelled() {
55            break;
56        }
57        match connect_once(
58            &gateway,
59            &device_id,
60            &projects,
61            project_map.clone(),
62            auth.clone(),
63            engine.clone(),
64            stop.clone(),
65            json_output,
66        )
67        .await
68        {
69            Ok(ConnectOutcome::Stopped) => break,
70            Ok(ConnectOutcome::Rejected(reason)) => return Err(anyhow!(reason)),
71            Ok(ConnectOutcome::Disconnected) => {
72                delay = Duration::from_secs(1);
73                emit_event(
74                    json_output,
75                    "close",
76                    json!({ "reason": format!("Disconnected. Reconnecting in {}s.", delay.as_secs()) }),
77                );
78                tokio::select! { _ = tokio::time::sleep(delay) => {}, _ = stop.cancelled() => break }
79            }
80            Err(error) => {
81                emit_event(
82                    json_output,
83                    "close",
84                    json!({ "reason": format!("{error}. Reconnecting in {}s.", delay.as_secs()) }),
85                );
86                tokio::select! { _ = tokio::time::sleep(delay) => {}, _ = stop.cancelled() => break }
87                delay = (delay * 2).min(Duration::from_secs(30));
88            }
89        }
90    }
91    engine.kill_all().await;
92    if !json_output {
93        println!("Disconnected.");
94    }
95    Ok(())
96}
97
98enum ConnectOutcome {
99    Stopped,
100    Rejected(String),
101    Disconnected,
102}
103
104#[allow(clippy::too_many_arguments)]
105async fn connect_once(
106    gateway: &str,
107    device_id: &str,
108    projects: &[ProjectEntry],
109    project_map: Arc<HashMap<String, ProjectEntry>>,
110    auth: Arc<AuthManager>,
111    engine: Arc<ToolEngine>,
112    stop: CancellationToken,
113    json_output: bool,
114) -> Result<ConnectOutcome> {
115    let token = auth.access_token().await?;
116    let mut url = Url::parse(gateway)?.join(&format!("/api/relay/{device_id}"))?;
117    url.set_scheme(if url.scheme() == "https" { "wss" } else { "ws" })
118        .map_err(|_| anyhow!("invalid relay URL"))?;
119    let mut request = url.as_str().into_client_request()?;
120    request.headers_mut().insert(
121        "authorization",
122        HeaderValue::from_str(&format!("Bearer {token}"))?,
123    );
124    let (mut socket, _) = connect_async(request)
125        .await
126        .context("Could not connect to the Exeora relay")?;
127    let can_prompt = !json_output
128        && std::io::IsTerminal::is_terminal(&std::io::stdin())
129        && std::io::IsTerminal::is_terminal(&std::io::stdout());
130    socket.send(Message::Text(serde_json::to_string(&json!({
131        "type": "hello", "protocolVersion": PROTOCOL_VERSION, "deviceId": device_id,
132        "cliVersion": CLI_VERSION, "platform": platform(),
133        "projects": projects.iter().map(|project| json!({ "id": project.id, "slug": project.slug })).collect::<Vec<_>>(),
134        "capabilities": { "prompt": can_prompt, "tools": ToolName::ALL.iter().map(ToString::to_string).collect::<Vec<_>>() },
135    }))?.into())).await?;
136    emit_event(json_output, "open", json!({}));
137    if !json_output {
138        println!("✓ Connected. Waiting for tool calls.");
139    }
140
141    let (out_tx, mut out_rx) = mpsc::unbounded_channel::<Value>();
142    let in_flight: InFlight = Arc::new(Mutex::new(HashMap::new()));
143    let mut tick = tokio::time::interval(Duration::from_millis(HEARTBEAT_INTERVAL_MS));
144    tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
145    let mut heartbeat_auto = false;
146    let mut last_ack = now_ms();
147    let mut last_presence = now_ms();
148
149    loop {
150        tokio::select! {
151            _ = stop.cancelled() => {
152                let _ = socket.close(None).await;
153                cancel_all(&in_flight).await;
154                engine.kill_all().await;
155                return Ok(ConnectOutcome::Stopped);
156            }
157            _ = tick.tick() => {
158                let now = now_ms();
159                if heartbeat_auto && now.saturating_sub(last_ack) > HEARTBEAT_TIMEOUT_MS {
160                    let _ = socket.close(None).await;
161                    break;
162                }
163                let frame = if heartbeat_auto { HEARTBEAT_REQUEST.to_owned() } else { json!({ "type": "heartbeat", "at": now }).to_string() };
164                socket.send(Message::Text(frame.into())).await?;
165                if heartbeat_auto && now.saturating_sub(last_presence) >= PRESENCE_SIGNAL_INTERVAL_MS {
166                    socket.send(Message::Text(json!({ "type": "presence", "at": now }).to_string().into())).await?;
167                    last_presence = now;
168                }
169            }
170            Some(outgoing) = out_rx.recv() => {
171                socket.send(Message::Text(outgoing.to_string().into())).await?;
172            }
173            incoming = socket.next() => {
174                let Some(incoming) = incoming else { break; };
175                match incoming? {
176                    Message::Text(text) => {
177                        let Ok(message) = serde_json::from_str::<Value>(&text) else { continue; };
178                        match message.get("type").and_then(Value::as_str) {
179                            Some("heartbeat.ack") => last_ack = now_ms(),
180                            Some("hello.ack") => {
181                                heartbeat_auto = message.get("heartbeatMode").and_then(Value::as_str) == Some("auto");
182                                last_ack = now_ms();
183                                if let Some(latest) = message.get("latestCliVersion").and_then(Value::as_str)
184                                    && is_outdated(CLI_VERSION, latest) {
185                                        let notice = format!("A newer Exeora CLI is available ({CLI_VERSION} → {latest}). Run `exeora upgrade`.");
186                                        emit_event(json_output, "notice", json!({ "message": notice }));
187                                        if !json_output { println!("{notice}"); }
188                                }
189                            }
190                            Some("cancel") => {
191                                if let Some(id) = message.get("requestId").and_then(Value::as_str)
192                                    && let Some(token) = in_flight.lock().await.get(id) {
193                                    token.cancel();
194                                }
195                            }
196                            Some("approval.request") => {
197                                let tx = out_tx.clone();
198                                tokio::spawn(handle_approval(message, tx, can_prompt, json_output));
199                            }
200                            Some("approval.resolved") => {}
201                            Some("shutdown") => {
202                                let reason = message.get("reason").and_then(Value::as_str).unwrap_or("The gateway closed the connection.");
203                                cancel_all(&in_flight).await;
204                                engine.kill_all().await;
205                                return Ok(ConnectOutcome::Rejected(reason.to_owned()));
206                            }
207                            Some("tool.call") => {
208                                spawn_tool_call(message, project_map.clone(), engine.clone(), in_flight.clone(), out_tx.clone(), json_output).await;
209                            }
210                            _ => {}
211                        }
212                    }
213                    Message::Ping(data) => socket.send(Message::Pong(data)).await?,
214                    Message::Close(frame) => {
215                        if let Some(frame) = frame
216                            && frame.code == tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode::Policy {
217                                return Ok(ConnectOutcome::Rejected(frame.reason.to_string()));
218                        }
219                        break;
220                    }
221                    _ => {}
222                }
223            }
224        }
225    }
226    cancel_all(&in_flight).await;
227    engine.kill_all().await;
228    Ok(ConnectOutcome::Disconnected)
229}
230
231async fn spawn_tool_call(
232    message: Value,
233    projects: Arc<HashMap<String, ProjectEntry>>,
234    engine: Arc<ToolEngine>,
235    in_flight: InFlight,
236    outgoing: mpsc::UnboundedSender<Value>,
237    json_output: bool,
238) {
239    let Some(request_id) = message
240        .get("requestId")
241        .and_then(Value::as_str)
242        .map(str::to_owned)
243    else {
244        return;
245    };
246    let Some(project_id) = message.get("projectId").and_then(Value::as_str) else {
247        return;
248    };
249    let started = now_ms();
250    let send_error = |code: ErrorCode, text: &str| {
251        let _ = outgoing.send(result_frame(
252            &request_id,
253            started,
254            Err(ExeoraError::new(code, text)),
255        ));
256    };
257    if message
258        .get("expiresAt")
259        .and_then(Value::as_u64)
260        .is_some_and(|expires| now_ms() > expires)
261    {
262        send_error(
263            ErrorCode::ToolTimeout,
264            "The request expired before it was received.",
265        );
266        return;
267    }
268    let Some(project) = projects.get(project_id).cloned() else {
269        send_error(
270            ErrorCode::UnknownProject,
271            "This machine does not serve that project. Run `exeora project add` there.",
272        );
273        return;
274    };
275    let Some(tool_name) = message
276        .get("tool")
277        .and_then(Value::as_str)
278        .map(str::to_owned)
279    else {
280        send_error(ErrorCode::UnknownTool, "Unsupported tool.");
281        return;
282    };
283    let Ok(tool) = tool_name.parse::<ToolName>() else {
284        send_error(ErrorCode::UnknownTool, "Unsupported tool.");
285        return;
286    };
287    let arguments = message
288        .get("arguments")
289        .cloned()
290        .unwrap_or_else(|| json!({}));
291    let remote = message
292        .get("policy")
293        .cloned()
294        .and_then(|value| serde_json::from_value::<CommandPolicy>(value).ok());
295    let (policy, problem) = effective_policy(&project.root, remote);
296    if let Some(problem) = problem {
297        emit_event(json_output, "error", json!({ "message": problem }));
298    }
299    let verdict = policy_allows(&policy, tool, &arguments);
300    if !verdict.allowed {
301        send_error(
302            ErrorCode::Forbidden,
303            verdict
304                .reason
305                .as_deref()
306                .unwrap_or("This project does not allow that."),
307        );
308        return;
309    }
310
311    let cancel = CancellationToken::new();
312    in_flight
313        .lock()
314        .await
315        .insert(request_id.clone(), cancel.clone());
316    emit_event(
317        json_output,
318        "call",
319        json!({ "tool": tool_name, "project": project.slug, "client": describe_client(message.get("client")) }),
320    );
321    if !json_output {
322        println!("→ {tool_name} ({})", project.slug);
323    }
324    tokio::spawn(async move {
325        let result = engine.execute(&project.root, tool, arguments, cancel).await;
326        in_flight.lock().await.remove(&request_id);
327        let elapsed = now_ms().saturating_sub(started);
328        let frame = result_frame(&request_id, started, result);
329        let ok = frame
330            .pointer("/result/ok")
331            .and_then(Value::as_bool)
332            .unwrap_or(false);
333        let _ = outgoing.send(frame);
334        emit_event(
335            json_output,
336            "result",
337            json!({ "tool": tool_name, "ok": ok, "durationMs": elapsed }),
338        );
339        if !json_output {
340            println!("{} {tool_name} {elapsed}ms", if ok { "✓" } else { "✗" });
341        }
342    });
343}
344
345async fn handle_approval(
346    message: Value,
347    outgoing: mpsc::UnboundedSender<Value>,
348    can_prompt: bool,
349    json_output: bool,
350) {
351    let Some(id) = message.get("id").and_then(Value::as_str).map(str::to_owned) else {
352        return;
353    };
354    if !can_prompt {
355        let _ = outgoing.send(json!({ "type": "approval.answer", "id": id, "approved": false }));
356        return;
357    }
358    let prompt = message
359        .get("prompt")
360        .and_then(Value::as_str)
361        .unwrap_or("Allow this tool call?")
362        .to_owned();
363    let approved = tokio::task::spawn_blocking(move || {
364        cliclack::confirm(prompt)
365            .initial_value(false)
366            .interact()
367            .unwrap_or(false)
368    })
369    .await
370    .unwrap_or(false);
371    emit_event(
372        json_output,
373        "approval",
374        json!({ "id": id, "approved": approved }),
375    );
376    let _ = outgoing.send(json!({ "type": "approval.answer", "id": id, "approved": approved }));
377}
378
379fn result_frame(request_id: &str, started: u64, result: Result<Value, ExeoraError>) -> Value {
380    let result = match result {
381        Ok(value)
382            if serde_json::to_vec(&value).is_ok_and(|bytes| bytes.len() <= MAX_RESULT_BYTES) =>
383        {
384            json!({ "ok": true, "value": value })
385        }
386        Ok(_) => json!({
387            "ok": false,
388            "error": {
389                "code": ErrorCode::ToolFailed.as_str(),
390                "message": format!("Tool result exceeded the {MAX_RESULT_BYTES}-byte protocol limit. Narrow the request and try again."),
391            }
392        }),
393        Err(error) => {
394            json!({ "ok": false, "error": { "code": error.code.as_str(), "message": error.message } })
395        }
396    };
397    json!({ "type": "tool.result", "requestId": request_id, "durationMs": now_ms().saturating_sub(started), "result": result })
398}
399
400async fn cancel_all(in_flight: &InFlight) {
401    let mut calls = in_flight.lock().await;
402    for token in calls.values() {
403        token.cancel();
404    }
405    calls.clear();
406}
407
408fn describe_client(value: Option<&Value>) -> Option<String> {
409    let value = value?;
410    match (
411        value.get("name").and_then(Value::as_str),
412        value.get("version").and_then(Value::as_str),
413    ) {
414        (Some(name), Some(version)) => Some(format!("{name} {version}")),
415        (Some(name), None) => Some(name.to_owned()),
416        (None, Some(version)) => Some(version.to_owned()),
417        _ => None,
418    }
419}
420
421fn emit_event(json_output: bool, event: &str, fields: Value) {
422    if !json_output {
423        return;
424    }
425    let mut value = json!({ "at": now_ms(), "event": event });
426    if let (Some(target), Some(source)) = (value.as_object_mut(), fields.as_object()) {
427        target.extend(source.clone());
428    }
429    println!("{value}");
430}
431
432fn is_outdated(current: &str, latest: &str) -> bool {
433    match (
434        semver::Version::parse(current),
435        semver::Version::parse(latest),
436    ) {
437        (Ok(current), Ok(latest)) => current < latest,
438        _ => false,
439    }
440}
441
442fn platform() -> &'static str {
443    if cfg!(target_os = "windows") {
444        "win32"
445    } else if cfg!(target_os = "macos") {
446        "darwin"
447    } else {
448        "linux"
449    }
450}
451
452#[cfg(test)]
453mod tests {
454    use super::result_frame;
455    use crate::protocol::MAX_RESULT_BYTES;
456    use serde_json::json;
457
458    #[test]
459    fn rejects_an_oversized_tool_result_before_it_reaches_the_socket() {
460        let frame = result_frame(
461            "req_test",
462            0,
463            Ok(json!({ "content": "x".repeat(MAX_RESULT_BYTES) })),
464        );
465        assert_eq!(frame["result"]["ok"], false);
466        assert_eq!(frame["result"]["error"]["code"], "TOOL_FAILED");
467    }
468}
exeora_cli/connection.rs

exeora_cli/
connection.rs
