feat: add Mosh-style optimistic input prediction for remote terminals

crates/okena-core/src/client/connection.rs

@@ -25,7 +25,8 @@ pub trait ConnectionHandler: Send + Sync + 'static {
         ws_sender: async_channel::Sender<WsClientMessage>,
     );
     /// Binary PTY output arrived — route to the terminal's emulator.
-    fn on_terminal_output(&self, prefixed_id: &str, data: &[u8]);
+    /// `acked_input_seq` is the server's echo of the latest input sequence (v2 protocol).
+    fn on_terminal_output(&self, prefixed_id: &str, data: &[u8], acked_input_seq: Option<u64>);
     /// Terminal removed — clean up platform terminal object.
     fn remove_terminal(&self, prefixed_id: &str);
     /// Pre-resize a terminal's grid to match the server's dimensions.
@@ -614,17 +615,22 @@ impl<H: ConnectionHandler> RemoteClient<H> {
         let writer_handle = tokio::spawn(async move {
             while let Ok(msg) = ws_rx_clone.recv().await {
                 // For SendText, prefer binary frame when stream_id is known
-                if let WsClientMessage::SendText { terminal_id, text } = &msg {
+                if let WsClientMessage::SendText { terminal_id, text, input_seq } = &msg {
                     let stream_id = stream_map_for_writer
                         .read()
                         .ok()
                         .and_then(|m| m.get(terminal_id).copied());
                     if let Some(sid) = stream_id {
-                        let frame = crate::ws::build_binary_frame(
-                            crate::ws::FRAME_TYPE_INPUT,
-                            sid,
-                            text.as_bytes(),
-                        );
+                        // Use v2 frame if we have a non-zero input_seq (prediction active)
+                        let frame = if *input_seq > 0 {
+                            crate::ws::build_input_frame_v2(sid, *input_seq, text.as_bytes())
+                        } else {
+                            crate::ws::build_binary_frame(
+                                crate::ws::FRAME_TYPE_INPUT,
+                                sid,
+                                text.as_bytes(),
+                            )
+                        };
                         if let Err(e) = futures::SinkExt::send(
                             &mut ws_write,
                             tungstenite::Message::Binary(frame.into()),
@@ -639,7 +645,7 @@ impl<H: ConnectionHandler> RemoteClient<H> {
                 }
 
                 let json = match &msg {
-                    WsClientMessage::SendText { terminal_id, text } => {
+                    WsClientMessage::SendText { terminal_id, text, .. } => {
                         serde_json::json!({
                             "type": "send_text",
                             "terminal_id": terminal_id,
@@ -694,16 +700,22 @@ impl<H: ConnectionHandler> RemoteClient<H> {
         loop {
             match futures::StreamExt::next(&mut ws_read).await {
                 Some(Ok(tungstenite::Message::Binary(data))) => {
-                    // Generic binary frame: [proto:1][type:1][stream_id:4 BE][payload...]
-                    if let Some((frame_type, stream_id, payload)) =
-                        crate::ws::parse_binary_frame(&data)
+                    // Parse binary frame (v1 or v2)
+                    if let Some((frame_type, stream_id, payload, acked_seq)) =
+                        crate::ws::parse_binary_frame_any(&data)
                     {
                         match frame_type {
-                            crate::ws::FRAME_TYPE_PTY | crate::ws::FRAME_TYPE_SNAPSHOT => {
-                                // Route PTY output or snapshot to the correct terminal
+                            crate::ws::FRAME_TYPE_PTY => {
+                                if let Some(remote_tid) = reverse_stream_map.get(&stream_id) {
+                                    let prefixed = make_prefixed_id(&config_id, remote_tid);
+                                    handler_clone.on_terminal_output(&prefixed, payload, acked_seq);
+                                }
+                            }
+                            crate::ws::FRAME_TYPE_SNAPSHOT => {
+                                // Snapshots implicitly ack everything
                                 if let Some(remote_tid) = reverse_stream_map.get(&stream_id) {
                                     let prefixed = make_prefixed_id(&config_id, remote_tid);
-                                    handler_clone.on_terminal_output(&prefixed, payload);
+                                    handler_clone.on_terminal_output(&prefixed, payload, Some(u64::MAX));
                                 }
                             }
                             _ => {

crates/okena-core/src/client/types.rs

@@ -22,8 +22,9 @@ pub enum ConnectionStatus {
 /// Messages sent from the UI thread to the WebSocket writer task.
 #[derive(Debug)]
 pub enum WsClientMessage {
-    /// Send text input to a remote terminal
-    SendText { terminal_id: String, text: String },
+    /// Send text input to a remote terminal.
+    /// `input_seq` is the latest prediction sequence (0 if not predicting).
+    SendText { terminal_id: String, text: String, input_seq: u64 },
     /// Resize a remote terminal
     Resize {
         terminal_id: String,
@@ -111,6 +112,7 @@ mod tests {
         let msg = WsClientMessage::SendText {
             terminal_id: "t1".to_string(),
             text: "hello".to_string(),
+            input_seq: 0,
         };
         let debug = format!("{:?}", msg);
         assert!(debug.contains("SendText"));

crates/okena-core/src/ws.rs

@@ -64,6 +64,7 @@ pub enum WsOutbound {
 // ── Binary frame protocol ──────────────────────────────────────────────────
 
 pub const PROTO_VERSION: u8 = 1;
+pub const PROTO_VERSION_2: u8 = 2;
 pub const FRAME_TYPE_PTY: u8 = 1; // server → client: live PTY output
 pub const FRAME_TYPE_SNAPSHOT: u8 = 2; // server → client: full screen redraw
 pub const FRAME_TYPE_INPUT: u8 = 3; // client → server: terminal input
@@ -105,6 +106,83 @@ pub fn build_pty_frame(stream_id: u32, data: &[u8]) -> Vec<u8> {
     build_binary_frame(FRAME_TYPE_PTY, stream_id, data)
 }
 
+// ── V2 binary frame protocol (with input sequence numbers) ─────────────────
+
+/// Build a v2 INPUT frame: [proto=2][type=3][stream_id:4 BE][input_seq:8 BE][payload...]
+pub fn build_input_frame_v2(stream_id: u32, input_seq: u64, data: &[u8]) -> Vec<u8> {
+    let mut frame = Vec::with_capacity(14 + data.len());
+    frame.push(PROTO_VERSION_2);
+    frame.push(FRAME_TYPE_INPUT);
+    frame.extend_from_slice(&stream_id.to_be_bytes());
+    frame.extend_from_slice(&input_seq.to_be_bytes());
+    frame.extend_from_slice(data);
+    frame
+}
+
+/// Parse a v2 INPUT frame.
+/// Returns (stream_id, input_seq, payload) or None if invalid.
+pub fn parse_input_frame_v2(data: &[u8]) -> Option<(u32, u64, &[u8])> {
+    if data.len() < 14 || data[0] != PROTO_VERSION_2 || data[1] != FRAME_TYPE_INPUT {
+        return None;
+    }
+    let stream_id = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
+    let input_seq = u64::from_be_bytes([data[6], data[7], data[8], data[9], data[10], data[11], data[12], data[13]]);
+    Some((stream_id, input_seq, &data[14..]))
+}
+
+/// Build a v2 PTY frame: [proto=2][type=1][stream_id:4 BE][last_input_seq:8 BE][payload...]
+pub fn build_pty_frame_v2(stream_id: u32, last_input_seq: u64, data: &[u8]) -> Vec<u8> {
+    let mut frame = Vec::with_capacity(14 + data.len());
+    frame.push(PROTO_VERSION_2);
+    frame.push(FRAME_TYPE_PTY);
+    frame.extend_from_slice(&stream_id.to_be_bytes());
+    frame.extend_from_slice(&last_input_seq.to_be_bytes());
+    frame.extend_from_slice(data);
+    frame
+}
+
+/// Parse a v2 PTY frame.
+/// Returns (stream_id, last_input_seq, payload) or None if invalid.
+pub fn parse_pty_frame_v2(data: &[u8]) -> Option<(u32, u64, &[u8])> {
+    if data.len() < 14 || data[0] != PROTO_VERSION_2 || data[1] != FRAME_TYPE_PTY {
+        return None;
+    }
+    let stream_id = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
+    let last_input_seq = u64::from_be_bytes([data[6], data[7], data[8], data[9], data[10], data[11], data[12], data[13]]);
+    Some((stream_id, last_input_seq, &data[14..]))
+}
+
+/// Parse a binary frame that could be v1 or v2.
+/// Returns (frame_type, stream_id, payload, acked_input_seq).
+/// For v1 frames, acked_input_seq is None. For v2 PTY/INPUT frames, it's Some(seq).
+pub fn parse_binary_frame_any(data: &[u8]) -> Option<(u8, u32, &[u8], Option<u64>)> {
+    if data.len() < 6 {
+        return None;
+    }
+    let proto = data[0];
+    let frame_type = data[1];
+
+    if proto == PROTO_VERSION_2 && (frame_type == FRAME_TYPE_PTY || frame_type == FRAME_TYPE_INPUT) {
+        // v2 frame with input_seq
+        if data.len() < 14 {
+            return None;
+        }
+        let stream_id = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
+        let seq = u64::from_be_bytes([data[6], data[7], data[8], data[9], data[10], data[11], data[12], data[13]]);
+        Some((frame_type, stream_id, &data[14..], Some(seq)))
+    } else if proto == PROTO_VERSION {
+        // v1 frame
+        let stream_id = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
+        Some((frame_type, stream_id, &data[6..], None))
+    } else if proto == PROTO_VERSION_2 && frame_type == FRAME_TYPE_SNAPSHOT {
+        // v2 snapshot — same format as v1 (no seq)
+        let stream_id = u32::from_be_bytes([data[2], data[3], data[4], data[5]]);
+        Some((frame_type, stream_id, &data[6..], None))
+    } else {
+        None
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -271,4 +349,69 @@ mod tests {
             assert_eq!(parsed_data, payload.as_bytes());
         }
     }
+
+    // ── V2 frame tests ──────────────────────────────────────────────────
+
+    #[test]
+    fn v2_input_frame_round_trip() {
+        let frame = build_input_frame_v2(42, 12345, b"hello");
+        let (sid, seq, payload) = parse_input_frame_v2(&frame).unwrap();
+        assert_eq!(sid, 42);
+        assert_eq!(seq, 12345);
+        assert_eq!(payload, b"hello");
+    }
+
+    #[test]
+    fn v2_pty_frame_round_trip() {
+        let frame = build_pty_frame_v2(7, 99, b"output");
+        let (sid, seq, payload) = parse_pty_frame_v2(&frame).unwrap();
+        assert_eq!(sid, 7);
+        assert_eq!(seq, 99);
+        assert_eq!(payload, b"output");
+    }
+
+    #[test]
+    fn v1_frames_still_parse() {
+        let frame = build_binary_frame(FRAME_TYPE_PTY, 1, b"data");
+        let (ft, sid, payload) = parse_binary_frame(&frame).unwrap();
+        assert_eq!(ft, FRAME_TYPE_PTY);
+        assert_eq!(sid, 1);
+        assert_eq!(payload, b"data");
+    }
+
+    #[test]
+    fn v2_parser_rejects_short_frames() {
+        assert!(parse_input_frame_v2(&[2, 3, 0, 0, 0, 0, 0, 0]).is_none());
+        assert!(parse_pty_frame_v2(&[2, 1, 0, 0, 0]).is_none());
+    }
+
+    #[test]
+    fn parse_binary_frame_any_v1() {
+        let frame = build_binary_frame(FRAME_TYPE_PTY, 5, b"test");
+        let (ft, sid, payload, seq) = parse_binary_frame_any(&frame).unwrap();
+        assert_eq!(ft, FRAME_TYPE_PTY);
+        assert_eq!(sid, 5);
+        assert_eq!(payload, b"test");
+        assert!(seq.is_none());
+    }
+
+    #[test]
+    fn parse_binary_frame_any_v2_pty() {
+        let frame = build_pty_frame_v2(5, 100, b"test");
+        let (ft, sid, payload, seq) = parse_binary_frame_any(&frame).unwrap();
+        assert_eq!(ft, FRAME_TYPE_PTY);
+        assert_eq!(sid, 5);
+        assert_eq!(payload, b"test");
+        assert_eq!(seq, Some(100));
+    }
+
+    #[test]
+    fn parse_binary_frame_any_v2_input() {
+        let frame = build_input_frame_v2(3, 50, b"keys");
+        let (ft, sid, payload, seq) = parse_binary_frame_any(&frame).unwrap();
+        assert_eq!(ft, FRAME_TYPE_INPUT);
+        assert_eq!(sid, 3);
+        assert_eq!(payload, b"keys");
+        assert_eq!(seq, Some(50));
+    }
 }

mobile/native/src/api/state.rs

@@ -81,6 +81,7 @@ pub async fn send_special_key(
         WsClientMessage::SendText {
             terminal_id,
             text,
+            input_seq: 0,
         },
     );
     Ok(())

mobile/native/src/api/terminal.rs

@@ -169,6 +169,7 @@ pub async fn send_text(conn_id: String, terminal_id: String, text: String) -> an
         WsClientMessage::SendText {
             terminal_id,
             text,
+            input_seq: 0,
         },
     );
     Ok(())

mobile/native/src/client/handler.rs

@@ -47,7 +47,7 @@ impl ConnectionHandler for MobileConnectionHandler {
             .insert(prefixed_id.to_string(), holder);
     }
 
-    fn on_terminal_output(&self, prefixed_id: &str, data: &[u8]) {
+    fn on_terminal_output(&self, prefixed_id: &str, data: &[u8], _acked_input_seq: Option<u64>) {
         *self.last_activity.lock() = Instant::now();
         if let Some(holder) = self.terminals.read().get(prefixed_id) {
             holder.process_output(data);
@@ -120,7 +120,7 @@ mod tests {
         let (tx, _rx) = async_channel::bounded(1);
 
         handler.create_terminal("conn1", "t1", "remote:conn1:t1", tx);
-        handler.on_terminal_output("remote:conn1:t1", b"hello");
+        handler.on_terminal_output("remote:conn1:t1", b"hello", None);
 
         let terminals = handler.terminals().read();
         let holder = terminals.get("remote:conn1:t1").unwrap();

src/action_dispatch.rs

@@ -119,14 +119,36 @@ impl ActionDispatcher {
                         });
                         return;
                     }
+                    ActionRequest::CreateTerminal { project_id } => {
+                        // Optimistically set focus for the expected new terminal path.
+                        // add_terminal() wraps root in Split{[old, new]}, so new terminal
+                        // will be at path [1]. If no layout existed, it's at root [].
+                        let pid = project_id.clone();
+                        workspace.update(cx, |ws, cx| {
+                            let path = if ws.project(&pid)
+                                .and_then(|p| p.layout.as_ref())
+                                .is_some()
+                            {
+                                vec![1]
+                            } else {
+                                vec![]
+                            };
+                            ws.set_focused_terminal(pid, path, cx);
+                        });
+                        // Don't return — action proceeds to be sent to server below
+                    }
                     _ => {}
                 }
 
+                log::info!("[dispatch] sending remote action to server");
                 let action = strip_remote_ids(action, connection_id);
                 let cid = connection_id.clone();
                 manager.update(cx, |rm, cx| {
+                    log::info!("[dispatch] inside manager.update, calling send_action");
                     rm.send_action(&cid, action, cx);
+                    log::info!("[dispatch] send_action returned");
                 });
+                log::info!("[dispatch] manager.update completed");
             }
         }
     }

src/app/headless.rs

@@ -254,14 +254,24 @@ impl HeadlessApp {
                     Err(_) => break,
                 };
 
+                log::info!("[headless] received bridge message: {:?}", msg.command);
                 let result = match msg.command {
                     RemoteCommand::Action(action) => {
-                        cx.update(|cx| {
-                            workspace.update(cx, |ws, cx| {
-                                execute_action(action, ws, &*backend, &terminals, cx)
-                                    .into_command_result()
-                            })
-                        })
+                        log::info!("[headless] executing action: {:?}", action);
+                        let r = cx.update(|cx| {
+                            log::info!("[headless] inside cx.update");
+                            let r = workspace.update(cx, |ws, cx| {
+                                log::info!("[headless] inside workspace.update, calling execute_action");
+                                let r = execute_action(action, ws, &*backend, &terminals, cx)
+                                    .into_command_result();
+                                log::info!("[headless] execute_action returned");
+                                r
+                            });
+                            log::info!("[headless] workspace.update returned");
+                            r
+                        });
+                        log::info!("[headless] cx.update returned");
+                        r
                     }
                     RemoteCommand::GetState => {
                         cx.update(|cx| {