diff --git a/crates/goose/src/acp/transport/connection.rs b/crates/goose/src/acp/transport/connection.rs index 3ab9179f..49d2a30a 100644 --- a/crates/goose/src/acp/transport/connection.rs +++ b/crates/goose/src/acp/transport/connection.rs @@ -1,9 +1,6 @@ -//! Connection-level state shared between HTTP and WebSocket transports. -//! -//! Each connection hosts one ACP agent task. All server→client messages for -//! the connection are multicast through a single broadcast channel; HTTP GET -//! SSE streams and WebSocket sinks subscribe to that channel. POSTs (and WS -//! text frames) forward client→server messages into the agent over an mpsc. +//! Per-connection state. Server→client messages fan out to a connection-scoped +//! stream, a per-session stream for each active `sessionId`, and an +//! all-outbound stream consumed by WebSocket. use std::{ collections::{HashMap, VecDeque}, @@ -11,43 +8,81 @@ use std::{ }; use anyhow::Result; +use serde_json::Value; use tokio::sync::{broadcast, mpsc, Mutex, RwLock}; use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; -use tracing::{error, info, warn}; +use tracing::{error, info, trace, warn}; use crate::acp::adapters::{ReceiverToAsyncRead, SenderToAsyncWrite}; use crate::acp::server_factory::AcpServer; -/// Broadcast capacity for agent→client messages. Large enough to buffer a -/// typical prompt's streaming notifications even if the subscriber is briefly -/// slow (e.g. during reconnect). const OUTBOUND_BROADCAST_CAPACITY: usize = 1024; -/// Maximum number of server→client messages to retain while no subscriber is -/// attached. In the HTTP flow the client opens `GET /acp` only after receiving -/// the initialize response, so any notifications or server-initiated requests -/// emitted by the agent in that window would otherwise be broadcast to zero -/// subscribers and permanently lost. We buffer them here and replay on the -/// first subscribe. On overflow the oldest message is dropped with a warning. +/// Buffers messages emitted before a subscriber attaches (e.g. session +/// notifications that land before the client opens the session GET stream). const PRE_SUBSCRIBE_BUFFER_CAPACITY: usize = 1024; +#[derive(Clone, Debug)] +pub(crate) enum ResponseRoute { + Connection, + Session(String), +} + +struct OutboundStream { + tx: broadcast::Sender, + pre_subscribe_buffer: Mutex>>, +} + +impl OutboundStream { + fn new() -> Self { + let (tx, _) = broadcast::channel(OUTBOUND_BROADCAST_CAPACITY); + Self { + tx, + pre_subscribe_buffer: Mutex::new(Some(VecDeque::new())), + } + } + + async fn push(&self, msg: String) { + let mut guard = self.pre_subscribe_buffer.lock().await; + match guard.as_mut() { + Some(buf) => { + if buf.len() >= PRE_SUBSCRIBE_BUFFER_CAPACITY { + warn!( + "Pre-subscribe buffer full ({} messages); dropping oldest", + PRE_SUBSCRIBE_BUFFER_CAPACITY + ); + buf.pop_front(); + } + buf.push_back(msg); + } + None => { + drop(guard); + let _ = self.tx.send(msg); + } + } + } + + async fn subscribe_with_replay(&self) -> (Vec, broadcast::Receiver) { + let mut guard = self.pre_subscribe_buffer.lock().await; + let receiver = self.tx.subscribe(); + let replay = guard.take().map(Vec::from).unwrap_or_default(); + (replay, receiver) + } +} + pub(crate) struct Connection { - /// Send client→server messages into the agent. pub to_agent_tx: mpsc::Sender, - /// Subscribe here to receive all server→client messages for this connection. - pub outbound_tx: broadcast::Sender, - /// Pulled exactly once during `initialize` to read the synchronous response - /// that must be returned as the HTTP 200 body before any broadcast - /// subscribers exist. `None` once consumed. + /// Consumed once by `handle_initialize` to read the synchronous initialize + /// response before the router task takes over. pub init_receiver: Mutex>>, - /// Set once the initialize handler has captured the initialize response and - /// handed ownership of the agent output pump over to the broadcast fan-out. pub init_complete: Mutex, - /// Handle to the agent task; aborted on connection termination. pub agent_handle: tokio::task::JoinHandle<()>, - /// Handle to the fan-out pump task; aborted on connection termination. - pub pump_handle: Mutex>>, - pre_subscribe_buffer: Arc>>>, + pub router_handle: Mutex>>, + + connection_stream: Arc, + session_streams: Arc>>>, + all_outbound: Arc, + pending_routes: Arc>>, } pub(crate) struct ConnectionRegistry { @@ -63,14 +98,9 @@ impl ConnectionRegistry { } } - /// Create a new connection, spawn the ACP agent task, and return - /// (connection_id, connection). The initialize request body should be sent - /// via `connection.to_agent_tx` and the synchronous initialize response - /// read via `consume_initialize_response`. pub async fn create_connection(&self) -> Result<(String, Arc)> { let (to_agent_tx, to_agent_rx) = mpsc::channel::(256); let (from_agent_tx, from_agent_rx) = mpsc::unbounded_channel::(); - let (outbound_tx, _) = broadcast::channel::(OUTBOUND_BROADCAST_CAPACITY); let agent = self.server.create_agent().await?; let connection_id = uuid::Uuid::new_v4().to_string(); @@ -89,12 +119,14 @@ impl ConnectionRegistry { let connection = Arc::new(Connection { to_agent_tx, - outbound_tx, init_receiver: Mutex::new(Some(from_agent_rx)), init_complete: Mutex::new(false), agent_handle, - pump_handle: Mutex::new(None), - pre_subscribe_buffer: Arc::new(Mutex::new(Some(VecDeque::new()))), + router_handle: Mutex::new(None), + connection_stream: Arc::new(OutboundStream::new()), + session_streams: Arc::new(RwLock::new(HashMap::new())), + all_outbound: Arc::new(OutboundStream::new()), + pending_routes: Arc::new(Mutex::new(HashMap::new())), }); self.connections @@ -116,10 +148,7 @@ impl ConnectionRegistry { } impl Connection { - /// After the synchronous initialize response has been consumed, spawn a - /// task that forwards all remaining agent output to the broadcast channel. - /// Idempotent. - pub async fn start_fanout(self: &Arc) { + pub async fn start_router(self: &Arc) { let mut complete = self.init_complete.lock().await; if *complete { return; @@ -127,49 +156,121 @@ impl Connection { let Some(mut rx) = self.init_receiver.lock().await.take() else { return; }; - let outbound_tx = self.outbound_tx.clone(); - let buffer = self.pre_subscribe_buffer.clone(); + + let me = self.clone(); let handle = tokio::spawn(async move { while let Some(msg) = rx.recv().await { - let mut buf_guard = buffer.lock().await; - match buf_guard.as_mut() { - Some(buf) => { - if buf.len() >= PRE_SUBSCRIBE_BUFFER_CAPACITY { - warn!( - "Pre-subscribe buffer full ({} messages); dropping oldest", - PRE_SUBSCRIBE_BUFFER_CAPACITY - ); - buf.pop_front(); - } - buf.push_back(msg); - } - None => { - drop(buf_guard); - let _ = outbound_tx.send(msg); - } - } + me.route_outbound(msg).await; } }); - *self.pump_handle.lock().await = Some(handle); + *self.router_handle.lock().await = Some(handle); *complete = true; } - pub async fn subscribe_with_replay(&self) -> (Vec, broadcast::Receiver) { - let mut guard = self.pre_subscribe_buffer.lock().await; - let receiver = self.outbound_tx.subscribe(); - let replay = guard.take().map(Vec::from).unwrap_or_default(); - (replay, receiver) + async fn route_outbound(self: &Arc, msg: String) { + self.all_outbound.push(msg.clone()).await; + + let parsed: Option = serde_json::from_str(&msg).ok(); + let target = match parsed.as_ref() { + Some(v) => self.classify(v).await, + None => Target::Connection, + }; + + match target { + Target::Connection => { + trace!(target = "connection", "→ connection-scoped stream"); + self.connection_stream.push(msg).await; + } + Target::Session(sid) => { + trace!(target = %sid, "→ session-scoped stream"); + let stream = self.get_or_create_session_stream(&sid).await; + stream.push(msg).await; + } + } + } + + async fn classify(self: &Arc, v: &Value) -> Target { + let has_method = v.get("method").is_some(); + let has_id = v.get("id").is_some(); + let has_result_or_error = v.get("result").is_some() || v.get("error").is_some(); + + if has_method { + if let Some(sid) = extract_session_id_from_params(v) { + return Target::Session(sid); + } + return Target::Connection; + } + + if has_id && has_result_or_error { + let id = v.get("id").cloned().unwrap_or(Value::Null); + let route = self.pending_routes.lock().await.remove(&id); + return match route { + Some(ResponseRoute::Session(sid)) => Target::Session(sid), + Some(ResponseRoute::Connection) | None => Target::Connection, + }; + } + + Target::Connection + } + + pub async fn record_pending_route(&self, id: Value, route: ResponseRoute) { + if id.is_null() { + return; + } + self.pending_routes.lock().await.insert(id, route); + } + + pub async fn subscribe_connection_stream(&self) -> (Vec, broadcast::Receiver) { + self.connection_stream.subscribe_with_replay().await + } + + pub async fn subscribe_session_stream( + &self, + session_id: &str, + ) -> Option<(Vec, broadcast::Receiver)> { + let stream = self.session_streams.read().await.get(session_id).cloned()?; + Some(stream.subscribe_with_replay().await) + } + + pub async fn ensure_session(&self, session_id: &str) { + self.get_or_create_session_stream(session_id).await; + } + + async fn get_or_create_session_stream(&self, session_id: &str) -> Arc { + if let Some(s) = self.session_streams.read().await.get(session_id) { + return s.clone(); + } + let mut w = self.session_streams.write().await; + w.entry(session_id.to_string()) + .or_insert_with(|| Arc::new(OutboundStream::new())) + .clone() + } + + pub async fn subscribe_all_outbound(&self) -> (Vec, broadcast::Receiver) { + self.all_outbound.subscribe_with_replay().await } - /// Terminate the connection: abort the agent task and the fan-out pump. pub async fn shutdown(&self) { self.agent_handle.abort(); - if let Some(h) = self.pump_handle.lock().await.take() { + if let Some(h) = self.router_handle.lock().await.take() { h.abort(); } } } +#[derive(Debug)] +enum Target { + Connection, + Session(String), +} + +fn extract_session_id_from_params(v: &Value) -> Option { + v.get("params") + .and_then(|p| p.get("sessionId")) + .and_then(|s| s.as_str()) + .map(|s| s.to_string()) +} + #[cfg(test)] mod tests { use super::*; @@ -179,7 +280,6 @@ mod tests { fn fake_connection() -> (Arc, mpsc::UnboundedSender) { let (to_agent_tx, _to_agent_rx) = mpsc::channel::(256); let (from_agent_tx, from_agent_rx) = mpsc::unbounded_channel::(); - let (outbound_tx, _) = broadcast::channel::(OUTBOUND_BROADCAST_CAPACITY); let agent_handle = tokio::spawn(async { std::future::pending::<()>().await; @@ -187,56 +287,121 @@ mod tests { let connection = Arc::new(Connection { to_agent_tx, - outbound_tx, init_receiver: Mutex::new(Some(from_agent_rx)), init_complete: Mutex::new(false), agent_handle, - pump_handle: Mutex::new(None), - pre_subscribe_buffer: Arc::new(Mutex::new(Some(VecDeque::new()))), + router_handle: Mutex::new(None), + connection_stream: Arc::new(OutboundStream::new()), + session_streams: Arc::new(RwLock::new(HashMap::new())), + all_outbound: Arc::new(OutboundStream::new()), + pending_routes: Arc::new(Mutex::new(HashMap::new())), }); (connection, from_agent_tx) } #[tokio::test] - async fn buffers_messages_emitted_before_first_subscribe() { + async fn buffers_connection_scoped_messages_before_first_subscribe() { let (conn, agent_tx) = fake_connection(); - conn.start_fanout().await; + conn.start_router().await; - agent_tx.send("one".to_string()).unwrap(); - agent_tx.send("two".to_string()).unwrap(); - agent_tx.send("three".to_string()).unwrap(); + agent_tx + .send(r#"{"id":1,"result":{"capabilities":{}}}"#.to_string()) + .unwrap(); - tokio::task::yield_now().await; tokio::time::sleep(Duration::from_millis(20)).await; - let (replay, _rx) = conn.subscribe_with_replay().await; - assert_eq!(replay, vec!["one", "two", "three"]); + let (replay, _rx) = conn.subscribe_connection_stream().await; + assert_eq!(replay.len(), 1); + assert!(replay[0].contains("\"capabilities\"")); conn.shutdown().await; } #[tokio::test] - async fn switches_to_live_broadcast_after_subscribe() { + async fn routes_session_scoped_notification_to_session_stream() { let (conn, agent_tx) = fake_connection(); - conn.start_fanout().await; + conn.start_router().await; - let (replay, mut rx) = conn.subscribe_with_replay().await; - assert!(replay.is_empty()); + conn.ensure_session("sess_abc").await; - agent_tx.send("live-one".to_string()).unwrap(); - agent_tx.send("live-two".to_string()).unwrap(); + let (_, mut rx) = conn.subscribe_session_stream("sess_abc").await.unwrap(); + agent_tx + .send( + r#"{"method":"session/update","params":{"sessionId":"sess_abc","update":{}}}"# + .to_string(), + ) + .unwrap(); - let got1 = timeout(Duration::from_secs(1), rx.recv()) + let got = timeout(Duration::from_secs(1), rx.recv()) .await .unwrap() .unwrap(); - let got2 = timeout(Duration::from_secs(1), rx.recv()) + assert!(got.contains("session/update")); + + let (replay, _) = conn.subscribe_connection_stream().await; + assert!( + replay.is_empty(), + "connection stream should not have session-scoped messages" + ); + + conn.shutdown().await; + } + + #[tokio::test] + async fn routes_response_using_pending_route_table() { + let (conn, agent_tx) = fake_connection(); + conn.start_router().await; + + conn.ensure_session("sess_xyz").await; + conn.record_pending_route( + Value::from(42), + ResponseRoute::Session("sess_xyz".to_string()), + ) + .await; + + let (_, mut rx) = conn.subscribe_session_stream("sess_xyz").await.unwrap(); + agent_tx + .send(r#"{"id":42,"result":{"stopReason":"end_turn"}}"#.to_string()) + .unwrap(); + + let got = timeout(Duration::from_secs(1), rx.recv()) .await .unwrap() .unwrap(); - assert_eq!(got1, "live-one"); - assert_eq!(got2, "live-two"); + assert!(got.contains("\"stopReason\"")); + + conn.shutdown().await; + } + + #[tokio::test] + async fn websocket_all_outbound_sees_everything() { + let (conn, agent_tx) = fake_connection(); + conn.start_router().await; + + agent_tx + .send(r#"{"id":1,"result":{}}"#.to_string()) + .unwrap(); + agent_tx + .send(r#"{"method":"session/update","params":{"sessionId":"s1"}}"#.to_string()) + .unwrap(); + + tokio::time::sleep(Duration::from_millis(20)).await; + + let (replay, _all_rx) = conn.subscribe_all_outbound().await; + assert_eq!(replay.len(), 2); + assert!(replay[0].contains("\"id\":1")); + assert!(replay[1].contains("session/update")); + + conn.shutdown().await; + } + + #[tokio::test] + async fn unknown_session_subscribe_returns_none() { + let (conn, _agent_tx) = fake_connection(); + conn.start_router().await; + + assert!(conn.subscribe_session_stream("nope").await.is_none()); conn.shutdown().await; } @@ -244,22 +409,17 @@ mod tests { #[tokio::test] async fn pre_subscribe_buffer_is_bounded() { let (conn, agent_tx) = fake_connection(); - conn.start_fanout().await; + conn.start_router().await; for i in 0..(PRE_SUBSCRIBE_BUFFER_CAPACITY + 50) { - agent_tx.send(format!("m{}", i)).unwrap(); + agent_tx + .send(format!(r#"{{"id":{},"result":{{}}}}"#, i)) + .unwrap(); } tokio::time::sleep(Duration::from_millis(50)).await; - let (replay, _rx) = conn.subscribe_with_replay().await; + let (replay, _rx) = conn.subscribe_connection_stream().await; assert_eq!(replay.len(), PRE_SUBSCRIBE_BUFFER_CAPACITY); - assert_eq!( - replay.last().unwrap(), - &format!("m{}", PRE_SUBSCRIBE_BUFFER_CAPACITY + 49) - ); - assert_eq!(replay.first().unwrap(), &format!("m{}", 50)); - - conn.shutdown().await; } } diff --git a/crates/goose/src/acp/transport/http.rs b/crates/goose/src/acp/transport/http.rs index 3694cfbe..be3c84f5 100644 --- a/crates/goose/src/acp/transport/http.rs +++ b/crates/goose/src/acp/transport/http.rs @@ -11,16 +11,9 @@ use serde_json::Value; use tokio::sync::broadcast; use tracing::{debug, error, info, trace}; -use super::connection::{Connection, ConnectionRegistry}; +use super::connection::{Connection, ConnectionRegistry, ResponseRoute}; use super::*; -/// POST /acp -/// -/// - `initialize`: creates a new connection, forwards the request, waits for -/// the synchronous initialize response from the agent, and returns it as a -/// 200 OK JSON body with the `Acp-Connection-Id` header set. -/// - All other messages: require `Acp-Connection-Id` (and `Acp-Session-Id` -/// for session-scoped methods), forward to the agent, return 202 Accepted. pub(crate) async fn handle_post( State(registry): State>, request: Request, @@ -92,6 +85,19 @@ pub(crate) async fn handle_post( return (StatusCode::BAD_REQUEST, "Invalid JSON-RPC message").into_response(); } + if let Some(sid) = session_id.as_deref() { + connection.ensure_session(sid).await; + } + if is_jsonrpc_request_with_id(&json_message) { + if let Some(id) = json_message.get("id") { + let route = match session_id.as_deref() { + Some(sid) => ResponseRoute::Session(sid.to_string()), + None => ResponseRoute::Connection, + }; + connection.record_pending_route(id.clone(), route).await; + } + } + let message_str = serde_json::to_string(&json_message).unwrap(); trace!(connection_id = %connection_id, payload = %message_str, "POST → agent"); if connection.to_agent_tx.send(message_str).await.is_err() { @@ -130,7 +136,6 @@ async fn handle_initialize(registry: Arc, json_message: Valu .into_response(); } - // Read exactly one message from the agent: the initialize response. let init_response = { let mut guard = connection.init_receiver.lock().await; let Some(rx) = guard.as_mut() else { @@ -158,7 +163,7 @@ async fn handle_initialize(registry: Arc, json_message: Valu } }; - connection.start_fanout().await; + connection.start_router().await; let mut response = ( StatusCode::OK, @@ -173,11 +178,6 @@ async fn handle_initialize(registry: Arc, json_message: Valu response } -/// GET /acp (no Upgrade) -/// -/// Opens the single long-lived SSE stream for a connection. All server→client -/// messages (responses + notifications + server-initiated requests) are -/// delivered here, correlated by their JSON-RPC body fields. pub(crate) async fn handle_get( registry: Arc, request: Request, @@ -202,13 +202,30 @@ pub(crate) async fn handle_get( return (StatusCode::NOT_FOUND, "Unknown Acp-Connection-Id").into_response(); }; - let (replay, receiver) = connection.subscribe_with_replay().await; + let session_id = header_value(&request, HEADER_SESSION_ID); + + let (replay, receiver) = match session_id.as_deref() { + Some(sid) => { + connection.ensure_session(sid).await; + connection + .subscribe_session_stream(sid) + .await + .expect("session stream exists after ensure_session") + } + None => connection.subscribe_connection_stream().await, + }; + let sse = build_sse_stream(connection.clone(), replay, receiver); let mut response = sse.into_response(); if let Ok(v) = HeaderValue::from_str(&connection_id) { response.headers_mut().insert(HEADER_CONNECTION_ID, v); } + if let Some(sid) = session_id { + if let Ok(v) = HeaderValue::from_str(&sid) { + response.headers_mut().insert(HEADER_SESSION_ID, v); + } + } response } @@ -244,7 +261,6 @@ fn build_sse_stream( ) } -/// DELETE /acp pub(crate) async fn handle_delete( State(registry): State>, request: Request, diff --git a/crates/goose/src/acp/transport/websocket.rs b/crates/goose/src/acp/transport/websocket.rs index 752cb320..6da9b4a6 100644 --- a/crates/goose/src/acp/transport/websocket.rs +++ b/crates/goose/src/acp/transport/websocket.rs @@ -11,13 +11,6 @@ use tracing::{debug, error, info, trace, warn}; use super::connection::ConnectionRegistry; use super::HEADER_CONNECTION_ID; -/// GET /acp with `Upgrade: websocket` -/// -/// Creates a new connection (same lifecycle as Streamable HTTP), upgrades to a -/// WebSocket, and runs a bidirectional message loop. The client still sends -/// `initialize` as the first WS text frame — unlike the HTTP path, the -/// initialize response is streamed back over the same WebSocket rather than -/// returned synchronously. pub(crate) async fn handle_ws_upgrade( registry: Arc, ws: WebSocketUpgrade, @@ -34,10 +27,7 @@ pub(crate) async fn handle_ws_upgrade( } }; - // WebSocket does not need the synchronous initialize split — start the - // broadcast fan-out immediately so the WS sink reads from the same stream - // of server→client messages as any HTTP SSE subscribers would. - connection.start_fanout().await; + connection.start_router().await; let conn_id_for_handler = connection_id.clone(); let registry_for_handler = registry.clone(); @@ -65,7 +55,7 @@ async fn run_ws( connection: Arc, ) { let (mut ws_tx, mut ws_rx) = socket.split(); - let (replay, mut outbound_rx) = connection.subscribe_with_replay().await; + let (replay, mut outbound_rx) = connection.subscribe_all_outbound().await; debug!(connection_id = %connection_id, "Starting WebSocket message loop"); diff --git a/ui/sdk/src/http-stream.ts b/ui/sdk/src/http-stream.ts index e934807d..efec232e 100644 --- a/ui/sdk/src/http-stream.ts +++ b/ui/sdk/src/http-stream.ts @@ -3,8 +3,6 @@ import type { AnyMessage, Stream } from "@agentclientprotocol/sdk"; const ACP_CONNECTION_HEADER = "Acp-Connection-Id"; const ACP_SESSION_HEADER = "Acp-Session-Id"; -// Enable via `globalThis.ACP_DEBUG = true`, `localStorage.ACP_DEBUG = "1"`, -// or `ACP_DEBUG=1` in the environment. function acpDebug(label: string, payload: unknown): void { const g = globalThis as { ACP_DEBUG?: unknown; @@ -21,7 +19,6 @@ function acpDebug(label: string, payload: unknown): void { console.debug(`[acp] ${label}`, payload); } -// Methods that are scoped to a session and require an Acp-Session-Id header. const SESSION_SCOPED_METHODS = new Set([ "session/prompt", "session/cancel", @@ -39,6 +36,10 @@ function messageParams(msg: AnyMessage): unknown { return (msg as { params?: unknown }).params; } +function messageResult(msg: AnyMessage): unknown { + return (msg as { result?: unknown }).result; +} + function isRequest(msg: AnyMessage): boolean { const m = msg as { method?: unknown; id?: unknown }; return typeof m.method === "string" && m.id !== undefined && m.id !== null; @@ -49,6 +50,16 @@ function isNotification(msg: AnyMessage): boolean { return typeof m.method === "string" && (m.id === undefined || m.id === null); } +function isResponse(msg: AnyMessage): boolean { + const m = msg as { method?: unknown; id?: unknown; result?: unknown; error?: unknown }; + return ( + m.method === undefined && + m.id !== undefined && + m.id !== null && + (m.result !== undefined || m.error !== undefined) + ); +} + function extractSessionId(value: unknown): string | null { if (value && typeof value === "object" && "sessionId" in value) { const sid = (value as { sessionId?: unknown }).sessionId; @@ -58,28 +69,19 @@ function extractSessionId(value: unknown): string | null { } /** - * Create a Stream that speaks the Streamable HTTP ACP transport. - * - * Protocol summary: - * - The first outbound message must be an `initialize` request, sent as a - * regular POST. The server responds synchronously with 200 OK, a JSON body - * containing the initialize response, and an `Acp-Connection-Id` header. - * - After initialize, we open a single long-lived GET SSE stream carrying - * all server → client messages (responses, notifications, server-initiated - * requests) for every session on the connection. - * - All subsequent POSTs carry `Acp-Connection-Id` and return 202 Accepted. - * Session-scoped methods must also carry `Acp-Session-Id`. - * - On close we send DELETE /acp with the connection header. + * Stream that speaks the ACP Streamable HTTP transport: a connection-scoped + * GET SSE stream plus a session-scoped stream per active `sessionId`. */ export function createHttpStream(serverUrl: string): Stream { const base = serverUrl.replace(/\/+$/, ""); const endpoint = `${base}/acp`; let connectionId: string | null = null; - let getStreamAbort: AbortController | null = null; + let connectionStreamAbort: AbortController | null = null; + const sessionStreamAborts = new Map(); + const openSessionStreams = new Set(); let closed = false; - // Readable-stream plumbing: enqueue-with-buffer until the consumer pulls. const inbox: AnyMessage[] = []; let pullResolve: (() => void) | null = null; @@ -99,9 +101,9 @@ export function createHttpStream(serverUrl: string): Stream { }); } - async function openGetStream() { + async function openConnectionGetStream() { if (!connectionId) return; - getStreamAbort = new AbortController(); + connectionStreamAbort = new AbortController(); const response = await fetch(endpoint, { method: "GET", @@ -109,23 +111,75 @@ export function createHttpStream(serverUrl: string): Stream { Accept: "text/event-stream", [ACP_CONNECTION_HEADER]: connectionId, }, - signal: getStreamAbort.signal, + signal: connectionStreamAbort.signal, }); if (!response.ok || !response.body) { throw new Error( - `Failed to open ACP GET stream: ${response.status} ${response.statusText}`, + `Failed to open ACP connection-scoped GET stream: ${response.status} ${response.statusText}`, ); } - void consumeSSE(response.body).catch((err) => { + void consumeSSE(response.body, "connection").catch((err) => { if (closed) return; // eslint-disable-next-line no-console - console.error("ACP GET stream error:", err); + console.error("ACP connection-scoped GET stream error:", err); }); } - async function consumeSSE(body: ReadableStream) { + async function ensureSessionGetStream(sessionId: string): Promise { + if (!connectionId) return; + if (openSessionStreams.has(sessionId)) return; + openSessionStreams.add(sessionId); + + const abort = new AbortController(); + sessionStreamAborts.set(sessionId, abort); + + let response: Response; + try { + response = await fetch(endpoint, { + method: "GET", + headers: { + Accept: "text/event-stream", + [ACP_CONNECTION_HEADER]: connectionId, + [ACP_SESSION_HEADER]: sessionId, + }, + signal: abort.signal, + }); + } catch (e) { + openSessionStreams.delete(sessionId); + sessionStreamAborts.delete(sessionId); + throw e; + } + + if (!response.ok || !response.body) { + openSessionStreams.delete(sessionId); + sessionStreamAborts.delete(sessionId); + throw new Error( + `Failed to open ACP session-scoped GET stream for ${sessionId}: ${response.status} ${response.statusText}`, + ); + } + + acpDebug("session GET stream open", { sessionId }); + void consumeSSE(response.body, `session:${sessionId}`) + .catch((err) => { + if (closed) return; + // eslint-disable-next-line no-console + console.error( + `ACP session-scoped GET stream error (${sessionId}):`, + err, + ); + }) + .finally(() => { + if (sessionStreamAborts.get(sessionId) === abort) { + sessionStreamAborts.delete(sessionId); + openSessionStreams.delete(sessionId); + acpDebug("session GET stream closed", { sessionId }); + } + }); + } + + async function consumeSSE(body: ReadableStream, label: string) { const reader = body.getReader(); const decoder = new TextDecoder(); let buffer = ""; @@ -140,17 +194,17 @@ export function createHttpStream(serverUrl: string): Stream { while ((idx = buffer.indexOf("\n\n")) >= 0) { const event = buffer.slice(0, idx); buffer = buffer.slice(idx + 2); - handleSseEvent(event); + handleSseEvent(event, label); } } - if (buffer.length > 0) handleSseEvent(buffer); + if (buffer.length > 0) handleSseEvent(buffer, label); } catch (e: unknown) { if (e instanceof DOMException && e.name === "AbortError") return; throw e; } } - function handleSseEvent(event: string) { + function handleSseEvent(event: string, label: string) { const dataLines: string[] = []; for (const line of event.split("\n")) { if (line.startsWith("data:")) { @@ -166,7 +220,22 @@ export function createHttpStream(serverUrl: string): Stream { return; } - acpDebug("SSE → client", msg); + acpDebug(`SSE → client (${label})`, msg); + handleInbound(msg); + } + + function handleInbound(msg: AnyMessage) { + if (isResponse(msg)) { + const sid = extractSessionId(messageResult(msg)); + if (sid && !openSessionStreams.has(sid)) { + ensureSessionGetStream(sid).catch((err) => { + if (closed) return; + // eslint-disable-next-line no-console + console.error("Failed to open session GET stream:", err); + }); + } + } + deliver(msg); } @@ -197,9 +266,9 @@ export function createHttpStream(serverUrl: string): Stream { const body = (await response.json()) as AnyMessage; acpDebug("initialize response", body); - await openGetStream(); - // Deliver the initialize response to the SDK *after* the GET stream is - // up, so any immediate server-initiated messages won't be missed. + // Open the connection-scoped GET stream before delivering the initialize + // response so we don't miss any immediate server-initiated messages. + await openConnectionGetStream(); deliver(body); } @@ -214,10 +283,11 @@ export function createHttpStream(serverUrl: string): Stream { [ACP_CONNECTION_HEADER]: connectionId, }; + let outboundSessionId: string | null = null; if (isRequest(msg) || isNotification(msg)) { - const sid = extractSessionId(messageParams(msg)); - if (sid) { - headers[ACP_SESSION_HEADER] = sid; + outboundSessionId = extractSessionId(messageParams(msg)); + if (outboundSessionId) { + headers[ACP_SESSION_HEADER] = outboundSessionId; } else if (isRequest(msg)) { const method = messageMethod(msg); if (method && SESSION_SCOPED_METHODS.has(method)) { @@ -226,6 +296,15 @@ export function createHttpStream(serverUrl: string): Stream { } } + if (outboundSessionId && messageMethod(msg) !== "session/load") { + try { + await ensureSessionGetStream(outboundSessionId); + } catch (err) { + // eslint-disable-next-line no-console + console.error("Failed to ensure session GET stream:", err); + } + } + acpDebug("POST → agent", msg); const response = await fetch(endpoint, { method: "POST", @@ -238,7 +317,6 @@ export function createHttpStream(serverUrl: string): Stream { `ACP POST failed: ${response.status} ${response.statusText}`, ); } - // Drain the body so the connection can be reused. await response.arrayBuffer().catch(() => undefined); } @@ -254,6 +332,16 @@ export function createHttpStream(serverUrl: string): Stream { } } + function abortAllStreams() { + connectionStreamAbort?.abort(); + connectionStreamAbort = null; + for (const a of sessionStreamAborts.values()) { + a.abort(); + } + sessionStreamAborts.clear(); + openSessionStreams.clear(); + } + const readable = new ReadableStream({ async pull(controller) { await waitForInbox(); @@ -267,7 +355,7 @@ export function createHttpStream(serverUrl: string): Stream { async cancel() { closed = true; await sendDelete(); - getStreamAbort?.abort(); + abortAllStreams(); if (pullResolve) { const r = pullResolve; pullResolve = null; @@ -296,8 +384,7 @@ export function createHttpStream(serverUrl: string): Stream { async close() { closed = true; await sendDelete(); - getStreamAbort?.abort(); - // Unblock any pending pull so the readable can close. + abortAllStreams(); if (pullResolve) { const r = pullResolve; pullResolve = null; @@ -307,7 +394,7 @@ export function createHttpStream(serverUrl: string): Stream { async abort() { closed = true; await sendDelete(); - getStreamAbort?.abort(); + abortAllStreams(); if (pullResolve) { const r = pullResolve; pullResolve = null;