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
10pub 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 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 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 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 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 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 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 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 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 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}