Vectrum
Sign In
341 lines · 13.3 KB · Rust
//! Remote connections, like T3 Code's environments: one copy of the app can
//! let others on the network see and drive its projects and threads. The
//! host runs a small WebSocket server guarded by a key; a client connects
//! with the host's address and that key, and its threads show up in the
//! client's sidebar under the host's name.
//!
//! The protocol is JSON over WebSocket text messages. A client says
//! [`ClientMessage::Hello`] with the key, gets the host's projects and
//! threads in [`ServerMessage::State`], subscribes to the threads it opens,
//! and receives each change as a [`ServerMessage::Thread`] carrying the
//! transcript from the first item that changed. The host also shares its
//! environment settings in [`ServerMessage::Settings`], which clients can
//! change with [`ClientMessage::UpdateSettings`].

pub mod client;
pub mod server;

use std::path::PathBuf;

use serde::{Deserialize, Serialize};

use crate::git::TurnChanges;
use crate::models::Provider;
use crate::session::{QueuedMessage, SavedItem};
use crate::settings::SettingsPatch;
use crate::worktree::PullRequest;

/// The port a host listens on unless told otherwise.
pub const DEFAULT_PORT: u16 = 47_823;
/// Bumped when the messages change incompatibly.
pub const PROTOCOL: u32 = 1;

/// A thread as the sidebar needs it.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(default)]
pub struct ThreadSummary {
    pub id: String,
    pub project: PathBuf,
    pub title: String,
    pub provider: Provider,
    pub instance: String,
    pub model: Option<String>,
    pub permission_mode: String,
    pub reasoning: String,
    pub fast_mode: bool,
    /// `starting`, `idle`, `working` or `exited`.
    pub status: String,
    /// A tool call or question waits for the user.
    pub attention: bool,
    pub settled: bool,
    pub pinned: bool,
    /// Milliseconds since the epoch, like every time here.
    pub snoozed_until: Option<u64>,
    pub created: u64,
    pub updated: u64,
    pub last_message: u64,
    /// The thread's worktree branch, if it runs in one.
    pub branch: Option<String>,
    pub total_cost: f64,
    /// Whether anything was sent in it yet.
    pub has_items: bool,
    /// For a subagent's thread, the thread that delegated its task.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub parent: Option<String>,
}

/// The rest of an open thread's state, sent with its transcript.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(default)]
pub struct ThreadDetail {
    pub cwd: PathBuf,
    pub active_model: Option<String>,
    pub thinking: bool,
    pub preparing: bool,
    /// How long the running turn has been going.
    pub turn_elapsed_ms: Option<u64>,
    pub context_tokens: Option<u64>,
    pub context_window: Option<u64>,
    pub queue: Vec<QueuedMessage>,
    pub limit_resets: Option<u64>,
    pub resume_at: Option<u64>,
    /// Bound for a worktree made with the first message.
    pub pending_worktree: bool,
    pub pull_request: Option<PullRequest>,
    pub turn_changes: Vec<(usize, TurnChanges)>,
    pub background_tasks: Vec<RemoteTask>,
}

/// A command, monitor or subagent the host's agent left running.
#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(default)]
pub struct RemoteTask {
    pub task_id: String,
    pub description: String,
    pub elapsed_ms: u64,
}

#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize)]
#[serde(default)]
pub struct ProjectInfo {
    pub path: PathBuf,
    pub name: String,
    /// `None` outside a git repository.
    pub branch: Option<String>,
    pub branches: Vec<String>,
    pub git: bool,
    pub added: u64,
    pub removed: u64,
}

/// A file sent along with a message.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub struct Upload {
    pub name: String,
    /// Base64.
    pub data: String,
}

