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        /// Explicit git start point (tag, commit, branch). When set, a new local
137        /// branch named `branch` is created from this ref.
138        from_ref: Option<String>,
139        /// If true, fetch from origin before creating the worktree.
140        fetch_latest: bool,
141    },
142    /// `repo_id = None` means all repos
143    WorkstreamList {
144        repo_id: Option<String>,
145    },
146    WorkstreamDelete {
147        workstream_id: String,
148    },
149
150    // ── Agents ────────────────────────────────────────────────────────────────
151    AgentSpawn {
152        workstream_id: String,
153        prompt: String,
154    },
155    /// Claim the caller's current tmux window as an agent (in-place conversion).
156    /// The daemon registers the agent and returns the command to exec.
157    AgentSpawnInPlace {
158        workstream_id: String,
159        /// The existing tmux window index (the caller's current pane)
160        tmux_window: u32,
161        /// Optional task description; `None` means run the agent interactively
162        prompt: Option<String>,
163    },
164    AgentKill {
165        agent_id: String,
166    },
167    AgentList {
168        workstream_id: String,
169    },
170
171    // ── Shells ────────────────────────────────────────────────────────────────
172    /// Spawn a new shell window in a workstream (creates a new tmux window).
173    ShellSpawn {
174        workstream_id: String,
175    },
176    /// Kill a shell session.
177    ShellKill {
178        shell_id: String,
179    },
180    /// List shell sessions for a workstream.
181    ShellList {
182        workstream_id: String,
183    },
184    /// Sent by `vex shell` to register itself with vexd after launching.
185    ShellRegister {
186        workstream_id: String,
187        tmux_window: u32,
188    },
189    /// Attach to a shell session's PTY stream.
190    /// After the response `ShellAttached`, the connection switches to PTY
191    /// streaming mode: vexd emits `PtyOutput` frames; client sends
192    /// `PtyInput` / `PtyResize` frames.
193    AttachShell {
194        shell_id: String,
195    },
196    /// Detach the current client from a shell session.
197    DetachShell {
198        shell_id: String,
199    },
200    /// Send keyboard input to a shell's PTY (base64-encoded bytes).
201    PtyInput {
202        shell_id: String,
203        /// base64-encoded bytes to write to the PTY master
204        data: String,
205    },
206    /// Resize a shell's PTY.
207    PtyResize {
208        shell_id: String,
209        cols: u16,
210        rows: u16,
211    },
212}
213
214#[derive(Debug, Clone, Serialize, Deserialize)]
215#[serde(tag = "type", content = "data")]
216pub enum Response {
217    // ── Existing ──────────────────────────────────────────────────────────────
218    Pong,
219    Ok,
220    DaemonStatus(DaemonStatus),
221    ClientInfo(ClientInfo),
222    /// Returned after PairCreate; contains the plaintext secret (one-time)
223    Pair(PairPayload),
224    PairedClient(PairedClient),
225    PairedClients(Vec<PairedClient>),
226    /// Returned by PairRevoke / PairRevokeAll, carrying the revoked count.
227    Revoked(u32),
228    Error(VexProtoError),
229
230    // ── Repos ─────────────────────────────────────────────────────────────────
231    RepoRegistered(Repository),
232    RepoList(Vec<Repository>),
233    RepoUnregistered,
234    RepoDefaultBranchSet,
235
236    // ── Workstreams ───────────────────────────────────────────────────────────
237    WorkstreamCreated(Workstream),
238    /// Full tree: repos → workstreams → agents
239    WorkstreamList(Vec<Repository>),
240    WorkstreamDeleted,
241
242    // ── Agents ────────────────────────────────────────────────────────────────
243    AgentSpawned(Agent),
244    /// Returned by `AgentSpawnInPlace`; client should `exec` the given command.
245    AgentSpawnedInPlace {
246        agent: Agent,
247        /// Shell command string to exec (replaces the caller's current process)
248        exec_cmd: String,
249    },
250    AgentKilled,
251    AgentList(Vec<Agent>),
252
253    // ── Shells ────────────────────────────────────────────────────────────────
254    ShellSpawned(ShellSession),
255    ShellKilled,
256    ShellList(Vec<ShellSession>),
257    /// Sent back to `vex shell` after `ShellRegister`; carries the assigned ID.
258    ShellRegistered {
259        shell_id: String,
260    },
261    /// Sent back after `AttachShell`; followed by streaming `PtyOutput` frames.
262    ShellAttached,
263    ShellDetached,
264    /// PTY output from the shell (base64-encoded bytes).
265    PtyOutput {
266        shell_id: String,
267        /// base64-encoded raw terminal bytes
268        data: String,
269    },
270    /// Emitted by vexd when a shell process exits.
271    ShellExited {
272        shell_id: String,
273        code: Option<i32>,
274    },
275}
276
277// ── PTY streaming ─────────────────────────────────────────────────────────────
278
279/// Bidirectional PTY streaming message.
280///
281/// Used on two channels:
282/// 1. `vex shell` ↔ vexd: supervisor sends `Out`/`Exited`; vexd sends `In`/`Resize`.
283/// 2. `vex attach` remote client ↔ vexd: vexd sends `Out`/`Exited`; client sends
284///    `In`/`Resize`.
285#[derive(Debug, Clone, Serialize, Deserialize)]
286#[serde(tag = "type", content = "data")]
287pub enum ShellMsg {
288    /// PTY output bytes (base64-encoded) from shell → vexd → attached clients.
289    Out { data: String },
290    /// PTY input bytes (base64-encoded) from attached client → vexd → shell.
291    In { data: String },
292    /// Terminal resize.
293    Resize { cols: u16, rows: u16 },
294    /// Shell process exited.
295    Exited { code: Option<i32> },
296}
297
298// ── Existing helper types ─────────────────────────────────────────────────────
299
300#[derive(Debug, Clone, Serialize, Deserialize)]
301pub struct DaemonStatus {
302    pub uptime_secs: u64,
303    pub connected_clients: u32,
304    pub version: String,
305}
306
307/// Returned by PairCreate — contains the plaintext secret for the new token.
308#[derive(Debug, Clone, Serialize, Deserialize)]
309pub struct PairPayload {
310    pub token_id: String,
311    pub token_secret: String,
312    /// Optional TCP host for encoding into a QR pairing string
313    pub host: Option<String>,
314}
315
316impl PairPayload {
317    /// Returns the pairing string in `<token_id>:<token_secret>` format.
318    pub fn pairing_string(&self) -> String {
319        format!("{}:{}", self.token_id, self.token_secret)
320    }
321}
322
323#[derive(Debug, Clone, Serialize, Deserialize)]
324pub struct PairedClient {
325    pub token_id: String,
326    pub label: Option<String>,
327    pub created_at: String,
328    pub expires_at: Option<String>,
329    pub last_seen: Option<String>,
330}
331
332#[derive(Debug, Clone, Serialize, Deserialize)]
333pub struct ClientInfo {
334    pub token_id: Option<String>,
335    pub is_local: bool,
336}
337
338#[derive(Debug, Clone, Serialize, Deserialize)]
339#[serde(rename_all = "snake_case")]
340pub enum Transport {
341    Unix,
342    Tcp,
343}
344
345#[derive(Debug, Clone, Serialize, Deserialize)]
346#[serde(tag = "code", content = "message")]
347pub enum VexProtoError {
348    Unauthorized,
349    LocalOnly,
350    NotFound,
351    Internal(String),
352}
353
354/// Sent by the client at the start of every TCP connection before any Command.
355#[derive(Debug, Clone, Serialize, Deserialize)]
356pub struct AuthToken {
357    pub token_id: String,
358    /// Plaintext hex-encoded 32-byte secret
359    pub token_secret: String,
360}
361
362// ── Framing ───────────────────────────────────────────────────────────────────
363
364pub mod framing {
365    use serde::{Deserialize, Serialize};
366    use std::io;
367    use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
368
369    #[derive(Debug)]
370    pub enum VexFrameError {
371        Io(io::Error),
372        Json(serde_json::Error),
373    }
374
375    impl std::fmt::Display for VexFrameError {
376        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
377            match self {
378                VexFrameError::Io(e) => write!(f, "IO error: {e}"),
379                VexFrameError::Json(e) => write!(f, "JSON error: {e}"),
380            }
381        }
382    }
383
384    impl std::error::Error for VexFrameError {}
385
386    impl From<io::Error> for VexFrameError {
387        fn from(e: io::Error) -> Self {
388            VexFrameError::Io(e)
389        }
390    }
391
392    impl From<serde_json::Error> for VexFrameError {
393        fn from(e: serde_json::Error) -> Self {
394            VexFrameError::Json(e)
395        }
396    }
397
398    /// Write a length-prefixed JSON frame to `w`.
399    pub async fn send<W, T>(w: &mut W, msg: &T) -> Result<(), VexFrameError>
400    where
401        W: AsyncWrite + Unpin,
402        T: Serialize,
403    {
404        let body = serde_json::to_vec(msg)?;
405        w.write_u32(body.len() as u32).await?;
406        w.write_all(&body).await?;
407        Ok(())
408    }
409
410    /// Maximum allowed frame size (16 MiB). Prevents OOM from malicious length prefixes.
411    const MAX_FRAME_SIZE: u32 = 16 * 1024 * 1024;
412
413    /// Read a length-prefixed JSON frame from `r`.
414    pub async fn recv<R, T>(r: &mut R) -> Result<T, VexFrameError>
415    where
416        R: AsyncRead + Unpin,
417        T: for<'de> Deserialize<'de>,
418    {
419        let len = r.read_u32().await?;
420        if len > MAX_FRAME_SIZE {
421            return Err(VexFrameError::Io(io::Error::new(
422                io::ErrorKind::InvalidData,
423                format!("frame too large: {len} bytes (max {MAX_FRAME_SIZE})"),
424            )));
425        }
426        let mut buf = vec![0u8; len as usize];
427        r.read_exact(&mut buf).await?;
428        Ok(serde_json::from_slice(&buf)?)
429    }
430}
431
432#[cfg(test)]
433mod tests {
434    use super::{Command, Response, ShellMsg, VexProtoError, framing};
435
436    #[tokio::test]
437    async fn framing_command_roundtrip() {
438        let (mut a, mut b) = tokio::io::duplex(4096);
439        framing::send(&mut a, &Command::Status).await.unwrap();
440        let recv: Command = framing::recv(&mut b).await.unwrap();
441        assert!(matches!(recv, Command::Status));
442    }
443
444    #[tokio::test]
445    async fn framing_shell_register_roundtrip() {
446        let (mut a, mut b) = tokio::io::duplex(4096);
447        let cmd = Command::ShellRegister {
448            workstream_id: "ws_abc123".to_string(),
449            tmux_window: 7,
450        };
451        framing::send(&mut a, &cmd).await.unwrap();
452        let recv: Command = framing::recv(&mut b).await.unwrap();
453        match recv {
454            Command::ShellRegister {
455                workstream_id,
456                tmux_window,
457            } => {
458                assert_eq!(workstream_id, "ws_abc123");
459                assert_eq!(tmux_window, 7);
460            }
461            other => panic!("unexpected: {other:?}"),
462        }
463    }
464
465    #[tokio::test]
466    async fn framing_shell_msg_out() {
467        let (mut a, mut b) = tokio::io::duplex(4096);
468        let msg = ShellMsg::Out {
469            data: "aGVsbG8=".to_string(),
470        };
471        framing::send(&mut a, &msg).await.unwrap();
472        let recv: ShellMsg = framing::recv(&mut b).await.unwrap();
473        match recv {
474            ShellMsg::Out { data } => assert_eq!(data, "aGVsbG8="),
475            other => panic!("unexpected: {other:?}"),
476        }
477    }
478
479    #[tokio::test]
480    async fn framing_shell_msg_resize() {
481        let (mut a, mut b) = tokio::io::duplex(4096);
482        framing::send(
483            &mut a,
484            &ShellMsg::Resize {
485                cols: 120,
486                rows: 40,
487            },
488        )
489        .await
490        .unwrap();
491        let recv: ShellMsg = framing::recv(&mut b).await.unwrap();
492        assert!(matches!(
493            recv,
494            ShellMsg::Resize {
495                cols: 120,
496                rows: 40
497            }
498        ));
499    }
500
501    #[tokio::test]
502    async fn framing_shell_msg_exited() {
503        let (mut a, mut b) = tokio::io::duplex(4096);
504        framing::send(&mut a, &ShellMsg::Exited { code: Some(1) })
505            .await
506            .unwrap();
507        let recv: ShellMsg = framing::recv(&mut b).await.unwrap();
508        assert!(matches!(recv, ShellMsg::Exited { code: Some(1) }));
509    }
510
511    #[tokio::test]
512    async fn framing_response_registered() {
513        let (mut a, mut b) = tokio::io::duplex(4096);
514        let resp = Response::ShellRegistered {
515            shell_id: "sh_aabbcc".to_string(),
516        };
517        framing::send(&mut a, &resp).await.unwrap();
518        let recv: Response = framing::recv(&mut b).await.unwrap();
519        match recv {
520            Response::ShellRegistered { shell_id } => assert_eq!(shell_id, "sh_aabbcc"),
521            other => panic!("unexpected: {other:?}"),
522        }
523    }
524
525    #[tokio::test]
526    async fn framing_response_error() {
527        let (mut a, mut b) = tokio::io::duplex(4096);
528        framing::send(&mut a, &Response::Error(VexProtoError::NotFound))
529            .await
530            .unwrap();
531        let recv: Response = framing::recv(&mut b).await.unwrap();
532        assert!(matches!(recv, Response::Error(VexProtoError::NotFound)));
533    }
534}