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