forge_protocol_client/
forge-protocol-client.rs1use std::path::PathBuf;
33use std::time::Duration;
34
35use clap::Parser;
36use forge_primitives::SessionSlot;
37use forge_server::Command;
38use forge_server::transport::envelope::{ClientMessage, Subject};
39use futures_util::{SinkExt, StreamExt};
40use tokio_tungstenite::tungstenite::Message;
41
42#[derive(Parser)]
43struct Args {
44 #[arg(long)]
46 url: String,
47 #[arg(long)]
49 script: PathBuf,
50 #[arg(long, default_value_t = 10)]
54 seconds: u64,
55}
56
57#[tokio::main]
58async fn main() -> anyhow::Result<()> {
59 let args = Args::parse();
60 let script = std::fs::read_to_string(&args.script)?;
61 let lines: Vec<&str> = script
62 .lines()
63 .map(str::trim)
64 .filter(|line| !line.is_empty() && !line.starts_with('#'))
65 .collect();
66
67 let (mut socket, _) = tokio_tungstenite::connect_async(&args.url).await?;
68 if let Some(Ok(Message::Text(greeting))) = next_within(&mut socket, args.seconds).await {
72 println!("{greeting}");
73 }
74
75 let mut reply_to = 1_u64;
76 for line in lines {
77 let message = parse_line(line, &mut reply_to)?;
78 let text = serde_json::to_string(&message)?;
79 socket.send(Message::Text(text.into())).await?;
80 }
81
82 let deadline = tokio::time::Instant::now() + Duration::from_secs(args.seconds);
85 loop {
86 let left = deadline.saturating_duration_since(tokio::time::Instant::now());
87 if left.is_zero() {
88 break;
89 }
90 let Some(Ok(msg)) = next_within(&mut socket, left.as_secs()).await else {
91 break;
92 };
93 match msg {
94 Message::Text(text) => println!("{text}"),
95 Message::Close(_) => break,
96 _ => {}
99 }
100 }
101 Ok(())
102}
103
104async fn next_within(
106 socket: &mut tokio_tungstenite::WebSocketStream<
107 tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>,
108 >,
109 seconds: u64,
110) -> Option<Result<Message, tokio_tungstenite::tungstenite::Error>> {
111 tokio::time::timeout(Duration::from_secs(seconds), socket.next()).await.ok().flatten()
112}
113
114fn parse_line(line: &str, reply_to: &mut u64) -> anyhow::Result<ClientMessage> {
120 let (verb, rest) = line.split_once(char::is_whitespace).unwrap_or((line, ""));
121 let rest = rest.trim();
122 match verb {
123 "subscribe" => {
127 Ok(ClientMessage::Subscribe { what: subject(rest)?, answering: true, review: false })
128 }
129 "unsubscribe" => Ok(ClientMessage::Unsubscribe { what: subject(rest)? }),
130 "more" => {
131 let mut parts = rest.split_whitespace();
132 let Some(seat) = parts.next() else {
133 anyhow::bail!("`more` wants a session subject");
134 };
135 let Subject::Session(conversation) = subject(&format!("session {seat}"))? else {
136 anyhow::bail!("`more` wants a session subject");
137 };
138 let turns = parts.next().map_or(Ok(10), str::parse)?;
139 Ok(ClientMessage::More { conversation, before: None, turns })
140 }
141 "commit_diff" => {
142 let mut parts = rest.split_whitespace();
143 let Some(seat) = parts.next() else {
144 anyhow::bail!("`commit_diff` wants a session subject and a sha");
145 };
146 let Subject::Session(session) = subject(&format!("session {seat}"))? else {
147 anyhow::bail!("`commit_diff` wants a session subject");
148 };
149 let Some(sha) = parts.next() else {
150 anyhow::bail!("`commit_diff` wants a sha");
151 };
152 Ok(ClientMessage::CommitDiff { session, sha: sha.to_owned() })
153 }
154 "file_content" => {
158 let mut parts = rest.split_whitespace();
159 let Some(seat) = parts.next() else {
160 anyhow::bail!("`file_content` wants a session subject and a path");
161 };
162 let Subject::Session(session) = subject(&format!("session {seat}"))? else {
163 anyhow::bail!("`file_content` wants a session subject");
164 };
165 let Some(path) = parts.next() else {
166 anyhow::bail!("`file_content` wants a path");
167 };
168 Ok(ClientMessage::FileContent {
169 session,
170 path: path.to_owned(),
171 sha: parts.next().map(str::to_owned),
172 })
173 }
174 "command" | "ask" => {
175 let command: Command = serde_json::from_str(rest)?;
176 let asked = if verb == "ask" {
177 let asked = Some(*reply_to);
178 *reply_to += 1;
179 asked
180 } else {
181 None
182 };
183 Ok(ClientMessage::Command { command: Box::new(command), reply_to: asked })
184 }
185 other => anyhow::bail!("`{other}` is not a verb this client knows"),
186 }
187}
188
189fn subject(rest: &str) -> anyhow::Result<Subject> {
192 if rest == "home" {
193 return Ok(Subject::Home);
194 }
195 if rest == "board" {
196 return Ok(Subject::Board);
197 }
198 if rest == "usage" {
199 return Ok(Subject::Usage);
200 }
201 if rest == "dictate_models" {
202 return Ok(Subject::DictateModels);
203 }
204 if rest == "extensions" {
205 return Ok(Subject::Extensions);
206 }
207 let Some(seat) = rest.strip_prefix("session ") else {
208 anyhow::bail!(
209 "a subject is `home`, `board`, `usage`, `dictate_models`, `extensions` or \
210 `session <org>/<project>/<label>`, not `{rest}`"
211 );
212 };
213 let mut parts = seat.split('/');
214 let (Some(org), Some(project), Some(label), None) =
215 (parts.next(), parts.next(), parts.next(), parts.next())
216 else {
217 anyhow::bail!("a session subject is <org>/<project>/<label>, not `{seat}`");
218 };
219 Ok(Subject::Session(SessionSlot::for_label(org, project, Some(label))))
220}