prospero_core/caliband/
client.rs

1//! Thin async client for a single caliband control socket.
2//!
3//! One request → one reply per connection (matching caliban's protocol): each
4//! call opens the Unix socket, writes one NDJSON request frame, reads one reply
5//! frame, and closes. Cheap for a local control plane and trivially debuggable.
6
7use std::path::PathBuf;
8
9use tokio::io::{AsyncWriteExt, BufReader};
10
11use crate::caliband::transport::{self, BoxConn, ConnectSpec, TlsClient};
12use crate::caliband::wire::{
13    AgentRecord, AttachInbound, CtlReply, CtlRequest, DaemonStatus, Endpoint, SpawnSpec,
14};
15use crate::caliband::{read_frame, write_frame};
16use crate::error::{CoreError, Result};
17
18/// Display form of an endpoint for error messages.
19fn endpoint_display(ep: &Endpoint) -> String {
20    match ep {
21        Endpoint::Unix { path } => path.display().to_string(),
22        Endpoint::Tcp { addr } => format!("tcp://{addr}"),
23    }
24}
25
26/// A client bound to one caliband control endpoint (Unix, or TCP+TLS+token).
27#[derive(Clone)]
28pub struct CalibandClient {
29    endpoint: Endpoint,
30    tls: Option<TlsClient>,
31    token: Option<String>,
32}
33
34impl CalibandClient {
35    /// Create a client for a local Unix control socket (credential-free default).
36    pub fn new(socket_path: impl Into<PathBuf>) -> Self {
37        Self {
38            endpoint: Endpoint::Unix {
39                path: socket_path.into(),
40            },
41            tls: None,
42            token: None,
43        }
44    }
45
46    /// Create a client that dials a TCP control endpoint over TLS + bearer token
47    /// (ADR 0051). `tls`/`token` are `None` only in plaintext/no-auth test setups.
48    pub fn connect_tcp(
49        addr: impl Into<String>,
50        tls: Option<TlsClient>,
51        token: Option<String>,
52    ) -> Self {
53        Self {
54            endpoint: Endpoint::Tcp { addr: addr.into() },
55            tls,
56            token,
57        }
58    }
59
60    /// The control endpoint this client targets.
61    pub fn endpoint(&self) -> &Endpoint {
62        &self.endpoint
63    }
64
65    /// Connect to the control endpoint, mapping connection failures to the
66    /// `CalibandUnreachable` error so callers can degrade the repo to
67    /// `Unreachable` rather than treating it as fatal.
68    async fn connect(&self) -> Result<BoxConn> {
69        transport::connect(&ConnectSpec {
70            endpoint: self.endpoint.clone(),
71            tls: self.tls.clone(),
72            token: self.token.clone(),
73        })
74        .await
75        .map_err(|source| CoreError::CalibandUnreachable {
76            endpoint: endpoint_display(&self.endpoint),
77            source,
78        })
79    }
80
81    /// Send one request and return the raw reply, surfacing `Error` replies as
82    /// typed [`CoreError`]s.
83    pub async fn request(&self, req: &CtlRequest) -> Result<CtlReply> {
84        let conn = self.connect().await?;
85        let (read_half, mut write_half) = tokio::io::split(conn);
86        write_frame(&mut write_half, req).await?;
87        let mut reader = BufReader::new(read_half);
88        let reply: CtlReply = read_frame(&mut reader).await?;
89        if let CtlReply::Error { error } = reply {
90            return Err(error.into());
91        }
92        Ok(reply)
93    }
94
95    /// List all agents registered with this daemon.
96    pub async fn list(&self) -> Result<Vec<AgentRecord>> {
97        match self.request(&CtlRequest::List).await? {
98            CtlReply::Listed { agents } => Ok(agents),
99            other => Err(unexpected("list", other)),
100        }
101    }
102
103    /// Spawn a new agent; returns `(id, per-agent endpoint)`.
104    pub async fn spawn(&self, spec: SpawnSpec) -> Result<(String, Endpoint)> {
105        match self.request(&CtlRequest::Spawn { spec }).await? {
106            CtlReply::Spawned { id, endpoint } => Ok((id, endpoint)),
107            other => Err(unexpected("spawn", other)),
108        }
109    }
110
111    /// Resolve an agent's per-agent endpoint for attaching.
112    pub async fn attach(&self, id: impl Into<String>) -> Result<Endpoint> {
113        match self.request(&CtlRequest::Attach { id: id.into() }).await? {
114            CtlReply::AttachAck { endpoint } => Ok(endpoint),
115            other => Err(unexpected("attach", other)),
116        }
117    }
118
119    /// Kill an agent.
120    pub async fn kill(&self, id: impl Into<String>) -> Result<()> {
121        match self.request(&CtlRequest::Kill { id: id.into() }).await? {
122            CtlReply::Killed => Ok(()),
123            other => Err(unexpected("kill", other)),
124        }
125    }
126
127    /// Respawn an agent; returns the new id.
128    pub async fn respawn(&self, id: impl Into<String>) -> Result<String> {
129        match self.request(&CtlRequest::Respawn { id: id.into() }).await? {
130            CtlReply::Respawned { id } => Ok(id),
131            other => Err(unexpected("respawn", other)),
132        }
133    }
134
135    /// Remove an agent from the registry.
136    pub async fn rm(&self, id: impl Into<String>, force: bool) -> Result<()> {
137        match self
138            .request(&CtlRequest::Rm {
139                id: id.into(),
140                force,
141            })
142            .await?
143        {
144            CtlReply::Removed => Ok(()),
145            other => Err(unexpected("rm", other)),
146        }
147    }
148
149    /// Probe daemon status.
150    pub async fn status(&self) -> Result<DaemonStatus> {
151        match self.request(&CtlRequest::Status).await? {
152            CtlReply::Status(s) => Ok(s),
153            other => Err(unexpected("status", other)),
154        }
155    }
156
157    /// Ask the daemon to drain and shut down.
158    pub async fn shutdown(&self) -> Result<()> {
159        match self.request(&CtlRequest::Shutdown).await? {
160            CtlReply::ShutdownAck => Ok(()),
161            other => Err(unexpected("shutdown", other)),
162        }
163    }
164
165    /// Open a streaming reader over a per-agent endpoint (from [`Self::attach`]),
166    /// dialing with this client's TLS/token. Lines read from this reader are
167    /// caliban stream-json frames.
168    pub async fn open_stream(&self, endpoint: &Endpoint) -> Result<BufReader<BoxConn>> {
169        let conn = transport::connect(&ConnectSpec {
170            endpoint: endpoint.clone(),
171            tls: self.tls.clone(),
172            token: self.token.clone(),
173        })
174        .await
175        .map_err(|source| CoreError::CalibandUnreachable {
176            endpoint: endpoint_display(endpoint),
177            source,
178        })?;
179        Ok(BufReader::new(conn))
180    }
181
182    /// Write a single inbound control frame to an interactive agent's per-agent
183    /// endpoint (from [`Self::attach`]). Opens a fresh write-only connection,
184    /// matching caliban's "all attach connections feed a shared inbox" model.
185    pub async fn send_inbound(&self, endpoint: &Endpoint, frame: &AttachInbound) -> Result<()> {
186        let mut conn = transport::connect(&ConnectSpec {
187            endpoint: endpoint.clone(),
188            tls: self.tls.clone(),
189            token: self.token.clone(),
190        })
191        .await
192        .map_err(|source| CoreError::CalibandUnreachable {
193            endpoint: endpoint_display(endpoint),
194            source,
195        })?;
196        let mut line = serde_json::to_vec(frame)?;
197        line.push(b'\n');
198        conn.write_all(&line)
199            .await
200            .map_err(|source| CoreError::CalibandUnreachable {
201                endpoint: endpoint_display(endpoint),
202                source,
203            })?;
204        // Flush explicitly: a TLS `BoxConn` wraps a userspace buffer, unlike the
205        // old raw UnixStream, so write_all alone may not reach the kernel.
206        conn.flush()
207            .await
208            .map_err(|source| CoreError::CalibandUnreachable {
209                endpoint: endpoint_display(endpoint),
210                source,
211            })
212    }
213}
214
215fn unexpected(op: &str, reply: CtlReply) -> CoreError {
216    CoreError::Protocol(format!("unexpected reply to {op}: {reply:?}"))
217}
218
219#[cfg(test)]
220mod tests {
221    use super::*;
222    use crate::caliband::wire::AttachInbound;
223    use tokio::io::{AsyncBufReadExt, BufReader};
224    use tokio::net::UnixListener;
225
226    #[tokio::test]
227    async fn send_inbound_writes_one_ndjson_frame() {
228        let dir = tempfile::tempdir().unwrap();
229        let sock = dir.path().join("a.sock");
230        let listener = UnixListener::bind(&sock).unwrap();
231        let server = tokio::spawn(async move {
232            let (stream, _) = listener.accept().await.unwrap();
233            let mut line = String::new();
234            BufReader::new(stream).read_line(&mut line).await.unwrap();
235            line
236        });
237        let client = CalibandClient::new(&sock);
238        client
239            .send_inbound(
240                &Endpoint::Unix { path: sock.clone() },
241                &AttachInbound::UserMessage { text: "go".into() },
242            )
243            .await
244            .unwrap();
245        assert_eq!(
246            server.await.unwrap().trim_end(),
247            r#"{"type":"UserMessage","text":"go"}"#
248        );
249    }
250
251    fn test_spec() -> SpawnSpec {
252        SpawnSpec {
253            label: None,
254            frontmatter_path: None,
255            initial_prompt: "hi".into(),
256            model: None,
257            provider: None,
258            tool_allowlist: None,
259            isolation_worktree: false,
260            inherit_hooks: true,
261            interactive: false,
262        }
263    }
264
265    #[tokio::test]
266    async fn client_round_trips_control_requests() {
267        use crate::testkit::FakeCaliband;
268        let dir = tempfile::tempdir().unwrap();
269        let sock = dir.path().join("ctl.sock");
270        let mut fake = FakeCaliband::start_at(&sock).await.unwrap();
271        let client = CalibandClient::new(&sock);
272        assert_eq!(client.endpoint(), &Endpoint::Unix { path: sock.clone() });
273
274        let (id, _endpoint) = client.spawn(test_spec()).await.unwrap();
275        assert!(client.list().await.unwrap().iter().any(|a| a.id == id));
276        let _ = client.attach(&id).await.unwrap();
277        assert!(client.status().await.unwrap().agents >= 1);
278        client.kill(&id).await.unwrap();
279
280        let (id2, _) = client.spawn(test_spec()).await.unwrap();
281        assert!(!client.respawn(&id2).await.unwrap().is_empty());
282
283        let (id3, _) = client.spawn(test_spec()).await.unwrap();
284        client.rm(&id3, true).await.unwrap();
285
286        // Error-reply path: an unknown id maps to AgentNotFound.
287        assert!(matches!(
288            client.kill("nope").await.unwrap_err(),
289            CoreError::AgentNotFound(_)
290        ));
291
292        client.shutdown().await.unwrap();
293        let _ = &mut fake;
294    }
295
296    #[tokio::test]
297    async fn connect_error_maps_to_unreachable() {
298        let client = CalibandClient::new("/nonexistent/dir/ctl.sock");
299        assert!(matches!(
300            client.list().await.unwrap_err(),
301            CoreError::CalibandUnreachable { .. }
302        ));
303    }
304
305    #[tokio::test]
306    async fn client_round_trips_over_tcp_tls_token() {
307        use crate::testkit::FakeCaliband;
308        let (fake, fixture) = FakeCaliband::start_tcp_tls("s3cr3t").await.unwrap();
309        let tls =
310            crate::caliband::transport::tls_client_from_pem(&fixture.ca_pem, "localhost").unwrap();
311        let client =
312            CalibandClient::connect_tcp(fixture.addr.clone(), Some(tls), Some("s3cr3t".into()));
313        assert!(matches!(client.endpoint(), Endpoint::Tcp { .. }));
314
315        let (id, _ep) = client.spawn(test_spec()).await.unwrap();
316        assert!(client.list().await.unwrap().iter().any(|a| a.id == id));
317        let _ = client.attach(&id).await.unwrap();
318        assert!(client.status().await.unwrap().agents >= 1);
319        client.kill(&id).await.unwrap();
320        client.shutdown().await.unwrap();
321        let _ = fake;
322    }
323
324    #[tokio::test]
325    async fn client_rejects_bad_token_over_tcp() {
326        use crate::testkit::FakeCaliband;
327        let (fake, fixture) = FakeCaliband::start_tcp_tls("right").await.unwrap();
328        let tls =
329            crate::caliband::transport::tls_client_from_pem(&fixture.ca_pem, "localhost").unwrap();
330        let client =
331            CalibandClient::connect_tcp(fixture.addr.clone(), Some(tls), Some("wrong".into()));
332        // The server rejects the token at accept-time and drops the connection,
333        // which surfaces on the client as a failed dial (CalibandUnreachable), a
334        // clean EOF before any reply (Protocol), or a mid-write connection reset
335        // over TLS (Io) — any of which correctly denies the bad token.
336        assert!(matches!(
337            client.list().await.unwrap_err(),
338            CoreError::CalibandUnreachable { .. } | CoreError::Protocol(_) | CoreError::Io(_)
339        ));
340        let _ = fake;
341    }
342}
prospero_core/caliband/client.rs

prospero_core/caliband/
client.rs