/// What a client asks of a thread on the host.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "action", rename_all = "snake_case")]
pub enum ThreadAction {
    Send { text: String, attachments: Vec<Upload> },
    Queue { text: String, attachments: Vec<Upload> },
    SteerQueued { ix: usize },
    RemoveQueued { ix: usize },
    Interrupt,
    /// `allow`, `always` or `deny`.
    Permission { ix: usize, decision: String },
    Answer { ix: usize, option: usize },
    /// Answers an AskUserQuestion form, one answer per question.
    AskSubmit { ix: usize, answers: Vec<String> },
    AskSkip { ix: usize },
    /// Sends a running command or subagent to the background.
    Background { ix: usize },
    StopTask { task_id: String },
    SetModel { model: Option<String> },
    SetPermissionMode { mode: String },
    SetReasoning { reasoning: String },
    SetFastMode { on: bool },
    Settle,
    Unsettle,
    Pin { pinned: bool },
    Snooze { until: Option<u64> },
    Rename { title: String },
    CancelResume,
    ScheduleResume { at: u64 },
    Delete,
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ClientMessage {
    Hello { key: String, name: String, protocol: u32 },
    Subscribe { thread: String },
    Unsubscribe { thread: String },
    Action { thread: String, action: ThreadAction },
    /// Starts a thread in one of the host's projects with its first message.
    Create {
        reference: u64,
        project: PathBuf,
        instance: String,
        provider: Provider,
        model: Option<String>,
        permission_mode: String,
        reasoning: String,
        fast_mode: bool,
        worktree: bool,
        /// The branch a worktree starts from.
        base: Option<String>,
        text: String,
        attachments: Vec<Upload>,
    },
    /// Changes the host's environment settings or its projects' overrides.
    UpdateSettings { patch: SettingsPatch },
}

#[derive(Clone, Debug, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ServerMessage {
    Welcome { host: String, version: String, protocol: u32 },
    Denied { reason: String },
    State { projects: Vec<ProjectInfo>, threads: Vec<ThreadSummary> },
    /// A subscribed thread changed: its state, and its transcript from item
    /// `from` on (earlier items are as last sent); `len` items in all.
    Thread { summary: ThreadSummary, detail: ThreadDetail, from: usize, len: usize, items: Vec<SavedItem> },
    Created { reference: u64, thread: String },
    Error { message: String },
    /// The host's environment settings and project overrides, sent on
    /// connecting and whenever they change. Older hosts don't send it.
    Settings { values: serde_json::Map<String, serde_json::Value> },
}

/// A new random key: 24 bytes, hex.
pub fn new_key() -> String {
    let mut bytes = [0u8; 24];
    let read = std::fs::File::open("/dev/urandom").and_then(|mut file| std::io::Read::read_exact(&mut file, &mut bytes));
    if read.is_err() {
        // Not as good, but never empty.
        let seed = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).unwrap_or_default().as_nanos();
        for (ix, byte) in bytes.iter_mut().enumerate() {
            *byte = (seed >> ((ix % 16) * 8)) as u8 ^ (ix as u8).wrapping_mul(151);
        }
    }
    bytes.iter().map(|byte| format!("{byte:02x}")).collect()
}

/// Compares keys in time independent of where they differ.
pub fn keys_match(a: &str, b: &str) -> bool {
    a.len() == b.len() && a.bytes().zip(b.bytes()).fold(0u8, |acc, (x, y)| acc | (x ^ y)) == 0
}

/// A pairing code bundles a host's address and key: `192.168.2.10:47823/<key>`.
pub fn pairing_code(address: &str, key: &str) -> String {
    format!("{address}/{key}")
}

/// Splits a pairing code, or an address alone, into address and key.
pub fn parse_pairing(code: &str) -> (String, Option<String>) {
    let code = code.trim().trim_start_matches("ws://").trim_start_matches("vectrum://");
    match code.split_once('/') {
        Some((address, key)) if !key.trim().is_empty() => (address.to_string(), Some(key.trim().to_string())),
        Some((address, _)) => (address.to_string(), None),
        None => (code.to_string(), None),
    }
}

/// `host` or `host:port`, with the default port filled in.
pub fn with_port(address: &str) -> String {
    let address = address.trim();
    let has_port = match address.rsplit_once(':') {
        // An IPv6 address without brackets has colons but no port.
        Some((host, port)) => !host.is_empty() && port.parse::<u16>().is_ok() && (!host.contains(':') || host.ends_with(']')),
        None => false,
    };
    if has_port { address.to_string() } else { format!("{address}:{DEFAULT_PORT}") }
}

