ashlar_control/
client.rs

1//! A client for a running app's control port. It implements [`Dispatch`], so anything that drives
2//! a headless session (the MCP server, the CLI) drives a live app the same way.
3
4use std::io::{BufReader, ErrorKind, Write};
5use std::net::{Ipv4Addr, SocketAddr, TcpStream};
6use std::path::{Path, PathBuf};
7use std::time::Duration;
8
9use ashlar_command::{CommandError, CommandInfo, Dispatch};
10use serde_json::{Value, json};
11
12use crate::ControlError;
13use crate::key::{key_path, read_key_file};
14use crate::lines::{Plain, ReadError, read_bounded_line};
15use crate::protocol::Reply;
16
17/// Longest reply line accepted (screenshots travel inline as base64).
18pub const MAX_REPLY: usize = 256 * 1024 * 1024;
19
20enum KeySource {
21    File(PathBuf),
22    Fixed(String),
23}
24
25struct Conn {
26    reader: BufReader<Plain<TcpStream>>,
27    writer: TcpStream,
28}
29
30/// A connection to `127.0.0.1:<port>` (the control port is loopback only, so the key never leaves
31/// the machine). Reconnects when the app closed the kept connection or restarted (a restart also
32/// means a new key, read again from the key file), but only before a request was sent: once any
33/// byte of a request reached the socket it is never sent again, so a mutating command cannot run
34/// twice when the app dies mid-request.
35pub struct Client {
36    port: u16,
37    key: KeySource,
38    conn: Option<Conn>,
39    next_id: u64,
40    timeout: Duration,
41}
42
43impl std::fmt::Debug for Client {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        f.debug_struct("Client").field("port", &self.port).field("connected", &self.conn.is_some()).finish_non_exhaustive()
46    }
47}
48
49impl Client {
50    /// Connect to the app on `port` whose key file is in `key_dir`; fails fast when nothing
51    /// listens or the key is refused.
52    pub fn connect(port: u16, key_dir: &Path) -> Result<Self, ControlError> {
53        Self::open(port, KeySource::File(key_path(key_dir, port)))
54    }
55
56    /// Connect with a key known in-process (tests, a parent that started the app).
57    pub fn connect_with_key(port: u16, key: impl Into<String>) -> Result<Self, ControlError> {
58        Self::open(port, KeySource::Fixed(key.into()))
59    }
60
61    fn open(port: u16, key: KeySource) -> Result<Self, ControlError> {
62        let mut c = Self { port, key, conn: None, next_id: 1, timeout: Duration::from_secs(90) };
63        c.reconnect()?;
64        Ok(c)
65    }
66
67    pub fn port(&self) -> u16 {
68        self.port
69    }
70
71    /// How long to wait for one reply (default 90 s; the app itself gives up after its own
72    /// reply timeout, 60 s by default).
73    pub fn set_timeout(&mut self, t: Duration) {
74        self.timeout = t;
75        if let Some(c) = &self.conn {
76            let _ = c.writer.set_read_timeout(Some(t));
77        }
78    }
79
80    fn key(&self) -> Result<String, ControlError> {
81        match &self.key {
82            KeySource::File(p) => read_key_file(p).map_err(ControlError::Key),
83            KeySource::Fixed(k) => Ok(k.clone()),
84        }
85    }
86
87    fn reconnect(&mut self) -> Result<(), ControlError> {
88        self.conn = None;
89        let key = self.key()?;
90        let addr = SocketAddr::from((Ipv4Addr::LOCALHOST, self.port));
91        let s = TcpStream::connect_timeout(&addr, Duration::from_millis(800))
92            .map_err(|e| ControlError::Io(format!("no app is listening on 127.0.0.1:{} ({e})", self.port)))?;
93        let _ = s.set_nodelay(true);
94        let _ = s.set_read_timeout(Some(self.timeout));
95        let _ = s.set_write_timeout(Some(Duration::from_secs(30)));
96        let read = s.try_clone().map_err(|e| ControlError::Io(e.to_string()))?;
97        self.conn = Some(Conn { reader: BufReader::new(Plain(read)), writer: s });
98        // Authenticate the connection once; later requests need not carry the key.
99        match self.roundtrip(&json!({"id": 0, "method": "auth", "key": key}))? {
100            Reply::Ok(_) => Ok(()),
101            Reply::Err { error, .. } => {
102                self.conn = None;
103                Err(if error == "unauthorized" { ControlError::Unauthorized } else { ControlError::Protocol(error) })
104            }
105        }
106    }
107
108    fn roundtrip(&mut self, req: &Value) -> Result<Reply, ControlError> {
109        match self.send(req) {
110            Ok(()) => self.receive(),
111            Err((_, e)) => Err(e),
112        }
113    }
114
115    /// Write one request line. `Err((sent_nothing, error))`: `sent_nothing` is true when no byte
116    /// of the request reached the socket, the only case that may be retried.
117    fn send(&mut self, req: &Value) -> Result<(), (bool, ControlError)> {
118        let port = self.port;
119        let Some(c) = self.conn.as_mut() else { return Err((true, ControlError::Io("not connected".into()))) };
120        let mut line = req.to_string().into_bytes();
121        line.push(b'\n');
122        let mut sent = 0usize;
123        while let Some(rest) = line.get(sent..).filter(|r| !r.is_empty()) {
124            match c.writer.write(rest) {
125                Ok(0) => return Err((sent == 0, ControlError::Io(format!("control connection to 127.0.0.1:{port} closed while sending")))),
126                Ok(n) => sent = sent.saturating_add(n),
127                Err(e) if e.kind() == ErrorKind::Interrupted => {}
128                Err(e) => return Err((sent == 0, ControlError::Io(format!("control connection to 127.0.0.1:{port} failed: {e}")))),
129            }
130        }
131        Ok(())
132    }
133
134    fn receive(&mut self) -> Result<Reply, ControlError> {
135        let Some(c) = self.conn.as_mut() else { return Err(ControlError::Io("not connected".into())) };
136        let mut buf = Vec::new();
137        match read_bounded_line(&mut c.reader, &mut buf, MAX_REPLY) {
138            Ok(true) => {}
139            Ok(false) => return Err(ControlError::Io("the app closed the control connection".into())),
140            Err(ReadError::TooLong) => return Err(ControlError::Protocol(format!("reply longer than {MAX_REPLY} bytes"))),
141            Err(ReadError::Io) => return Err(ControlError::Io("no reply from the app (timed out or connection lost)".into())),
142        }
143        let v: Value = serde_json::from_slice(&buf).map_err(|e| ControlError::Protocol(e.to_string()))?;
144        Reply::from_json(&v).ok_or_else(|| ControlError::Protocol("reply without a boolean `ok`".into()))
145    }
146
147    /// Whether the kept connection is known dead: the app closed it (it quit or restarted) or sent
148    /// something unasked. Checked without blocking, before a request is written.
149    fn stale(&self) -> bool {
150        let Some(c) = self.conn.as_ref() else { return true };
151        if !c.reader.buffer().is_empty() || c.writer.set_nonblocking(true).is_err() {
152            return true;
153        }
154        let mut probe = [0u8; 1];
155        let alive = matches!(c.writer.peek(&mut probe), Err(e) if e.kind() == ErrorKind::WouldBlock);
156        c.writer.set_nonblocking(false).is_err() || !alive
157    }
158
159    /// Send one request and return the app's reply; `Err` only for transport failures.
160    pub fn request(&mut self, method: &str, params: Value) -> Result<Reply, ControlError> {
161        let id = self.next_id;
162        self.next_id = self.next_id.wrapping_add(1);
163        let req = json!({"id": id, "method": method, "params": params});
164        // A kept connection the app closed (it quit or restarted) is replaced before sending.
165        if self.stale() {
166            self.reconnect()?;
167        }
168        let r = match self.send(&req) {
169            Ok(()) => self.receive(),
170            // Nothing was sent: the app cannot have seen the request, so one retry is safe.
171            Err((true, _)) => self.reconnect().and_then(|()| self.roundtrip(&req)),
172            // Part of the request may have reached the app: never send it again.
173            Err((false, e)) => Err(e),
174        };
175        r.inspect_err(|_| self.conn = None)
176    }
177
178    /// Call one method: its `result`, or the app's error as [`ControlError::Remote`].
179    pub fn call(&mut self, method: &str, params: Value) -> Result<Value, ControlError> {
180        match self.request(method, params)? {
181            Reply::Ok(v) => Ok(v),
182            Reply::Err { error, detail } => Err(ControlError::Remote { error, detail }),
183        }
184    }
185}
186
187impl Dispatch for Client {
188    /// The app's commands; empty when it cannot be reached (`Dispatch` has no error channel here;
189    /// use [`Client::call`] with `engine.commands` to see why).
190    fn commands(&self) -> Vec<CommandInfo> {
191        // `commands` takes `&self`: talk on a separate short-lived connection.
192        let key = match self.key() {
193            Ok(k) => k,
194            Err(_) => return Vec::new(),
195        };
196        let Ok(mut c) = Client::connect_with_key(self.port, key) else { return Vec::new() };
197        c.timeout = self.timeout;
198        c.call("engine.commands", json!({})).ok().and_then(|v| serde_json::from_value(v).ok()).unwrap_or_default()
199    }
200
201    fn execute(&mut self, id: &str, params: Value) -> Result<Value, CommandError> {
202        self.call("engine.execute", json!({"command": id, "params": params})).map_err(|e| e.into_command_error(id))
203    }
204}
205
206impl ControlError {
207    /// As a [`CommandError`]: the app's own when it sent one, otherwise [`CommandError::Failed`].
208    pub fn into_command_error(self, id: &str) -> CommandError {
209        match self {
210            ControlError::Remote { detail: Some(d), error } => {
211                serde_json::from_value(d).unwrap_or(CommandError::Failed { id: id.to_string(), msg: error })
212            }
213            e => CommandError::Failed { id: id.to_string(), msg: e.to_string() },
214        }
215    }
216}
ashlar_control/client.rs

ashlar_control/
client.rs
