lattice_ai/acp/connection.rs
1//! ACP connection adapter over the `agent-client-protocol` crate.
2//!
3//! The crate frames JSON-RPC over stdio itself (newline-delimited JSON) and drives a
4//! connection through a closure-based API: `Client.builder()...connect_with(transport,
5//! async |cx| { ... })`. That closure owns the connection for as long as it runs — once it
6//! returns, the whole connection (including its background dispatch loop) shuts down. That
7//! doesn't fit the shape Tasks 4/5/7 need: a `Connection` handle with independent async
8//! methods (`initialize`, `new_session`, `prompt`) callable at any time from any task.
9//!
10//! This module bridges the two shapes: [`Connection::spawn`] starts a background "driver"
11//! task that runs the crate's `connect_with` closure as a command loop for the connection's
12//! whole lifetime, and the public methods send commands into that loop over a channel and
13//! await the reply.
14//!
15//! ## Threading
16//!
17//! `agent-client-protocol` builds its connection state on `futures::channel::mpsc` (not
18//! tokio) and requires spawned work to be `Send + 'static` (see the crate's internal
19//! `Task::new` bound). That makes the whole stack executor-agnostic and `Send`, so the driver
20//! task runs on a plain `tokio::spawn` — no dedicated thread or `LocalSet` is required.
21//!
22//! `Connection::spawn` takes generic tokio `AsyncRead`/`AsyncWrite` halves and adapts them to
23//! the `futures::io` traits the crate's `ByteStreams` transport expects via
24//! `tokio_util::compat`.
25//!
26//! ## Notifications
27//!
28//! `session/update` notifications are handled once, connection-wide, via a single
29//! `on_receive_notification` handler registered on the builder (the crate does not scope
30//! notification handlers per-session at this layer). Every notification the agent sends for
31//! the life of the connection is forwarded to the `mpsc::UnboundedReceiver<SessionNotification>`
32//! returned by [`Connection::spawn`]; callers that run multiple sessions on one connection
33//! must filter by `SessionNotification::session_id` themselves. The channel is unbounded (not
34//! bounded) because `on_receive_notification` runs inside the crate's single dispatch loop —
35//! see the comment at the channel construction site in [`Connection::spawn`] for why a bounded
36//! channel would risk deadlocking every in-flight request.
37
38use std::sync::Arc;
39
40use agent_client_protocol::schema::ProtocolVersion;
41use agent_client_protocol::schema::v1::{
42 CancelNotification, ContentBlock, InitializeRequest, NewSessionRequest, PromptRequest,
43 RequestPermissionOutcome, RequestPermissionRequest, RequestPermissionResponse, TextContent,
44};
45use agent_client_protocol::{ByteStreams, Client, Responder};
46use tokio::io::{AsyncRead, AsyncWrite};
47use tokio::sync::{mpsc, oneshot};
48use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
49
50use crate::acp::error::{AiError, Result};
51
52/// Re-exported so the supervisor (AU‑4) can inspect the permission request +
53/// build the response without depending on `agent-client-protocol` directly.
54pub use agent_client_protocol::schema::v1::RequestPermissionRequest as PermissionRequestPayload;
55/// Re-exported so callers (Task 6) can match on `session/update` payloads without depending
56/// on `agent-client-protocol` directly.
57pub use agent_client_protocol::schema::v1::SessionNotification;
58
59/// AU‑4: an agent→client `session/request_permission` request routed out to the
60/// supervisor, carrying the [`Responder`] it must answer. The supervisor
61/// classifies the tool call (auto-allow reads / review file edits) and calls
62/// `responder.respond(...)` — possibly much later, after a diff verdict. The
63/// [`Responder`] is `Send + 'static` (its `send_fn` is a boxed `FnOnce + Send`),
64/// so answering off the connection's dispatch loop is sound.
65pub struct PermissionRequest {
66 pub request: RequestPermissionRequest,
67 pub responder: Responder<RequestPermissionResponse>,
68}
69
70/// A lattice-local session identifier.
71///
72/// Kept distinct from `agent_client_protocol::schema::v1::SessionId` (which wraps an
73/// `Arc<str>` and carries protocol-schema derives) so the rest of `lattice-ai` doesn't need
74/// to depend on the wire type directly.
75#[derive(Debug, Clone, PartialEq, Eq, Hash)]
76pub struct SessionId(pub String);
77
78/// A command sent from a `Connection`'s public async methods to the driver task.
79enum DriverCommand {
80 Initialize {
81 reply: oneshot::Sender<Result<()>>,
82 },
83 NewSession {
84 cwd: String,
85 reply: oneshot::Sender<Result<SessionId>>,
86 },
87 Prompt {
88 session_id: SessionId,
89 text: String,
90 reply: oneshot::Sender<Result<()>>,
91 },
92 /// AU‑3: interrupt the active turn without ending the session. Sent by
93 /// [`Connection::cancel`]; the driver forwards an ACP `session/cancel`
94 /// notification. Handled concurrently with an in-flight `Prompt` (which
95 /// the driver runs via `cx.spawn`), so it can actually interrupt.
96 Cancel {
97 session_id: SessionId,
98 reply: oneshot::Sender<Result<()>>,
99 },
100}
101
102/// A handle to a live ACP connection driven by the `agent-client-protocol` crate.
103///
104/// Constructed with [`Connection::spawn`], which owns the transport and runs the crate's
105/// JSON-RPC dispatch loop on a background task for the connection's whole lifetime.
106///
107/// ## Single-flight, single-session limitation
108///
109/// The driver's command loop (spawned by [`Connection::spawn`]) awaits the reply to
110/// `initialize` / `new_session` before pulling the next command off the channel. `prompt`
111/// is the exception: the loop runs the prompt turn on a `cx.spawn`ed task (AU‑3) so a
112/// concurrent `Cancel` (`session/cancel`) can actually interrupt it — a `block_task().await`
113/// inline would pin the loop for the whole turn and starve the very cancel meant to stop it.
114/// Notifications are connection-wide, not scoped to a session (see the module-level docs). A
115/// future multi-session `Connection` needs to (a) filter delivered notifications by
116/// `SessionNotification::session_id` and (b) generalise the per-command spawning so one
117/// session's in-flight request never blocks another session's commands.
118pub struct Connection {
119 commands: mpsc::Sender<DriverCommand>,
120}
121
122impl Connection {
123 /// Spawn a driver task that adapts `reader`/`writer` into an ACP transport and drives
124 /// the `agent-client-protocol` client connection over it.
125 ///
126 /// Returns a `Connection` handle plus two receivers: one that yields every
127 /// `session/update` notification, and one (AU‑4) that yields every
128 /// agent→client `session/request_permission` request the supervisor must
129 /// answer.
130 pub fn spawn<R, W>(
131 reader: R,
132 writer: W,
133 ) -> (
134 Arc<Connection>,
135 mpsc::UnboundedReceiver<SessionNotification>,
136 mpsc::UnboundedReceiver<PermissionRequest>,
137 )
138 where
139 R: AsyncRead + Unpin + Send + 'static,
140 W: AsyncWrite + Unpin + Send + 'static,
141 {
142 let (cmd_tx, mut cmd_rx) = mpsc::channel::<DriverCommand>(32);
143 // Unbounded, not bounded: `on_receive_notification` below runs *inside* the crate's
144 // single dispatch loop, which also routes responses to in-flight `initialize` /
145 // `new_session` / `prompt` calls. A bounded channel's `.send().await` would block that
146 // loop (and therefore every pending/future request) if the notification consumer ever
147 // falls behind by a full channel's worth of messages — very plausible during
148 // token-by-token `session/update` streaming. Spawning a task per notification instead
149 // would avoid blocking but let streaming chunks complete out of order, corrupting
150 // assistant message text; unbounded `send` is both non-blocking and order-preserving.
151 let (notif_tx, notif_rx) = mpsc::unbounded_channel::<SessionNotification>();
152 // AU‑4: agent→client `session/request_permission` requests. Same
153 // unbounded, order-preserving, synchronous-send rationale as the
154 // notification channel: the `on_receive_request` handler runs inside the
155 // dispatch loop and must not await. The supervisor answers each via the
156 // carried `Responder`.
157 let (perm_tx, perm_rx) = mpsc::unbounded_channel::<PermissionRequest>();
158
159 tokio::spawn(async move {
160 let transport = ByteStreams::new(writer.compat_write(), reader.compat());
161
162 let result = Client
163 .builder()
164 .on_receive_notification(
165 async move |notification: SessionNotification, _cx| {
166 // Synchronous, non-blocking send on the unbounded channel: this
167 // handler runs inside the crate's single dispatch loop, so it must
168 // never await here (see the channel-construction comment above).
169 // Ignore send errors: if the caller dropped the receiver they've
170 // opted out of notifications, not out of the connection.
171 let _ = notif_tx.send(notification);
172 Ok(())
173 },
174 agent_client_protocol::on_receive_notification!(),
175 )
176 .on_receive_request(
177 async move |request: RequestPermissionRequest, responder, _cx| {
178 // Route to the supervisor (never await in-loop). If the
179 // receiver is gone we MUST still answer, or the agent hangs
180 // its turn forever — respond `Cancelled` (the protocol's
181 // "no decision" outcome).
182 match perm_tx.send(PermissionRequest { request, responder }) {
183 Ok(()) => Ok(()),
184 Err(mpsc::error::SendError(PermissionRequest {
185 responder, ..
186 })) => responder.respond(RequestPermissionResponse::new(
187 RequestPermissionOutcome::Cancelled,
188 )),
189 }
190 },
191 agent_client_protocol::on_receive_request!(),
192 )
193 .connect_with(transport, async move |cx| {
194 while let Some(cmd) = cmd_rx.recv().await {
195 match cmd {
196 DriverCommand::Initialize { reply } => {
197 let outcome = cx
198 .send_request(InitializeRequest::new(ProtocolVersion::V1))
199 .block_task()
200 .await
201 .map(|_response| ())
202 .map_err(map_acp_error);
203 let _ = reply.send(outcome);
204 }
205 DriverCommand::NewSession { cwd, reply } => {
206 let outcome = cx
207 .send_request(NewSessionRequest::new(cwd))
208 .block_task()
209 .await
210 .map(|response| SessionId(response.session_id.0.to_string()))
211 .map_err(map_acp_error);
212 let _ = reply.send(outcome);
213 }
214 DriverCommand::Prompt {
215 session_id,
216 text,
217 reply,
218 } => {
219 // AU‑3: run the prompt turn on a spawned task rather than
220 // `block_task().await`-ing it inline. `block_task` blocks the
221 // command loop until the whole turn completes; a mid-turn
222 // `Cancel` would then sit unprocessed in the channel until the
223 // very turn it means to interrupt finishes — no interrupt at
224 // all. Spawning frees the loop to deliver `Cancel` concurrently
225 // (the crate's documented "block_task only inside cx.spawn"
226 // pattern). The prompt reply still resolves the caller's
227 // oneshot; the supervisor already awaits it on its own task.
228 let spawn_result = cx.spawn({
229 let cx = cx.clone();
230 async move {
231 let outcome = cx
232 .send_request(PromptRequest::new(
233 session_id.0,
234 vec![ContentBlock::Text(TextContent::new(text))],
235 ))
236 .block_task()
237 .await
238 .map(|_response| ())
239 .map_err(map_acp_error);
240 let _ = reply.send(outcome);
241 Ok(())
242 }
243 });
244 if let Err(err) = spawn_result {
245 tracing::debug!(%err, "ACP prompt spawn failed");
246 }
247 }
248 DriverCommand::Cancel { session_id, reply } => {
249 // `session/cancel` is a fire-and-forget notification (no
250 // reply); the agent stops the active turn and emits a
251 // `Cancelled` stop reason on the normal update stream. The
252 // reply oneshot only reports whether the notification was
253 // handed to the transport.
254 let outcome = cx
255 .send_notification(CancelNotification::new(session_id.0))
256 .map_err(map_acp_error);
257 let _ = reply.send(outcome);
258 }
259 }
260 }
261 Ok(())
262 })
263 .await;
264
265 if let Err(err) = result {
266 tracing::debug!(%err, "ACP connection driver exited");
267 }
268 });
269
270 (Arc::new(Connection { commands: cmd_tx }), notif_rx, perm_rx)
271 }
272
273 /// Send the ACP `initialize` handshake.
274 pub async fn initialize(&self) -> Result<()> {
275 self.call(|reply| DriverCommand::Initialize { reply }).await
276 }
277
278 /// Create a new session rooted at `cwd`.
279 pub async fn new_session(&self, cwd: &str) -> Result<SessionId> {
280 let cwd = cwd.to_string();
281 self.call(|reply| DriverCommand::NewSession { cwd, reply })
282 .await
283 }
284
285 /// Send a text prompt to `session`.
286 pub async fn prompt(&self, session: &SessionId, text: &str) -> Result<()> {
287 let session_id = session.clone();
288 let text = text.to_string();
289 self.call(|reply| DriverCommand::Prompt {
290 session_id,
291 text,
292 reply,
293 })
294 .await
295 }
296
297 /// AU‑3: interrupt the active turn on `session` without ending it.
298 ///
299 /// Forwards an ACP `session/cancel` notification; the agent stops its
300 /// current turn and reports a `Cancelled` stop reason on the update
301 /// stream. The session stays open and usable for the next prompt.
302 /// `Ok(())` means the notification reached the transport, not that the
303 /// agent has finished cancelling.
304 pub async fn cancel(&self, session: &SessionId) -> Result<()> {
305 let session_id = session.clone();
306 self.call(|reply| DriverCommand::Cancel { session_id, reply })
307 .await
308 }
309
310 /// Send `make_command(reply)` to the driver task and await its reply.
311 async fn call<T>(
312 &self,
313 make_command: impl FnOnce(oneshot::Sender<Result<T>>) -> DriverCommand,
314 ) -> Result<T> {
315 let (reply_tx, reply_rx) = oneshot::channel();
316 self.commands
317 .send(make_command(reply_tx))
318 .await
319 .map_err(|_| AiError::Transport("ACP connection driver task is gone".into()))?;
320 reply_rx.await.map_err(|_| {
321 AiError::Transport("ACP connection driver dropped the reply channel".into())
322 })?
323 }
324}
325
326fn map_acp_error(err: agent_client_protocol::Error) -> AiError {
327 AiError::Protocol(err.to_string())
328}
329
330#[cfg(test)]
331mod tests {
332 use std::time::Duration;
333
334 use serde_json::{Value, json};
335 use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
336
337 use super::*;
338
339 /// Drives the "agent" side of a mocked duplex ACP connection: reads newline-delimited
340 /// JSON-RPC requests, replies with canned responses, and pushes an unsolicited
341 /// `session/update` notification right after answering `session/new` — mirroring how a
342 /// real agent streams updates once a session exists.
343 async fn run_mock_peer(peer: tokio::io::DuplexStream) {
344 let mut reader = BufReader::new(peer);
345 let mut line = String::new();
346 loop {
347 line.clear();
348 match reader.read_line(&mut line).await {
349 Ok(0) | Err(_) => break,
350 Ok(_) => {}
351 }
352 let trimmed = line.trim();
353 if trimmed.is_empty() {
354 continue;
355 }
356 let request: Value = match serde_json::from_str(trimmed) {
357 Ok(v) => v,
358 Err(_) => continue,
359 };
360 let method = request.get("method").and_then(Value::as_str).unwrap_or("");
361 let id = request.get("id").cloned().unwrap_or(Value::Null);
362
363 let response = match method {
364 "initialize" => Some(json!({
365 "jsonrpc": "2.0",
366 "id": id,
367 "result": { "protocolVersion": 1 },
368 })),
369 "session/new" => Some(json!({
370 "jsonrpc": "2.0",
371 "id": id,
372 "result": { "sessionId": "sess-1" },
373 })),
374 "session/prompt" => Some(json!({
375 "jsonrpc": "2.0",
376 "id": id,
377 "result": { "stopReason": "end_turn" },
378 })),
379 _ => None,
380 };
381
382 if let Some(response) = response {
383 write_line(reader.get_mut(), &response).await;
384 }
385
386 if method == "session/new" {
387 let notification = json!({
388 "jsonrpc": "2.0",
389 "method": "session/update",
390 "params": {
391 "sessionId": "sess-1",
392 "update": {
393 "sessionUpdate": "agent_message_chunk",
394 "content": { "type": "text", "text": "hello from agent" },
395 },
396 },
397 });
398 write_line(reader.get_mut(), ¬ification).await;
399 }
400 }
401 }
402
403 async fn write_line(writer: &mut tokio::io::DuplexStream, value: &Value) {
404 let mut line = value.to_string();
405 line.push('\n');
406 let _ = writer.write_all(line.as_bytes()).await;
407 }
408
409 fn spawn_connection_with_mock_peer() -> (
410 Arc<Connection>,
411 mpsc::UnboundedReceiver<SessionNotification>,
412 ) {
413 let (ours, mock) = tokio::io::duplex(8192);
414 let (reader, writer) = tokio::io::split(ours);
415 let (connection, notif_rx, _perm_rx) = Connection::spawn(reader, writer);
416 tokio::spawn(run_mock_peer(mock));
417 (connection, notif_rx)
418 }
419
420 #[tokio::test]
421 async fn initialize_completes_the_handshake() {
422 let (connection, _notif_rx) = spawn_connection_with_mock_peer();
423 connection
424 .initialize()
425 .await
426 .expect("initialize should succeed");
427 }
428
429 #[tokio::test]
430 async fn new_session_returns_the_peers_session_id() {
431 let (connection, _notif_rx) = spawn_connection_with_mock_peer();
432 connection
433 .initialize()
434 .await
435 .expect("initialize should succeed");
436
437 let session = connection
438 .new_session("/tmp")
439 .await
440 .expect("new_session should succeed");
441 assert_eq!(session.0, "sess-1");
442 }
443
444 #[tokio::test]
445 async fn session_update_notification_reaches_the_receiver() {
446 let (connection, mut notif_rx) = spawn_connection_with_mock_peer();
447 connection
448 .initialize()
449 .await
450 .expect("initialize should succeed");
451 connection
452 .new_session("/tmp")
453 .await
454 .expect("new_session should succeed");
455
456 let notification = tokio::time::timeout(Duration::from_secs(5), notif_rx.recv())
457 .await
458 .expect("a session/update notification should arrive before the timeout")
459 .expect("the notification channel should still be open");
460
461 assert_eq!(notification.session_id.0.to_string(), "sess-1");
462 }
463
464 #[tokio::test]
465 async fn prompt_round_trips_through_the_mock_peer() {
466 let (connection, _notif_rx) = spawn_connection_with_mock_peer();
467 connection
468 .initialize()
469 .await
470 .expect("initialize should succeed");
471 let session = connection
472 .new_session("/tmp")
473 .await
474 .expect("new_session should succeed");
475
476 connection
477 .prompt(&session, "hi")
478 .await
479 .expect("prompt should succeed");
480 }
481
482 /// AU‑3: `cancel` forwards `session/cancel` and leaves the session open —
483 /// a subsequent prompt still round-trips. Exercises the driver's
484 /// concurrent-prompt restructure: the `Cancel` command is serviced
485 /// without the loop being pinned on a prior turn.
486 #[tokio::test]
487 async fn cancel_leaves_the_session_usable() {
488 let (connection, _notif_rx) = spawn_connection_with_mock_peer();
489 connection
490 .initialize()
491 .await
492 .expect("initialize should succeed");
493 let session = connection
494 .new_session("/tmp")
495 .await
496 .expect("new_session should succeed");
497
498 connection
499 .cancel(&session)
500 .await
501 .expect("cancel should reach the transport");
502
503 // The session was interrupted, not ended: the next prompt still works.
504 connection
505 .prompt(&session, "again")
506 .await
507 .expect("prompt after cancel should succeed");
508 }
509}