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,
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 ok = result.is_ok();
328        let elapsed = now_ms().saturating_sub(started);
329        let _ = outgoing.send(result_frame(&request_id, started, result));
330        emit_event(
331            json_output,
332            "result",
333            json!({ "tool": tool_name, "ok": ok, "durationMs": elapsed }),
334        );
335        if !json_output {
336            println!("{} {tool_name} {elapsed}ms", if ok { "✓" } else { "✗" });
337        }
338    });
339}
340
341async fn handle_approval(
342    message: Value,
343    outgoing: mpsc::UnboundedSender<Value>,
344    can_prompt: bool,
345    json_output: bool,
346) {
347    let Some(id) = message.get("id").and_then(Value::as_str).map(str::to_owned) else {
348        return;
349    };
350    if !can_prompt {
351        let _ = outgoing.send(json!({ "type": "approval.answer", "id": id, "approved": false }));
352        return;
353    }
354    let prompt = message
355        .get("prompt")
356        .and_then(Value::as_str)
357        .unwrap_or("Allow this tool call?")
358        .to_owned();
359    let approved = tokio::task::spawn_blocking(move || {
360        cliclack::confirm(prompt)
361            .initial_value(false)
362            .interact()
363            .unwrap_or(false)
364    })
365    .await
366    .unwrap_or(false);
367    emit_event(
368        json_output,
369        "approval",
370        json!({ "id": id, "approved": approved }),
371    );
372    let _ = outgoing.send(json!({ "type": "approval.answer", "id": id, "approved": approved }));
373}
374
375fn result_frame(request_id: &str, started: u64, result: Result<Value, ExeoraError>) -> Value {
376    let result = match result {
377        Ok(value) => json!({ "ok": true, "value": value }),
378        Err(error) => {
379            json!({ "ok": false, "error": { "code": error.code.as_str(), "message": error.message } })
380        }
381    };
382    json!({ "type": "tool.result", "requestId": request_id, "durationMs": now_ms().saturating_sub(started), "result": result })
383}
384
385async fn cancel_all(in_flight: &InFlight) {
386    let mut calls = in_flight.lock().await;
387    for token in calls.values() {
388        token.cancel();
389    }
390    calls.clear();
391}
392
393fn describe_client(value: Option<&Value>) -> Option<String> {
394    let value = value?;
395    match (
396        value.get("name").and_then(Value::as_str),
397        value.get("version").and_then(Value::as_str),
398    ) {
399        (Some(name), Some(version)) => Some(format!("{name} {version}")),
400        (Some(name), None) => Some(name.to_owned()),
401        (None, Some(version)) => Some(version.to_owned()),
402        _ => None,
403    }
404}
405
406fn emit_event(json_output: bool, event: &str, fields: Value) {
407    if !json_output {
408        return;
409    }
410    let mut value = json!({ "at": now_ms(), "event": event });
411    if let (Some(target), Some(source)) = (value.as_object_mut(), fields.as_object()) {
412        target.extend(source.clone());
413    }
414    println!("{value}");
415}
416
417fn is_outdated(current: &str, latest: &str) -> bool {
418    match (
419        semver::Version::parse(current),
420        semver::Version::parse(latest),
421    ) {
422        (Ok(current), Ok(latest)) => current < latest,
423        _ => false,
424    }
425}
426
427fn platform() -> &'static str {
428    if cfg!(target_os = "windows") {
429        "win32"
430    } else if cfg!(target_os = "macos") {
431        "darwin"
432    } else {
433        "linux"
434    }
435}
exeora_cli/connection.rs

exeora_cli/
connection.rs
