Skip to main content

vex_client/
lib.rs

1use anyhow::{Result, bail};
2use tokio::io;
3use tokio::net::TcpStream;
4use uuid::Uuid;
5use vex_proto::{
6    AgentEntry, ClientMessage, Frame, RepoEntry, ServerMessage, ShellInfo, WorkstreamInfo,
7    read_frame, send_client_message,
8};
9
10/// Typed async client for communicating with the vex daemon over TCP.
11pub struct DaemonClient {
12    port: u16,
13}
14
15impl DaemonClient {
16    pub fn new(port: u16) -> Self {
17        Self { port }
18    }
19
20    pub fn port(&self) -> u16 {
21        self.port
22    }
23
24    /// Open a raw TCP connection to the daemon.
25    pub async fn connect(&self) -> Result<TcpStream> {
26        TcpStream::connect(("127.0.0.1", self.port))
27            .await
28            .map_err(|e| {
29                anyhow::anyhow!(
30                    "could not connect to daemon on port {}: {} (is the daemon running?)",
31                    self.port,
32                    e
33                )
34            })
35    }
36
37    /// One-shot request-response: connect, send, read one control frame.
38    pub async fn request(&self, msg: &ClientMessage) -> Result<ServerMessage> {
39        let stream = self.connect().await?;
40        let (mut reader, mut writer) = io::split(stream);
41        send_client_message(&mut writer, msg).await?;
42        match read_frame(&mut reader).await? {
43            Some(Frame::Control(data)) => {
44                let resp: ServerMessage = serde_json::from_slice(&data)?;
45                Ok(resp)
46            }
47            Some(Frame::Data(_)) => bail!("unexpected data frame"),
48            None => bail!("server closed connection"),
49        }
50    }
51
52    // ── Typed convenience methods ──────────────────────────────────
53
54    pub async fn shell_create(
55        &self,
56        shell: Option<String>,
57        repo: Option<String>,
58        workstream: Option<String>,
59    ) -> Result<Uuid> {
60        let resp = self
61            .request(&ClientMessage::CreateShell {
62                shell,
63                repo,
64                workstream,
65            })
66            .await?;
67        match resp {
68            ServerMessage::ShellCreated { id } => Ok(id),
69            ServerMessage::Error { message } => bail!("{}", message),
70            other => bail!("unexpected response: {:?}", other),
71        }
72    }
73
74    pub async fn shell_list(&self) -> Result<Vec<ShellInfo>> {
75        let resp = self.request(&ClientMessage::ListShells).await?;
76        match resp {
77            ServerMessage::Shells { shells } => Ok(shells),
78            ServerMessage::Error { message } => bail!("{}", message),
79            other => bail!("unexpected response: {:?}", other),
80        }
81    }
82
83    pub async fn shell_kill(&self, id: Uuid) -> Result<()> {
84        let resp = self.request(&ClientMessage::KillShell { id }).await?;
85        match resp {
86            ServerMessage::ShellEnded { .. } => Ok(()),
87            ServerMessage::Error { message } => bail!("{}", message),
88            _ => Ok(()),
89        }
90    }
91
92    pub async fn agent_list(&self) -> Result<Vec<AgentEntry>> {
93        let resp = self.request(&ClientMessage::AgentList).await?;
94        match resp {
95            ServerMessage::AgentListResponse { agents } => Ok(agents),
96            ServerMessage::Error { message } => bail!("{}", message),
97            other => bail!("unexpected response: {:?}", other),
98        }
99    }
100
101    pub async fn agent_notifications(&self) -> Result<Vec<AgentEntry>> {
102        let resp = self.request(&ClientMessage::AgentNotifications).await?;
103        match resp {
104            ServerMessage::AgentListResponse { agents } => Ok(agents),
105            ServerMessage::Error { message } => bail!("{}", message),
106            other => bail!("unexpected response: {:?}", other),
107        }
108    }
109
110    pub async fn agent_spawn(&self, repo: &str, workstream: Option<&str>) -> Result<Uuid> {
111        let resp = self
112            .request(&ClientMessage::AgentSpawn {
113                repo: repo.to_string(),
114                workstream: workstream.map(String::from),
115            })
116            .await?;
117        match resp {
118            ServerMessage::ShellCreated { id } => Ok(id),
119            ServerMessage::Error { message } => bail!("{}", message),
120            other => bail!("unexpected response: {:?}", other),
121        }
122    }
123
124    pub async fn agent_prompt(&self, shell_id: Uuid, text: &str) -> Result<()> {
125        let resp = self
126            .request(&ClientMessage::AgentPrompt {
127                shell_id,
128                text: text.to_string(),
129            })
130            .await?;
131        match resp {
132            ServerMessage::AgentPromptSent { .. } => Ok(()),
133            ServerMessage::Error { message } => bail!("{}", message),
134            _ => Ok(()),
135        }
136    }
137
138    pub async fn repo_list(&self) -> Result<Vec<RepoEntry>> {
139        let resp = self.request(&ClientMessage::RepoList).await?;
140        match resp {
141            ServerMessage::Repos { repos } => Ok(repos),
142            ServerMessage::Error { message } => bail!("{}", message),
143            other => bail!("unexpected response: {:?}", other),
144        }
145    }
146
147    pub async fn repo_add(&self, name: &str, path: &std::path::Path) -> Result<()> {
148        let resp = self
149            .request(&ClientMessage::RepoAdd {
150                name: name.to_string(),
151                path: path.to_path_buf(),
152            })
153            .await?;
154        match resp {
155            ServerMessage::RepoAdded { .. } => Ok(()),
156            ServerMessage::Error { message } => bail!("{}", message),
157            other => bail!("unexpected response: {:?}", other),
158        }
159    }
160
161    pub async fn workstream_create(&self, repo: &str, name: &str) -> Result<()> {
162        let resp = self
163            .request(&ClientMessage::WorkstreamCreate {
164                repo: repo.to_string(),
165                name: name.to_string(),
166            })
167            .await?;
168        match resp {
169            ServerMessage::WorkstreamCreated { .. } => Ok(()),
170            ServerMessage::Error { message } => bail!("{}", message),
171            other => bail!("unexpected response: {:?}", other),
172        }
173    }
174
175    pub async fn workstream_list(&self, repo: Option<&str>) -> Result<Vec<WorkstreamInfo>> {
176        let resp = self
177            .request(&ClientMessage::WorkstreamList {
178                repo: repo.map(String::from),
179            })
180            .await?;
181        match resp {
182            ServerMessage::Workstreams { workstreams } => Ok(workstreams),
183            ServerMessage::Error { message } => bail!("{}", message),
184            other => bail!("unexpected response: {:?}", other),
185        }
186    }
187}