/// This machine's IPv4 addresses on the network, for showing where to connect.
pub fn local_addresses() -> Vec<String> {
    let output = std::process::Command::new("hostname").arg("-I").stderr(std::process::Stdio::null()).output();
    let text = output.map(|output| String::from_utf8_lossy(&output.stdout).to_string()).unwrap_or_default();
    text.split_whitespace().filter(|address| address.parse::<std::net::Ipv4Addr>().is_ok()).map(str::to_string).collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[::core::prelude::v1::test]
    fn pairing_codes() {
        assert_eq!(parse_pairing("192.168.2.10:47823/abc"), ("192.168.2.10:47823".into(), Some("abc".into())));
        assert_eq!(parse_pairing(" 10.0.0.2 "), ("10.0.0.2".into(), None));
        assert_eq!(with_port("10.0.0.2"), format!("10.0.0.2:{DEFAULT_PORT}"));
        assert_eq!(with_port("10.0.0.2:9000"), "10.0.0.2:9000");
        assert_eq!(with_port("box.local"), format!("box.local:{DEFAULT_PORT}"));
        assert_eq!(with_port("[::1]:9000"), "[::1]:9000");
    }

    #[::core::prelude::v1::test]
    fn keys() {
        let key = new_key();
        assert_eq!(key.len(), 48);
        assert_ne!(key, new_key());
        assert!(keys_match(&key, &key.clone()));
        assert!(!keys_match(&key, &new_key()));
        assert!(!keys_match("a", "ab"));
    }

    #[::core::prelude::v1::test]
    fn messages_round_trip() {
        let message = ClientMessage::Action {
            thread: "t1".into(),
            action: ThreadAction::Send { text: "hi".into(), attachments: Vec::new() },
        };
        let text = serde_json::to_string(&message).unwrap();
        assert_eq!(serde_json::from_str::<ClientMessage>(&text).unwrap(), message);

        let patch: SettingsPatch = serde_json::from_str(r#"{"values": {"workspace": "worktree"}}"#).unwrap();
        let message = ClientMessage::UpdateSettings { patch };
        let text = serde_json::to_string(&message).unwrap();
        assert_eq!(serde_json::from_str::<ClientMessage>(&text).unwrap(), message);
        let values = crate::settings::Settings::default().environment_values();
        let text = serde_json::to_string(&ServerMessage::Settings { values: values.clone() }).unwrap();
        match serde_json::from_str::<ServerMessage>(&text).unwrap() {
            ServerMessage::Settings { values: received } => assert_eq!(received, values),
            _ => panic!("not settings"),
        }
    }

    #[::core::prelude::v1::test]
    fn client_and_server_talk() {
        use std::time::Duration;

        use super::client::{ClientEvent, ConnectionStatus, RemoteClient};
        use super::server::{RemoteServer, ServerEvent};

        let port = 47_900 + (std::process::id() % 90) as u16;
        let key = new_key();
        let server = RemoteServer::start(port, key.clone(), "test-host".into()).unwrap();
        fn wait<T>(receiver: &smol::channel::Receiver<T>) -> Option<T> {
            smol::block_on(async {
                smol::future::or(async { receiver.recv().await.ok() }, async {
                    smol::Timer::after(Duration::from_secs(10)).await;
                    None
                })
                .await
            })
        }

        let client = RemoteClient::connect(&format!("127.0.0.1:{port}"), key, "test-client".into());
        assert!(matches!(wait(&client.events), Some(ClientEvent::Status(ConnectionStatus::Connecting))));
        match wait(&client.events) {
            Some(ClientEvent::Status(ConnectionStatus::Connected(host))) => assert_eq!(host, "test-host"),
            _ => panic!("didn't connect"),
        }
        let conn = match wait(&server.events) {
            Some(ServerEvent::Connected { conn, name, .. }) => {
                assert_eq!(name, "test-client");
                conn
            }
            _ => panic!("server didn't see the client"),
        };
        client.send(&ClientMessage::Subscribe { thread: "t1".into() });
        match wait(&server.events) {
            Some(ServerEvent::Message { message: ClientMessage::Subscribe { thread }, .. }) => assert_eq!(thread, "t1"),
            _ => panic!("server didn't get the message"),
        }
        server.send(conn, &ServerMessage::Created { reference: 7, thread: "t2".into() });
        match wait(&client.events) {
            Some(ClientEvent::Message(ServerMessage::Created { reference, thread })) => {
                assert_eq!((reference, thread.as_str()), (7, "t2"))
            }
            _ => panic!("client didn't get the reply"),
        }

        let intruder = RemoteClient::connect(&format!("127.0.0.1:{port}"), "wrong".into(), "intruder".into());
        assert!(matches!(wait(&intruder.events), Some(ClientEvent::Status(ConnectionStatus::Connecting))));
        assert!(matches!(wait(&intruder.events), Some(ClientEvent::Status(ConnectionStatus::Denied(_)))));
    }
}