Skip to main content

forge_protocol_client/
forge-protocol-client.rs

1//! The protocol client: the instrument that speaks the socket from the other
2//! side of it.
3//!
4//! It is what `curl` cannot be for a socket. It connects to a running forge,
5//! sends a script of messages as a client would, and prints everything the
6//! server says back as one line of JSON each - so a person debugging a client
7//! has both halves of the conversation in front of them rather than a guess
8//! about which end is wrong.
9//!
10//! ```text
11//! forge-protocol-client --url ws://127.0.0.1:8790/socket --script ./walk.txt
12//! ```
13//!
14//! A script is the protocol in miniature rather than a vocabulary of its own:
15//! a verb and its arguments, and raw JSON where a message carries a payload.
16//!
17//! ```text
18//! # a comment
19//! subscribe home
20//! subscribe session Busytools/forge/lead
21//! more Busytools/forge/lead 10
22//! commit_diff Busytools/forge/lead 4f6870839
23//! file_content Busytools/forge/lead crates/forge-server/src/lib.rs
24//! file_content Busytools/forge/lead crates/forge-server/src/lib.rs 4f6870839
25//! command {"prompt": {"key": {"org": "Busytools", "project": "forge", "label": "lead"}, "text": "hello", "attachments": []}}
26//! ask {"submit_review": {"project": "forge", "branch": "main", "summary": null, "thread_ids": [], "origin": {"org": "Busytools", "project": "forge", "label": "lead"}}}
27//! ```
28//!
29//! `ask` is `command` with a reply asked for, so the awaited `Reply` path is
30//! exercised rather than only the fire-and-forget one.
31
32use 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    /// The socket to speak to, e.g. `ws://127.0.0.1:8790/socket`.
45    #[arg(long)]
46    url: String,
47    /// A file of script lines, one per line; blank lines and `#` are skipped.
48    #[arg(long)]
49    script: PathBuf,
50    /// How long to keep reading after the script has been sent. The client
51    /// exits at this bound whether or not the server has stopped talking, so
52    /// a run against a quiet server ends rather than hanging.
53    #[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    // The greeting first, so a reader sees the server speak before it is
69    // asked anything - and sees WHAT answered, which is the whole reason the
70    // greeting exists.
71    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    // One deadline for the whole run rather than one per read, so a chatty
83    // server cannot keep this alive past the bound the caller asked for.
84    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            // A ping or a pong carries nothing to say, and the socket is
97            // still open.
98            _ => {}
99        }
100    }
101    Ok(())
102}
103
104/// The socket's next message, or `None` once it stops arriving.
105async 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
114/// One script line as a message to send.
115///
116/// The lines name what the protocol has rather than inventing a second
117/// vocabulary: a verb and its subject, and raw JSON where a message carries a
118/// payload, so the instrument sends what a client would send.
119fn 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        // This instrument drives the socket the way a view does, so it
124        // declares itself able to answer: a client that cannot show a prompt
125        // must not be counted as one, but this one walks the reply path.
126        "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 <subject> <path> [<sha>]`: the path is one word (a
155        // path with a space in it is not a line this instrument can carry),
156        // and the sha is the scope's revision - absent for the branch tip.
157        "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
189/// `home`, `board`, `usage`, `dictate_models`, `extensions`, or
190/// `session <org>/<project>/<label>`.
191fn 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}