Skip to main content

vex_cli/proto/
mod.rs

1use serde::{Deserialize, Serialize};
2
3/// Default port vexd listens on for TLS TCP connections.
4pub const DEFAULT_TCP_PORT: u16 = 7422;
5
6// ── Domain types ──────────────────────────────────────────────────────────────
7
8fn default_branch_fallback() -> String {
9    "main".to_string()
10}
11
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct Repository {
14    pub id: String,
15    pub name: String,
16    /// Absolute path to the git repository on disk
17    pub path: String,
18    /// Default branch used when creating a workstream without an explicit branch.
19    /// Falls back to "main" for repos persisted before this field was added.
20    #[serde(default = "default_branch_fallback")]
21    pub default_branch: String,
22    pub registered_at: u64,
23    pub workstreams: Vec<Workstream>,
24}
25
26#[derive(Debug, Clone, Serialize, Deserialize)]
27pub struct Workstream {
28    pub id: String,
29    pub name: String,
30    pub repo_id: String,
31    pub branch: String,
32    /// Absolute path: `$VEX_HOME/worktrees/<workstream_id>`
33    pub worktree_path: String,
34    /// Always `"vex-<workstream_id>"`
35    pub tmux_session: String,
36    pub status: WorkstreamStatus,
37    pub agents: Vec<Agent>,
38    /// Active and recent shell sessions in this workstream
39    #[serde(default)]
40    pub shells: Vec<ShellSession>,
41    pub created_at: u64,
42}
43
44// ── Shell session ─────────────────────────────────────────────────────────────
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
47pub struct ShellSession {
48    pub id: String,
49    pub workstream_id: String,
50    /// tmux window index that hosts this shell
51    pub tmux_window: u32,
52    pub status: ShellStatus,
53    pub started_at: u64,
54    pub exited_at: Option<u64>,
55    pub exit_code: Option<i32>,
56}
57
58#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
59pub enum ShellStatus {
60    /// Shell is running and accepting PTY I/O
61    Active,
62    /// Shell is running but no client is currently attached
63    Detached,
64    /// Shell process has exited
65    Exited,
66}
67
68#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
69pub enum WorkstreamStatus {
70    Idle,
71    Running,
72    Stopped,
73}
74
75#[derive(Debug, Clone, Serialize, Deserialize)]
76pub struct Agent {
77    pub id: String,
78    pub workstream_id: String,
79    /// Window index in the tmux session
80    pub tmux_window: u32,
81    pub prompt: String,
82    pub status: AgentStatus,
83    pub exit_code: Option<i32>,
84    pub spawned_at: u64,
85    pub exited_at: Option<u64>,
86}
87
88#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
89pub enum AgentStatus {
90    Running,
91    Exited,
92    Failed,
93}
94
95// ── Wire types ────────────────────────────────────────────────────────────────
96
97#[derive(Debug, Clone, Serialize, Deserialize)]
98#[serde(tag = "type", content = "data")]
99pub enum Command {
100    // ── Existing ──────────────────────────────────────────────────────────────
101    Status,
102    Whoami,
103    PairCreate {
104        label: Option<String>,
105        /// Expiry in seconds from now
106        expire_secs: Option<u64>,
107    },
108    PairList,
109    PairRevoke {
110        id: String,
111    },
112    PairRevokeAll,
113
114    // ── Repos (LocalOnly) ─────────────────────────────────────────────────────
115    /// Register a git repository. Unix-socket only (LocalOnly on TCP).
116    RepoRegister {
117        path: String,
118    },
119    RepoList,
120    RepoUnregister {
121        repo_id: String,
122    },
123    /// Update the default branch stored for a repo. Unix-socket only.
124    RepoSetDefaultBranch {
125        repo_id: String,
126        branch: String,
127    },
128
129    // ── Workstreams ───────────────────────────────────────────────────────────
130    WorkstreamCreate {
131        repo_id: String,
132        /// Workstream name. `None` = use the resolved branch name.
133        name: Option<String>,
134        /// Branch to check out. `None` = use the repo's `default_branch`.
135        branch: Option<String>,
136        /// If true, fetch `origin/<branch>` and fast-forward the local branch
137        /// before creating the worktree.
138        fetch_latest: bool,
139    },
140    /// `repo_id = None` means all repos
141    WorkstreamList {
142        repo_id: Option<String>,
143    },
144    WorkstreamDelete {
145        workstream_id: String,
146    },
147
148    // ── Agents ────────────────────────────────────────────────────────────────
149    AgentSpawn {
150        workstream_id: String,
151        prompt: String,
152    },
153    /// Claim the caller's current tmux window as an agent (in-place conversion).
154    /// The daemon registers the agent and returns the command to exec.
155    AgentSpawnInPlace {
156        workstream_id: String,
157        /// The existing tmux window index (the caller's current pane)
158        tmux_window: u32,
159        /// Optional task description; `None` means run the agent interactively
160        prompt: Option<String>,
161    },
162    AgentKill {
163        agent_id: String,
164    },
165    AgentList {
166        workstream_id: String,
167    },
168
169    // ── Shells ────────────────────────────────────────────────────────────────
170    /// Spawn a new shell window in a workstream (creates a new tmux window).
171    ShellSpawn {
172        workstream_id: String,
173    },
174    /// Kill a shell session.
175    ShellKill {
176        shell_id: String,
177    },
178    /// List shell sessions for a workstream.
179    ShellList {
180        workstream_id: String,
181    },
182    /// Sent by `vex shell` to register itself with vexd after launching.
183    ShellRegister {
184        workstream_id: String,
185        tmux_window: u32,
186    },
187    /// Attach to a shell session's PTY stream.
188    /// After the response `ShellAttached`, the connection switches to PTY
189    /// streaming mode: vexd emits `PtyOutput` frames; client sends
190    /// `PtyInput` / `PtyResize` frames.
191    AttachShell {
192        shell_id: String,
193    },
194    /// Detach the current client from a shell session.
195    DetachShell {
196        shell_id: String,
197    },
198    /// Send keyboard input to a shell's PTY (base64-encoded bytes).
199    PtyInput {
200        shell_id: String,
201        /// base64-encoded bytes to write to the PTY master
202        data: String,
203    },
204    /// Resize a shell's PTY.
205    PtyResize {
206        shell_id: String,
207        cols: u16,
208        rows: u16,
209    },
210}
211
212#[derive(Debug, Clone, Serialize, Deserialize)]
213#[serde(tag = "type", content = "data")]
214pub enum Response {
215    // ── Existing ──────────────────────────────────────────────────────────────
216    Pong,
217    Ok,
218    DaemonStatus(DaemonStatus),
219    ClientInfo(ClientInfo),
220    /// Returned after PairCreate; contains the plaintext secret (one-time)
221    Pair(PairPayload),
222    PairedClient(PairedClient),
223    PairedClients(Vec<PairedClient>),
224    /// Returned by PairRevoke / PairRevokeAll, carrying the revoked count.
225    Revoked(u32),
226    Error(VexProtoError),
227
228    // ── Repos ─────────────────────────────────────────────────────────────────
229    RepoRegistered(Repository),
230    RepoList(Vec<Repository>),
231    RepoUnregistered,
232    RepoDefaultBranchSet,
233
234    // ── Workstreams ───────────────────────────────────────────────────────────
235    WorkstreamCreated(Workstream),
236    /// Full tree: repos → workstreams → agents
237    WorkstreamList(Vec<Repository>),
238    WorkstreamDeleted,
239
240    // ── Agents ────────────────────────────────────────────────────────────────
241    AgentSpawned(Agent),
242    /// Returned by `AgentSpawnInPlace`; client should `exec` the given command.
243    AgentSpawnedInPlace {
244        agent: Agent,
245        /// Shell command string to exec (replaces the caller's current process)
246        exec_cmd: String,
247    },
248    AgentKilled,
249    AgentList(Vec<Agent>),
250
251    // ── Shells ────────────────────────────────────────────────────────────────
252    ShellSpawned(ShellSession),
253    ShellKilled,
254    ShellList(Vec<ShellSession>),
255    /// Sent back to `vex shell` after `ShellRegister`; carries the assigned ID.
256    ShellRegistered {
257        shell_id: String,
258    },
259    /// Sent back after `AttachShell`; followed by streaming `PtyOutput` frames.
260    ShellAttached,
261    ShellDetached,
262    /// PTY output from the shell (base64-encoded bytes).
263    PtyOutput {
264        shell_id: String,
265        /// base64-encoded raw terminal bytes
266        data: String,
267    },
268    /// Emitted by vexd when a shell process exits.
269    ShellExited {
270        shell_id: String,
271        code: Option<i32>,
272    },
273}
274
275// ── PTY streaming ─────────────────────────────────────────────────────────────
276
277/// Bidirectional PTY streaming message.
278///
279/// Used on two channels:
280/// 1. `vex shell` ↔ vexd: supervisor sends `Out`/`Exited`; vexd sends `In`/`Resize`.
281/// 2. `vex attach` remote client ↔ vexd: vexd sends `Out`/`Exited`; client sends
282///    `In`/`Resize`.
283#[derive(Debug, Clone, Serialize, Deserialize)]
284#[serde(tag = "type", content = "data")]
285pub enum ShellMsg {
286    /// PTY output bytes (base64-encoded) from shell → vexd → attached clients.
287    Out { data: String },
288    /// PTY input bytes (base64-encoded) from attached client → vexd → shell.
289    In { data: String },
290    /// Terminal resize.
291    Resize { cols: u16, rows: u16 },
292    /// Shell process exited.
293    Exited { code: Option<i32> },
294}
295
296// ── Existing helper types ─────────────────────────────────────────────────────
297
298#[derive(Debug, Clone, Serialize, Deserialize)]
299pub struct DaemonStatus {
300    pub uptime_secs: u64,
301    pub connected_clients: u32,
302    pub version: String,
303}
304
305/// Returned by PairCreate — contains the plaintext secret for the new token.
306#[derive(Debug, Clone, Serialize, Deserialize)]
307pub struct PairPayload {
308    pub token_id: String,
309    pub token_secret: String,
310    /// Optional TCP host for encoding into a QR pairing string
311    pub host: Option<String>,
312}
313
314impl PairPayload {
315    /// Returns the pairing string in `<token_id>:<token_secret>` format.
316    pub fn pairing_string(&self) -> String {
317        format!("{}:{}", self.token_id, self.token_secret)
318    }
319}
320
321#[derive(Debug, Clone, Serialize, Deserialize)]
322pub struct PairedClient {
323    pub token_id: String,
324    pub label: Option<String>,
325    pub created_at: String,
326    pub expires_at: Option<String>,
327    pub last_seen: Option<String>,
328}
329
330#[derive(Debug, Clone, Serialize, Deserialize)]
331pub struct ClientInfo {
332    pub token_id: Option<String>,
333    pub is_local: bool,
334}
335
336#[derive(Debug, Clone, Serialize, Deserialize)]
337#[serde(rename_all = "snake_case")]
338pub enum Transport {
339    Unix,
340    Tcp,
341}
342
343#[derive(Debug, Clone, Serialize, Deserialize)]
344#[serde(tag = "code", content = "message")]
345pub enum VexProtoError {
346    Unauthorized,
347    LocalOnly,
348    NotFound,
349    Internal(String),
350}
351
352/// Sent by the client at the start of every TCP connection before any Command.
353#[derive(Debug, Clone, Serialize, Deserialize)]
354pub struct AuthToken {
355    pub token_id: String,
356    /// Plaintext hex-encoded 32-byte secret
357    pub token_secret: String,
358}
359
360// ── Framing ───────────────────────────────────────────────────────────────────
361
362pub mod framing {
363    use serde::{Deserialize, Serialize};
364    use std::io;
365    use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
366
367    #[derive(Debug)]
368    pub enum VexFrameError {
369        Io(io::Error),
370        Json(serde_json::Error),
371    }
372
373    impl std::fmt::Display for VexFrameError {
374        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
375            match self {
376                VexFrameError::Io(e) => write!(f, "IO error: {e}"),
377                VexFrameError::Json(e) => write!(f, "JSON error: {e}"),
378            }
379        }
380    }
381
382    impl std::error::Error for VexFrameError {}
383
384    impl From<io::Error> for VexFrameError {
385        fn from(e: io::Error) -> Self {
386            VexFrameError::Io(e)
387        }
388    }
389
390    impl From<serde_json::Error> for VexFrameError {
391        fn from(e: serde_json::Error) -> Self {
392            VexFrameError::Json(e)
393        }
394    }
395
396    /// Write a length-prefixed JSON frame to `w`.
397    pub async fn send<W, T>(w: &mut W, msg: &T) -> Result<(), VexFrameError>
398    where
399        W: AsyncWrite + Unpin,
400        T: Serialize,
401    {
402        let body = serde_json::to_vec(msg)?;
403        w.write_u32(body.len() as u32).await?;
404        w.write_all(&body).await?;
405        Ok(())
406    }
407
408    /// Maximum allowed frame size (16 MiB). Prevents OOM from malicious length prefixes.
409    const MAX_FRAME_SIZE: u32 = 16 * 1024 * 1024;
410
411    /// Read a length-prefixed JSON frame from `r`.
412    pub async fn recv<R, T>(r: &mut R) -> Result<T, VexFrameError>
413    where
414        R: AsyncRead + Unpin,
415        T: for<'de> Deserialize<'de>,
416    {
417        let len = r.read_u32().await?;
418        if len > MAX_FRAME_SIZE {
419            return Err(VexFrameError::Io(io::Error::new(
420                io::ErrorKind::InvalidData,
421                format!("frame too large: {len} bytes (max {MAX_FRAME_SIZE})"),
422            )));
423        }
424        let mut buf = vec![0u8; len as usize];
425        r.read_exact(&mut buf).await?;
426        Ok(serde_json::from_slice(&buf)?)
427    }
428}
429
430#[cfg(test)]
431mod tests {
432    use super::{Command, Response, ShellMsg, VexProtoError, framing};
433
434    #[tokio::test]
435    async fn framing_command_roundtrip() {
436        let (mut a, mut b) = tokio::io::duplex(4096);
437        framing::send(&mut a, &Command::Status).await.unwrap();
438        let recv: Command = framing::recv(&mut b).await.unwrap();
439        assert!(matches!(recv, Command::Status));
440    }
441
442    #[tokio::test]
443    async fn framing_shell_register_roundtrip() {
444        let (mut a, mut b) = tokio::io::duplex(4096);
445        let cmd = Command::ShellRegister {
446            workstream_id: "ws_abc123".to_string(),
447            tmux_window: 7,
448        };
449        framing::send(&mut a, &cmd).await.unwrap();
450        let recv: Command = framing::recv(&mut b).await.unwrap();
451        match recv {
452            Command::ShellRegister {
453                workstream_id,
454                tmux_window,
455            } => {
456                assert_eq!(workstream_id, "ws_abc123");
457                assert_eq!(tmux_window, 7);
458            }
459            other => panic!("unexpected: {other:?}"),
460        }
461    }
462
463    #[tokio::test]
464    async fn framing_shell_msg_out() {
465        let (mut a, mut b) = tokio::io::duplex(4096);
466        let msg = ShellMsg::Out {
467            data: "aGVsbG8=".to_string(),
468        };
469        framing::send(&mut a, &msg).await.unwrap();
470        let recv: ShellMsg = framing::recv(&mut b).await.unwrap();
471        match recv {
472            ShellMsg::Out { data } => assert_eq!(data, "aGVsbG8="),
473            other => panic!("unexpected: {other:?}"),
474        }
475    }
476
477    #[tokio::test]
478    async fn framing_shell_msg_resize() {
479        let (mut a, mut b) = tokio::io::duplex(4096);
480        framing::send(
481            &mut a,
482            &ShellMsg::Resize {
483                cols: 120,
484                rows: 40,
485            },
486        )
487        .await
488        .unwrap();
489        let recv: ShellMsg = framing::recv(&mut b).await.unwrap();
490        assert!(matches!(
491            recv,
492            ShellMsg::Resize {
493                cols: 120,
494                rows: 40
495            }
496        ));
497    }
498
499    #[tokio::test]
500    async fn framing_shell_msg_exited() {
501        let (mut a, mut b) = tokio::io::duplex(4096);
502        framing::send(&mut a, &ShellMsg::Exited { code: Some(1) })
503            .await
504            .unwrap();
505        let recv: ShellMsg = framing::recv(&mut b).await.unwrap();
506        assert!(matches!(recv, ShellMsg::Exited { code: Some(1) }));
507    }
508
509    #[tokio::test]
510    async fn framing_response_registered() {
511        let (mut a, mut b) = tokio::io::duplex(4096);
512        let resp = Response::ShellRegistered {
513            shell_id: "sh_aabbcc".to_string(),
514        };
515        framing::send(&mut a, &resp).await.unwrap();
516        let recv: Response = framing::recv(&mut b).await.unwrap();
517        match recv {
518            Response::ShellRegistered { shell_id } => assert_eq!(shell_id, "sh_aabbcc"),
519            other => panic!("unexpected: {other:?}"),
520        }
521    }
522
523    #[tokio::test]
524    async fn framing_response_error() {
525        let (mut a, mut b) = tokio::io::duplex(4096);
526        framing::send(&mut a, &Response::Error(VexProtoError::NotFound))
527            .await
528            .unwrap();
529        let recv: Response = framing::recv(&mut b).await.unwrap();
530        assert!(matches!(recv, Response::Error(VexProtoError::NotFound)));
531    }
532}