agent-orchestrator/agor-pty/src/protocol.rs

use std::collections::HashMap;
use serde::{Deserialize, Serialize};

/// Messages sent from client → daemon.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ClientMessage {
    Auth {
        token: String,
    },
    CreateSession {
        id: String,
        shell: Option<String>,
        cwd: Option<String>,
        env: Option<HashMap<String, String>>,
        cols: u16,
        rows: u16,
    },
    WriteInput {
        session_id: String,
        /// Raw bytes encoded as a base64 string.
        data: String,
    },
    Resize {
        session_id: String,
        cols: u16,
        rows: u16,
    },
    Subscribe {
        session_id: String,
    },
    Unsubscribe {
        session_id: String,
    },
    CloseSession {
        session_id: String,
    },
    ListSessions,
    Ping,
}

/// Messages sent from daemon → client.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum DaemonMessage {
    AuthResult {
        ok: bool,
    },
    SessionCreated {
        session_id: String,
        pid: u32,
    },
    /// PTY output bytes base64-encoded so they survive JSON transport.
    SessionOutput {
        session_id: String,
        data: String,
    },
    SessionClosed {
        session_id: String,
        exit_code: Option<i32>,
    },
    SessionList {
        sessions: Vec<SessionInfo>,
    },
    Pong,
    Error {
        message: String,
    },
}

/// Snapshot of a running or recently-exited session.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SessionInfo {
    pub id: String,
    pub pid: u32,
    pub shell: String,
    pub cwd: String,
    pub cols: u16,
    pub rows: u16,
    /// Unix timestamp of session creation.
    pub created_at: u64,
    pub alive: bool,
}

/// Encode raw bytes as base64 for embedding in JSON.
pub fn encode_output(bytes: &[u8]) -> String {
    use std::fmt::Write as FmtWrite;
    // Manual base64 — avoids adding a new crate; the hex crate IS available but
    // base64 better compresses binary PTY data (33% overhead vs 100% for hex).
    // We use a simple lookup table implementation that stays within 300 lines.
    const TABLE: &[u8; 64] =
        b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
    let mut out = String::with_capacity((bytes.len() + 2) / 3 * 4);
    let mut i = 0;
    while i + 2 < bytes.len() {
        let b0 = bytes[i] as usize;
        let b1 = bytes[i + 1] as usize;
        let b2 = bytes[i + 2] as usize;
        let _ = write!(
            out,
            "{}{}{}{}",
            TABLE[b0 >> 2] as char,
            TABLE[((b0 & 3) << 4) | (b1 >> 4)] as char,
            TABLE[((b1 & 0xf) << 2) | (b2 >> 6)] as char,
            TABLE[b2 & 0x3f] as char
        );
        i += 3;
    }
    let rem = bytes.len() - i;
    if rem == 1 {
        let b0 = bytes[i] as usize;
        let _ = write!(
            out,
            "{}{}==",
            TABLE[b0 >> 2] as char,
            TABLE[(b0 & 3) << 4] as char
        );
    } else if rem == 2 {
        let b0 = bytes[i] as usize;
        let b1 = bytes[i + 1] as usize;
        let _ = write!(
            out,
            "{}{}{}=",
            TABLE[b0 >> 2] as char,
            TABLE[((b0 & 3) << 4) | (b1 >> 4)] as char,
            TABLE[(b1 & 0xf) << 2] as char
        );
    }
    out
}

/// Decode base64 string back to bytes. Returns error string on invalid input.
pub fn decode_input(s: &str) -> Result<Vec<u8>, String> {
    fn val(c: u8) -> Result<u8, String> {
        match c {
            b'A'..=b'Z' => Ok(c - b'A'),
            b'a'..=b'z' => Ok(c - b'a' + 26),
            b'0'..=b'9' => Ok(c - b'0' + 52),
            b'+' => Ok(62),
            b'/' => Ok(63),
            b'=' => Ok(0),
            _ => Err(format!("invalid base64 char: {c}")),
        }
    }
    let bytes = s.as_bytes();
    if bytes.len() % 4 != 0 {
        return Err("base64 length not a multiple of 4".into());
    }
    let mut out = Vec::with_capacity(bytes.len() / 4 * 3);
    let mut i = 0;
    while i < bytes.len() {
        let v0 = val(bytes[i])?;
        let v1 = val(bytes[i + 1])?;
        let v2 = val(bytes[i + 2])?;
        let v3 = val(bytes[i + 3])?;
        out.push((v0 << 2) | (v1 >> 4));
        if bytes[i + 2] != b'=' {
            out.push((v1 << 4) | (v2 >> 2));
        }
        if bytes[i + 3] != b'=' {
            out.push((v2 << 6) | v3);
        }
        i += 4;
    }
    Ok(out)
}