Skip to main content

vex_client/
lib.rs

1use anyhow::{Result, bail};
2use tokio::io;
3use tokio::net::TcpStream;
4use vex_proto::{
5    AgentId, AgentInfo, ClientMessage, CommitInfo, DaemonStatusInfo, Frame, RepoBranchInfo,
6    RepoInfo, RepoStatusInfo, ServerMessage, ShellId, ShellInfo, SkillInfo, WorkstreamId,
7    WorkstreamInfo, read_frame, send_client_message,
8};
9
10/// Typed async client for communicating with the vex daemon.
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        let port = self.port();
27        TcpStream::connect(("127.0.0.1", port)).await.map_err(|e| {
28            anyhow::anyhow!(
29                "could not connect to daemon on port {}: {} (is the daemon running?)",
30                port,
31                e
32            )
33        })
34    }
35
36    /// One-shot request-response: connect, send, read one control frame.
37    pub async fn request(&self, msg: &ClientMessage) -> Result<ServerMessage> {
38        let stream = self.connect().await?;
39        let (mut reader, mut writer) = io::split(stream);
40        send_client_message(&mut writer, msg).await?;
41        match read_frame(&mut reader).await? {
42            Some(Frame::Control(data)) => {
43                let resp: ServerMessage = serde_json::from_slice(&data)?;
44                Ok(resp)
45            }
46            Some(Frame::PtyData { .. }) => bail!("unexpected pty data frame"),
47            Some(Frame::Event(_)) => bail!("unexpected event frame"),
48            None => bail!("server closed connection"),
49        }
50    }
51
52    // ── Shell methods ───────────────────────────────────────────
53
54    pub async fn shell_spawn(
55        &self,
56        workdir: Option<String>,
57        name: Option<String>,
58        workstream_id: Option<WorkstreamId>,
59        attach: bool,
60    ) -> Result<ShellId> {
61        let resp = self
62            .request(&ClientMessage::ShellSpawn {
63                attach,
64                workdir,
65                name,
66                workstream_id,
67            })
68            .await?;
69        match resp {
70            ServerMessage::ShellCreated { id } => Ok(id),
71            ServerMessage::Error { message } => bail!("{}", message),
72            other => bail!("unexpected response: {:?}", other),
73        }
74    }
75
76    pub async fn shell_list(&self) -> Result<Vec<ShellInfo>> {
77        let resp = self.request(&ClientMessage::ShellList).await?;
78        match resp {
79            ServerMessage::ShellList { shells } => Ok(shells),
80            ServerMessage::Error { message } => bail!("{}", message),
81            other => bail!("unexpected response: {:?}", other),
82        }
83    }
84
85    pub async fn shell_kill(&self, id: ShellId) -> Result<()> {
86        let resp = self.request(&ClientMessage::ShellKill { id }).await?;
87        match resp {
88            ServerMessage::ShellEnded { .. } | ServerMessage::Ok => Ok(()),
89            ServerMessage::Error { message } => bail!("{}", message),
90            _ => Ok(()),
91        }
92    }
93
94    pub async fn shell_send(
95        &self,
96        id: ShellId,
97        text: Option<String>,
98        keys: Option<Vec<String>>,
99    ) -> Result<()> {
100        let resp = self
101            .request(&ClientMessage::ShellSend { id, text, keys })
102            .await?;
103        match resp {
104            ServerMessage::Ok => Ok(()),
105            ServerMessage::Error { message } => bail!("{}", message),
106            _ => Ok(()),
107        }
108    }
109
110    pub async fn shell_screenshot(&self, id: ShellId) -> Result<String> {
111        let resp = self.request(&ClientMessage::ShellScreenshot { id }).await?;
112        match resp {
113            ServerMessage::ShellScreenshot { content, .. } => Ok(content),
114            ServerMessage::Error { message } => bail!("{}", message),
115            other => bail!("unexpected response: {:?}", other),
116        }
117    }
118
119    pub async fn shell_logs(&self, id: ShellId, last: Option<usize>) -> Result<String> {
120        let resp = self.request(&ClientMessage::ShellLogs { id, last }).await?;
121        match resp {
122            ServerMessage::ShellLogs { content, .. } => Ok(content),
123            ServerMessage::Error { message } => bail!("{}", message),
124            other => bail!("unexpected response: {:?}", other),
125        }
126    }
127
128    // ── Agent methods ───────────────────────────────────────────
129
130    pub async fn agent_spawn(
131        &self,
132        workdir: Option<String>,
133        prompt: Option<String>,
134        workstream_id: Option<WorkstreamId>,
135    ) -> Result<(AgentId, ShellId)> {
136        let resp = self
137            .request(&ClientMessage::AgentSpawn {
138                workdir,
139                prompt,
140                workstream_id,
141            })
142            .await?;
143        match resp {
144            ServerMessage::AgentCreated { id, shell_id } => Ok((id, shell_id)),
145            ServerMessage::Error { message } => bail!("{}", message),
146            other => bail!("unexpected response: {:?}", other),
147        }
148    }
149
150    pub async fn agent_list(&self) -> Result<Vec<AgentInfo>> {
151        let resp = self.request(&ClientMessage::AgentList).await?;
152        match resp {
153            ServerMessage::AgentList { agents } => Ok(agents),
154            ServerMessage::Error { message } => bail!("{}", message),
155            other => bail!("unexpected response: {:?}", other),
156        }
157    }
158
159    pub async fn agent_send(&self, id: AgentId, message: &str) -> Result<()> {
160        let resp = self
161            .request(&ClientMessage::AgentSend {
162                id,
163                message: message.to_string(),
164            })
165            .await?;
166        match resp {
167            ServerMessage::Ok => Ok(()),
168            ServerMessage::Error { message } => bail!("{}", message),
169            _ => Ok(()),
170        }
171    }
172
173    pub async fn agent_kill(&self, id: AgentId) -> Result<()> {
174        let resp = self.request(&ClientMessage::AgentKill { id }).await?;
175        match resp {
176            ServerMessage::Ok => Ok(()),
177            ServerMessage::Error { message } => bail!("{}", message),
178            _ => Ok(()),
179        }
180    }
181
182    pub async fn agent_pause(&self, id: AgentId) -> Result<()> {
183        let resp = self.request(&ClientMessage::AgentPause { id }).await?;
184        match resp {
185            ServerMessage::Ok => Ok(()),
186            ServerMessage::Error { message } => bail!("{}", message),
187            _ => Ok(()),
188        }
189    }
190
191    pub async fn agent_resume(&self, id: AgentId, prompt: Option<String>) -> Result<()> {
192        let resp = self
193            .request(&ClientMessage::AgentResume { id, prompt })
194            .await?;
195        match resp {
196            ServerMessage::Ok => Ok(()),
197            ServerMessage::Error { message } => bail!("{}", message),
198            _ => Ok(()),
199        }
200    }
201
202    pub async fn agent_history(&self, id: AgentId) -> Result<Vec<String>> {
203        let resp = self.request(&ClientMessage::AgentHistory { id }).await?;
204        match resp {
205            ServerMessage::AgentHistory { lines, .. } => Ok(lines),
206            ServerMessage::Error { message } => bail!("{}", message),
207            other => bail!("unexpected response: {:?}", other),
208        }
209    }
210
211    pub async fn agent_screenshot(&self, id: AgentId) -> Result<String> {
212        let resp = self.request(&ClientMessage::AgentScreenshot { id }).await?;
213        match resp {
214            ServerMessage::AgentScreenshot { content, .. } => Ok(content),
215            ServerMessage::Error { message } => bail!("{}", message),
216            other => bail!("unexpected response: {:?}", other),
217        }
218    }
219
220    // ── Workstream methods ──────────────────────────────────────
221
222    pub async fn ws_create(
223        &self,
224        name: &str,
225        branch: Option<&str>,
226        repo: Option<&str>,
227        issue: Option<&str>,
228        from: Option<&str>,
229    ) -> Result<WorkstreamId> {
230        let resp = self
231            .request(&ClientMessage::WsCreate {
232                name: name.to_string(),
233                branch: branch.map(String::from),
234                repo: repo.map(String::from),
235                issue: issue.map(String::from),
236                from: from.map(String::from),
237            })
238            .await?;
239        match resp {
240            ServerMessage::WsCreated { id } => Ok(id),
241            ServerMessage::Error { message } => bail!("{}", message),
242            other => bail!("unexpected response: {:?}", other),
243        }
244    }
245
246    pub async fn ws_list(&self) -> Result<Vec<WorkstreamInfo>> {
247        let resp = self.request(&ClientMessage::WsList).await?;
248        match resp {
249            ServerMessage::WsList { workstreams } => Ok(workstreams),
250            ServerMessage::Error { message } => bail!("{}", message),
251            other => bail!("unexpected response: {:?}", other),
252        }
253    }
254
255    pub async fn ws_show(
256        &self,
257        id: WorkstreamId,
258    ) -> Result<(
259        WorkstreamInfo,
260        Vec<ShellInfo>,
261        Vec<AgentInfo>,
262        Vec<vex_proto::EventEntry>,
263    )> {
264        let resp = self.request(&ClientMessage::WsShow { id }).await?;
265        match resp {
266            ServerMessage::WsDetail {
267                workstream,
268                shells,
269                agents,
270                events,
271            } => Ok((workstream, shells, agents, events)),
272            ServerMessage::Error { message } => bail!("{}", message),
273            other => bail!("unexpected response: {:?}", other),
274        }
275    }
276
277    pub async fn ws_switch(&self, id: WorkstreamId) -> Result<()> {
278        let resp = self.request(&ClientMessage::WsSwitch { id }).await?;
279        match resp {
280            ServerMessage::Ok => Ok(()),
281            ServerMessage::Error { message } => bail!("{}", message),
282            _ => Ok(()),
283        }
284    }
285
286    pub async fn ws_pause(&self, id: Option<WorkstreamId>) -> Result<()> {
287        let resp = self.request(&ClientMessage::WsPause { id }).await?;
288        match resp {
289            ServerMessage::Ok => Ok(()),
290            ServerMessage::Error { message } => bail!("{}", message),
291            _ => Ok(()),
292        }
293    }
294
295    pub async fn ws_resume(&self, id: Option<WorkstreamId>) -> Result<()> {
296        let resp = self.request(&ClientMessage::WsResume { id }).await?;
297        match resp {
298            ServerMessage::Ok => Ok(()),
299            ServerMessage::Error { message } => bail!("{}", message),
300            _ => Ok(()),
301        }
302    }
303
304    pub async fn ws_complete(&self, id: Option<WorkstreamId>) -> Result<()> {
305        let resp = self.request(&ClientMessage::WsComplete { id }).await?;
306        match resp {
307            ServerMessage::Ok => Ok(()),
308            ServerMessage::Error { message } => bail!("{}", message),
309            _ => Ok(()),
310        }
311    }
312
313    pub async fn ws_abandon(&self, id: Option<WorkstreamId>) -> Result<()> {
314        let resp = self.request(&ClientMessage::WsAbandon { id }).await?;
315        match resp {
316            ServerMessage::Ok => Ok(()),
317            ServerMessage::Error { message } => bail!("{}", message),
318            _ => Ok(()),
319        }
320    }
321
322    pub async fn ws_note(&self, id: WorkstreamId, text: &str) -> Result<()> {
323        let resp = self
324            .request(&ClientMessage::WsNote {
325                id,
326                text: text.to_string(),
327            })
328            .await?;
329        match resp {
330            ServerMessage::Ok => Ok(()),
331            ServerMessage::Error { message } => bail!("{}", message),
332            _ => Ok(()),
333        }
334    }
335
336    // ── Repo methods ────────────────────────────────────────────
337
338    pub async fn repo_list(&self) -> Result<Vec<RepoInfo>> {
339        let resp = self.request(&ClientMessage::RepoList).await?;
340        match resp {
341            ServerMessage::RepoList { repos } => Ok(repos),
342            ServerMessage::Error { message } => bail!("{}", message),
343            other => bail!("unexpected response: {:?}", other),
344        }
345    }
346
347    pub async fn repo_status(&self, path: Option<&str>) -> Result<RepoStatusInfo> {
348        let resp = self
349            .request(&ClientMessage::RepoStatus {
350                path: path.map(String::from),
351            })
352            .await?;
353        match resp {
354            ServerMessage::RepoStatus { status } => Ok(status),
355            ServerMessage::Error { message } => bail!("{}", message),
356            other => bail!("unexpected response: {:?}", other),
357        }
358    }
359
360    pub async fn repo_branches(&self, path: Option<&str>) -> Result<Vec<RepoBranchInfo>> {
361        let resp = self
362            .request(&ClientMessage::RepoBranches {
363                path: path.map(String::from),
364            })
365            .await?;
366        match resp {
367            ServerMessage::RepoBranches { branches } => Ok(branches),
368            ServerMessage::Error { message } => bail!("{}", message),
369            other => bail!("unexpected response: {:?}", other),
370        }
371    }
372
373    pub async fn repo_log(
374        &self,
375        path: Option<&str>,
376        last: Option<usize>,
377    ) -> Result<Vec<CommitInfo>> {
378        let resp = self
379            .request(&ClientMessage::RepoLog {
380                path: path.map(String::from),
381                last,
382            })
383            .await?;
384        match resp {
385            ServerMessage::RepoLog { commits } => Ok(commits),
386            ServerMessage::Error { message } => bail!("{}", message),
387            other => bail!("unexpected response: {:?}", other),
388        }
389    }
390
391    pub async fn repo_watch(&self, path: &str) -> Result<()> {
392        let resp = self
393            .request(&ClientMessage::RepoWatch {
394                path: path.to_string(),
395            })
396            .await?;
397        match resp {
398            ServerMessage::Ok => Ok(()),
399            ServerMessage::Error { message } => bail!("{}", message),
400            _ => Ok(()),
401        }
402    }
403
404    pub async fn repo_unwatch(&self, path: &str) -> Result<()> {
405        let resp = self
406            .request(&ClientMessage::RepoUnwatch {
407                path: path.to_string(),
408            })
409            .await?;
410        match resp {
411            ServerMessage::Ok => Ok(()),
412            ServerMessage::Error { message } => bail!("{}", message),
413            _ => Ok(()),
414        }
415    }
416
417    // ── Skill methods ───────────────────────────────────────────
418
419    pub async fn skill_list(&self) -> Result<Vec<SkillInfo>> {
420        let resp = self.request(&ClientMessage::SkillList).await?;
421        match resp {
422            ServerMessage::SkillList { skills } => Ok(skills),
423            ServerMessage::Error { message } => bail!("{}", message),
424            other => bail!("unexpected response: {:?}", other),
425        }
426    }
427
428    pub async fn skill_info(&self, name: &str) -> Result<SkillInfo> {
429        let resp = self
430            .request(&ClientMessage::SkillInfo {
431                name: name.to_string(),
432            })
433            .await?;
434        match resp {
435            ServerMessage::SkillDetail { skill } => Ok(skill),
436            ServerMessage::Error { message } => bail!("{}", message),
437            other => bail!("unexpected response: {:?}", other),
438        }
439    }
440
441    // ── GitHub methods ──────────────────────────────────────────
442
443    pub async fn gh_auth(&self) -> Result<(String, Vec<String>)> {
444        let resp = self.request(&ClientMessage::GhAuth).await?;
445        match resp {
446            ServerMessage::GhAuthResult { user, scopes } => Ok((user, scopes)),
447            ServerMessage::Error { message } => bail!("{}", message),
448            other => bail!("unexpected response: {:?}", other),
449        }
450    }
451
452    // ── Daemon methods ──────────────────────────────────────────
453
454    pub async fn daemon_status(&self) -> Result<DaemonStatusInfo> {
455        let resp = self.request(&ClientMessage::DaemonStatus).await?;
456        match resp {
457            ServerMessage::DaemonStatus { info } => Ok(info),
458            ServerMessage::Error { message } => bail!("{}", message),
459            other => bail!("unexpected response: {:?}", other),
460        }
461    }
462}