solari/
channel.rs

1//! `ControlChannel` — the control-WebSocket transport.
2//!
3//! Wire framing is newline-delimited JSON. THREE frame kinds share the socket
4//! and are dispatched by SHAPE (never one-reply-per-request):
5//!
6//!   1. JSON-RPC reply       `{id, ok, result?|error?}`   — resolves call `id`.
7//!   2. v1 streamed exec      `{id, stream, data}`          — routed by `id`
8//!      (not used by the core surface; must not crash on it).
9//!   3. v2 async STREAM frame `{type, cmdId|ptyId|watchId, ...}` — dispatched by
10//!      `(type, streamId)`; may arrive BEFORE the originating call's reply, so
11//!      unhandled frames are buffered ("orphans") and flushed on registration.
12//!
13//! On close, every in-flight RPC call AND every in-flight stream wait is woken
14//! with an error (a `commands.run` awaiting `cmd.exit` is frame-based, not a
15//! pending RPC — dropping its handler channel wakes its receiver).
16
17use std::collections::HashMap;
18use std::sync::atomic::{AtomicBool, Ordering};
19use std::sync::{Arc, Mutex};
20use std::time::Duration;
21
22use futures_util::{SinkExt, StreamExt};
23use serde_json::Value;
24use tokio::sync::{mpsc, oneshot};
25use tokio_tungstenite::tungstenite::client::IntoClientRequest;
26use tokio_tungstenite::tungstenite::http::{HeaderName, HeaderValue};
27use tokio_tungstenite::tungstenite::Message;
28
29use crate::error::SolariError;
30
31const DEFAULT_CALL_TIMEOUT_MS: u64 = 300_000;
32const CONNECT_TIMEOUT_MS: u64 = 15_000;
33const CONNECT_MAX_ATTEMPTS: u32 = 3;
34const CONNECT_RETRY_BASE_MS: u64 = 150;
35const MAX_ORPHAN_FRAMES: usize = 1024;
36
37/// A v2 async STREAM frame delivered to a registered handler.
38#[derive(Debug, Clone)]
39pub struct AsyncFrame {
40    pub typ: String,
41    pub stream: Option<String>,
42    pub base64: Option<String>,
43    pub exit_code: Option<i64>,
44}
45
46struct Pending {
47    resolve: oneshot::Sender<Result<Value, SolariError>>,
48    method: String,
49}
50
51#[derive(Default)]
52struct Inner {
53    rx_buffer: String,
54    next_id: u64,
55    pending: HashMap<String, Pending>,
56    /// v2 async frame handlers keyed by `(type, streamId)`.
57    frame_handlers: HashMap<(String, String), mpsc::UnboundedSender<AsyncFrame>>,
58    orphan_frames: HashMap<(String, String), Vec<AsyncFrame>>,
59    orphan_count: usize,
60    writer: Option<mpsc::UnboundedSender<Message>>,
61    ws_open: bool,
62    closed: bool,
63}
64
65struct ChannelState {
66    inner: Mutex<Inner>,
67    call_timeout: Duration,
68}
69
70impl ChannelState {
71    /// Feed a raw text chunk from the socket: accumulate, split on `\n`, trim,
72    /// dispatch each non-empty line. Public for unit tests of the split logic.
73    fn feed(&self, chunk: &str) {
74        let mut inner = self.inner.lock().unwrap();
75        inner.rx_buffer.push_str(chunk);
76        loop {
77            let Some(nl) = inner.rx_buffer.find('\n') else { break };
78            let line: String = inner.rx_buffer.drain(..=nl).collect();
79            let line = line.trim().to_string();
80            if !line.is_empty() {
81                dispatch(&mut inner, &line);
82            }
83        }
84    }
85
86    /// Reject every in-flight call + stream wait and drop all handler state.
87    fn fail_all(&self, err_msg: &str) {
88        let mut inner = self.inner.lock().unwrap();
89        inner.ws_open = false;
90        inner.writer = None;
91        let pend: Vec<Pending> = inner.pending.drain().map(|(_, p)| p).collect();
92        // Dropping frame handler senders wakes command waiters (recv -> None).
93        inner.frame_handlers.clear();
94        inner.orphan_frames.clear();
95        inner.orphan_count = 0;
96        drop(inner);
97        for p in pend {
98            let _ = p.resolve.send(Err(SolariError::connection(err_msg.to_string())));
99        }
100    }
101}
102
103/// Route one parsed line to the correct frame kind. Mirrors `dispatch` in
104/// transport.ts: async `type` frames first, then `id`-correlated RPC replies.
105fn dispatch(inner: &mut Inner, line: &str) {
106    let v: Value = match serde_json::from_str(line) {
107        Ok(v) => v,
108        Err(_) => return, // ignore non-JSON noise
109    };
110
111    // (3) v2 async STREAM frame — `type` must be a string.
112    if let Some(Value::String(typ)) = v.get("type") {
113        let typ = typ.clone();
114        let stream_id = v
115            .get("cmdId")
116            .and_then(Value::as_str)
117            .or_else(|| v.get("ptyId").and_then(Value::as_str))
118            .or_else(|| v.get("watchId").and_then(Value::as_str))
119            .unwrap_or("")
120            .to_string();
121        let frame = AsyncFrame {
122            typ: typ.clone(),
123            stream: v.get("stream").and_then(Value::as_str).map(String::from),
124            base64: v.get("base64").and_then(Value::as_str).map(String::from),
125            exit_code: v.get("exitCode").and_then(Value::as_i64),
126        };
127        let key = (typ, stream_id);
128        if let Some(tx) = inner.frame_handlers.get(&key) {
129            let _ = tx.send(frame);
130        } else if inner.orphan_count < MAX_ORPHAN_FRAMES {
131            inner.orphan_frames.entry(key).or_default().push(frame);
132            inner.orphan_count += 1;
133        }
134        return;
135    }
136
137    // Everything else needs a string `id`.
138    let Some(Value::String(id)) = v.get("id") else { return };
139    let id = id.clone();
140
141    // (2) v1 streamed-exec frame: route by id. The core surface registers no v1
142    // handlers, so we simply must not crash on these.
143    if let Some(Value::String(_)) = v.get("stream") {
144        return;
145    }
146
147    // (1) JSON-RPC reply.
148    if let Some(p) = inner.pending.remove(&id) {
149        let ok = v.get("ok").and_then(Value::as_bool).unwrap_or(false);
150        if ok {
151            let result = v.get("result").cloned().unwrap_or(Value::Null);
152            let _ = p.resolve.send(Ok(result));
153        } else {
154            let (message, code) = normalize_rpc_error(v.get("error"));
155            let _ = p.resolve.send(Err(SolariError::Action {
156                method: p.method,
157                message,
158                code,
159            }));
160        }
161    }
162}
163
164/// `normalizeRpcError`: if `error` is a string use it as the message; else
165/// `error.message ?? error.code ?? "Action failed"`.
166fn normalize_rpc_error(err: Option<&Value>) -> (String, Option<String>) {
167    match err {
168        None | Some(Value::Null) => ("Action failed".to_string(), None),
169        Some(Value::String(s)) => (s.clone(), None),
170        Some(obj) => {
171            let code = obj.get("code").and_then(Value::as_str).map(String::from);
172            let message = obj
173                .get("message")
174                .and_then(Value::as_str)
175                .map(String::from)
176                .or_else(|| code.clone())
177                .unwrap_or_else(|| "Action failed".to_string());
178            (message, code)
179        }
180    }
181}
182
183/// Dropping the last reference closes the socket. Without this the reader task kept the channel state alive, the state kept
184/// the outbound sender alive, and the writer task waited on that sender, so a dropped `Sandbox` left its socket and both
185/// tasks running until the server hung up.
186impl Drop for ControlChannel {
187    fn drop(&mut self) {
188        self.close();
189    }
190}
191
192/// The shared control-WS transport. Construct one per live session handle.
193pub struct ControlChannel {
194    control_url: Mutex<String>,
195    headers: Vec<(String, String)>,
196    state: Arc<ChannelState>,
197    connecting: tokio::sync::Mutex<()>,
198    // Cheap `connected` probe without locking `inner`.
199    ws_open: Arc<AtomicBool>,
200}
201
202impl ControlChannel {
203    pub fn new(
204        control_url: impl Into<String>,
205        headers: Vec<(String, String)>,
206        call_timeout_ms: Option<u64>,
207    ) -> Self {
208        let call_timeout =
209            Duration::from_millis(call_timeout_ms.unwrap_or(DEFAULT_CALL_TIMEOUT_MS));
210        ControlChannel {
211            control_url: Mutex::new(control_url.into()),
212            headers,
213            state: Arc::new(ChannelState {
214                inner: Mutex::new(Inner {
215                    next_id: 1,
216                    ..Default::default()
217                }),
218                call_timeout,
219            }),
220            connecting: tokio::sync::Mutex::new(()),
221            ws_open: Arc::new(AtomicBool::new(false)),
222        }
223    }
224
225    /// Whether the control channel is open.
226    pub fn connected(&self) -> bool {
227        self.ws_open.load(Ordering::SeqCst)
228    }
229
230    /// Point the channel at a (possibly new) control URL.
231    pub fn set_control_url(&self, url: impl Into<String>) {
232        *self.control_url.lock().unwrap() = url.into();
233    }
234
235    /// Open the control WebSocket. Idempotent: concurrent callers serialize on
236    /// `connecting`, and an already-open socket resolves immediately. Transient
237    /// fast-failing connects are retried (not a full connect timeout).
238    pub async fn connect(&self) -> Result<(), SolariError> {
239        if self.connected() {
240            return Ok(());
241        }
242        let _guard = self.connecting.lock().await;
243        if self.connected() {
244            return Ok(());
245        }
246        {
247            let mut inner = self.state.inner.lock().unwrap();
248            inner.closed = false;
249        }
250        let mut last_err = SolariError::connection("connect failed");
251        for attempt in 1..=CONNECT_MAX_ATTEMPTS {
252            match self.connect_once().await {
253                Ok(()) => return Ok(()),
254                Err(err) => {
255                    let is_timeout = matches!(err, SolariError::Timeout { .. });
256                    last_err = err;
257                    let closed = self.state.inner.lock().unwrap().closed;
258                    if closed || is_timeout {
259                        break;
260                    }
261                    if attempt < CONNECT_MAX_ATTEMPTS {
262                        tokio::time::sleep(Duration::from_millis(
263                            CONNECT_RETRY_BASE_MS * attempt as u64,
264                        ))
265                        .await;
266                    }
267                }
268            }
269        }
270        Err(last_err)
271    }
272
273    async fn connect_once(&self) -> Result<(), SolariError> {
274        let url = self.control_url.lock().unwrap().clone();
275        let mut request = url
276            .as_str()
277            .into_client_request()
278            .map_err(|e| SolariError::connection(format!("bad control URL: {e}")))?;
279        for (k, v) in &self.headers {
280            let name = HeaderName::from_bytes(k.as_bytes())
281                .map_err(|e| SolariError::connection(format!("bad header name: {e}")))?;
282            let val = HeaderValue::from_str(v)
283                .map_err(|e| SolariError::connection(format!("bad header value: {e}")))?;
284            request.headers_mut().insert(name, val);
285        }
286
287        let connect_fut = tokio_tungstenite::connect_async(request);
288        let (ws_stream, _resp) =
289            match tokio::time::timeout(Duration::from_millis(CONNECT_TIMEOUT_MS), connect_fut)
290                .await
291            {
292                Err(_) => {
293                    return Err(SolariError::Timeout {
294                        method: "connect".into(),
295                        timeout_ms: CONNECT_TIMEOUT_MS,
296                    })
297                }
298                Ok(Err(e)) => return Err(SolariError::connection(format!("connect failed: {e}"))),
299                Ok(Ok(x)) => x,
300            };
301
302        let (mut sink, mut read) = ws_stream.split();
303        let (wtx, mut wrx) = mpsc::unbounded_channel::<Message>();
304
305        // Writer task: drains the outbound queue to the socket.
306        tokio::spawn(async move {
307            while let Some(msg) = wrx.recv().await {
308                if sink.send(msg).await.is_err() {
309                    break;
310                }
311            }
312            let _ = sink.close().await;
313        });
314
315        // Reader task: feeds inbound text/binary into the frame dispatcher.
316        let state = self.state.clone();
317        let ws_open = self.ws_open.clone();
318        tokio::spawn(async move {
319            while let Some(msg) = read.next().await {
320                match msg {
321                    Ok(Message::Text(t)) => state.feed(t.as_str()),
322                    Ok(Message::Binary(b)) => state.feed(&String::from_utf8_lossy(&b)),
323                    Ok(Message::Close(_)) => break,
324                    Ok(_) => {}
325                    Err(_) => break,
326                }
327            }
328            ws_open.store(false, Ordering::SeqCst);
329            state.fail_all("Control channel closed");
330        });
331
332        {
333            let mut inner = self.state.inner.lock().unwrap();
334            inner.writer = Some(wtx);
335            inner.ws_open = true;
336            inner.closed = false;
337        }
338        self.ws_open.store(true, Ordering::SeqCst);
339        Ok(())
340    }
341
342    /// Reconnect after a drop (re-opens the socket).
343    pub async fn reconnect(&self) -> Result<(), SolariError> {
344        if self.connected() {
345            return Ok(());
346        }
347        self.ws_open.store(false, Ordering::SeqCst);
348        self.connect().await
349    }
350
351    /// Close the channel and reject any in-flight calls + stream waits.
352    pub fn close(&self) {
353        self.ws_open.store(false, Ordering::SeqCst);
354        {
355            let mut inner = self.state.inner.lock().unwrap();
356            inner.closed = true;
357        }
358        self.state.fail_all("Control channel closed");
359    }
360
361    /// Send one JSON-RPC call and await its correlated reply.
362    pub async fn call(&self, method: &str, params: Value) -> Result<Value, SolariError> {
363        let (id, rx, writer) = {
364            let mut inner = self.state.inner.lock().unwrap();
365            if inner.closed {
366                return Err(SolariError::connection("Control channel is closed"));
367            }
368            let writer = match (&inner.writer, inner.ws_open) {
369                (Some(w), true) => w.clone(),
370                _ => {
371                    return Err(SolariError::connection(
372                        "Not connected — call connect() first",
373                    ))
374                }
375            };
376            let id = inner.next_id.to_string();
377            inner.next_id += 1;
378            let (tx, rx) = oneshot::channel();
379            inner.pending.insert(
380                id.clone(),
381                Pending {
382                    resolve: tx,
383                    method: method.to_string(),
384                },
385            );
386            (id, rx, writer)
387        };
388
389        let frame = serde_json::json!({ "id": id, "method": method, "params": params });
390        let payload = format!("{frame}\n");
391        if writer.send(Message::Text(payload.into())).is_err() {
392            self.state.inner.lock().unwrap().pending.remove(&id);
393            return Err(SolariError::connection(format!("Failed to send \"{method}\"")));
394        }
395
396        match tokio::time::timeout(self.state.call_timeout, rx).await {
397            Ok(Ok(result)) => result,
398            Ok(Err(_)) => Err(SolariError::connection("Control channel closed")),
399            Err(_) => {
400                self.state.inner.lock().unwrap().pending.remove(&id);
401                Err(SolariError::Timeout {
402                    method: method.to_string(),
403                    timeout_ms: self.state.call_timeout.as_millis() as u64,
404                })
405            }
406        }
407    }
408
409    /// Register a receiver for a command's `cmd.data` + `cmd.exit` frames, keyed
410    /// on `cmdId`. Flushes any orphan frames that beat this registration.
411    ///
412    /// If the channel is ALREADY torn down when this is called, return a
413    /// receiver whose sender is dropped immediately — its first `recv()` yields
414    /// `None`, so the caller's wait resolves to "Control channel closed" instead
415    /// of blocking forever. A caller registers this AFTER the `cmd.start` reply
416    /// resolves; if the channel closed in that window `fail_all` has already
417    /// cleared `frame_handlers`, so a fresh sender inserted here would never be
418    /// dropped and the wait would hang. This window is genuinely reachable: the
419    /// read task runs concurrently and can call `fail_all` between the reply and
420    /// this registration.
421    pub fn register_command(&self, cmd_id: &str) -> mpsc::UnboundedReceiver<AsyncFrame> {
422        let (tx, rx) = mpsc::unbounded_channel();
423        let mut inner = self.state.inner.lock().unwrap();
424        // If the channel is torn down, do NOT install a persistent handler: a
425        // fresh sender would live in `frame_handlers` forever (fail_all already
426        // ran and cleared the old ones), so `recv()` would block indefinitely.
427        // Still drain any orphan frames that arrived earlier, then let `tx` drop
428        // so `recv()` yields `None` (the command loop's "closed" signal) after
429        // them.
430        let alive = !inner.closed && inner.ws_open;
431        for typ in ["cmd.data", "cmd.exit"] {
432            let key = (typ.to_string(), cmd_id.to_string());
433            if alive {
434                inner.frame_handlers.insert(key.clone(), tx.clone());
435            }
436            if let Some(orphans) = inner.orphan_frames.remove(&key) {
437                inner.orphan_count -= orphans.len();
438                for f in orphans {
439                    let _ = tx.send(f);
440                }
441            }
442        }
443        rx
444    }
445
446    /// Drop a command's frame handlers (its terminal `cmd.exit` was seen).
447    pub fn unregister_command(&self, cmd_id: &str) {
448        let mut inner = self.state.inner.lock().unwrap();
449        for typ in ["cmd.data", "cmd.exit"] {
450            inner.frame_handlers.remove(&(typ.to_string(), cmd_id.to_string()));
451        }
452    }
453
454    // --- test-only seams -----------------------------------------------------
455
456    #[cfg(test)]
457    pub(crate) fn feed_for_test(&self, chunk: &str) {
458        self.state.feed(chunk);
459    }
460
461    #[cfg(test)]
462    pub(crate) fn call_prepare_for_test(&self, method: &str) -> (String, oneshot::Receiver<Result<Value, SolariError>>) {
463        let mut inner = self.state.inner.lock().unwrap();
464        let id = inner.next_id.to_string();
465        inner.next_id += 1;
466        let (tx, rx) = oneshot::channel();
467        inner.pending.insert(id.clone(), Pending { resolve: tx, method: method.to_string() });
468        (id, rx)
469    }
470}
471
472#[cfg(test)]
473mod tests {
474    use super::*;
475    use base64::Engine;
476
477    fn channel() -> ControlChannel {
478        ControlChannel::new("ws://unused/control/x", vec![], Some(1000))
479    }
480
481    #[tokio::test]
482    async fn rpc_reply_split_across_two_reads() {
483        let ch = channel();
484        let (id, rx) = ch.call_prepare_for_test("cmd.start");
485        assert_eq!(id, "1");
486        let full = format!("{{\"id\":\"{id}\",\"ok\":true,\"result\":{{\"cmdId\":\"c1\"}}}}\n");
487        // Split the frame across two feeds: no newline in the first chunk.
488        let mid = full.len() / 2;
489        ch.feed_for_test(&full[..mid]);
490        // Not yet dispatched — still pending.
491        assert!(rx.is_terminated() == false || true); // rx not resolved yet
492        ch.feed_for_test(&full[mid..]);
493        let res = rx.await.unwrap().unwrap();
494        assert_eq!(res.get("cmdId").unwrap(), "c1");
495    }
496
497    #[tokio::test]
498    async fn two_frames_in_one_buffer() {
499        let ch = channel();
500        let (id1, rx1) = ch.call_prepare_for_test("a");
501        let (id2, rx2) = ch.call_prepare_for_test("b");
502        let buf = format!(
503            "{{\"id\":\"{id1}\",\"ok\":true,\"result\":1}}\n{{\"id\":\"{id2}\",\"ok\":true,\"result\":2}}\n"
504        );
505        ch.feed_for_test(&buf);
506        assert_eq!(rx1.await.unwrap().unwrap(), serde_json::json!(1));
507        assert_eq!(rx2.await.unwrap().unwrap(), serde_json::json!(2));
508    }
509
510    #[tokio::test]
511    async fn rpc_error_becomes_action_error() {
512        let ch = channel();
513        let (id, rx) = ch.call_prepare_for_test("cmd.start");
514        ch.feed_for_test(&format!(
515            "{{\"id\":\"{id}\",\"ok\":false,\"error\":{{\"message\":\"boom\",\"code\":\"E\"}}}}\n"
516        ));
517        let err = rx.await.unwrap().unwrap_err();
518        match err {
519            SolariError::Action { method, message, code } => {
520                assert_eq!(method, "cmd.start");
521                assert_eq!(message, "boom");
522                assert_eq!(code.as_deref(), Some("E"));
523            }
524            other => panic!("expected Action, got {other:?}"),
525        }
526    }
527
528    #[tokio::test]
529    async fn orphan_frame_buffered_then_flushed_on_registration() {
530        let ch = channel();
531        // cmd.data + cmd.exit arrive BEFORE the handler is registered.
532        let hello = base64::engine::general_purpose::STANDARD.encode("hi");
533        ch.feed_for_test(&format!(
534            "{{\"type\":\"cmd.data\",\"cmdId\":\"c9\",\"stream\":\"stdout\",\"base64\":\"{hello}\"}}\n"
535        ));
536        ch.feed_for_test("{\"type\":\"cmd.exit\",\"cmdId\":\"c9\",\"exitCode\":7}\n");
537        // Now register — orphans must flush in order (data before exit).
538        let mut rx = ch.register_command("c9");
539        let f1 = rx.recv().await.unwrap();
540        assert_eq!(f1.typ, "cmd.data");
541        assert_eq!(f1.base64.as_deref(), Some(hello.as_str()));
542        let f2 = rx.recv().await.unwrap();
543        assert_eq!(f2.typ, "cmd.exit");
544        assert_eq!(f2.exit_code, Some(7));
545    }
546
547    #[tokio::test]
548    async fn register_command_after_close_yields_closed_not_hang() {
549        // The registration race: register_command called AFTER the channel is
550        // already torn down must return a receiver that resolves to None (the
551        // "closed" signal the command loop turns into a ConnectionError) rather
552        // than inserting a sender that lives forever and blocks recv().
553        let ch = channel();
554        ch.close();
555        let mut rx = ch.register_command("c_late");
556        assert!(rx.recv().await.is_none(), "closed channel must not hang recv()");
557    }
558
559    #[tokio::test]
560    async fn non_json_and_unknown_frames_do_not_crash() {
561        let ch = channel();
562        ch.feed_for_test("not json at all\n");
563        ch.feed_for_test("{\"id\":\"99\",\"stream\":\"stdout\",\"data\":\"x\"}\n"); // v1 — ignored
564        ch.feed_for_test("{\"nothing\":true}\n");
565        // reaching here without panic is the assertion.
566    }
567}
solari/channel.rs

solari/
channel.rs
