feat: ACP streamable http spec compliance (#9034)

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<String>,
+    pre_subscribe_buffer: Mutex<Option<VecDeque<String>>>,
+}
+
+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<String>, broadcast::Receiver<String>) {
+        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<String>,
-    /// Subscribe here to receive all server→client messages for this connection.
-    pub outbound_tx: broadcast::Sender<String>,
-    /// 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<Option<mpsc::UnboundedReceiver<String>>>,
-    /// 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<bool>,
-    /// 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<Option<tokio::task::JoinHandle<()>>>,
-    pre_subscribe_buffer: Arc<Mutex<Option<VecDeque<String>>>>,
+    pub router_handle: Mutex<Option<tokio::task::JoinHandle<()>>>,
+
+    connection_stream: Arc<OutboundStream>,
+    session_streams: Arc<RwLock<HashMap<String, Arc<OutboundStream>>>>,
+    all_outbound: Arc<OutboundStream>,
+    pending_routes: Arc<Mutex<HashMap<Value, ResponseRoute>>>,
 }
 
 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<Connection>)> {
         let (to_agent_tx, to_agent_rx) = mpsc::channel::<String>(256);
         let (from_agent_tx, from_agent_rx) = mpsc::unbounded_channel::<String>();
-        let (outbound_tx, _) = broadcast::channel::<String>(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<Self>) {
+    pub async fn start_router(self: &Arc<Self>) {
         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<String>, broadcast::Receiver<String>) {
-        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<Self>, msg: String) {
+        self.all_outbound.push(msg.clone()).await;
+
+        let parsed: Option<Value> = 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<Self>, 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<String>, broadcast::Receiver<String>) {
+        self.connection_stream.subscribe_with_replay().await
+    }
+
+    pub async fn subscribe_session_stream(
+        &self,
+        session_id: &str,
+    ) -> Option<(Vec<String>, broadcast::Receiver<String>)> {
+        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<OutboundStream> {
+        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<String>, broadcast::Receiver<String>) {
+        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<String> {
+    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<Connection>, mpsc::UnboundedSender<String>) {
         let (to_agent_tx, _to_agent_rx) = mpsc::channel::<String>(256);
         let (from_agent_tx, from_agent_rx) = mpsc::unbounded_channel::<String>();
-        let (outbound_tx, _) = broadcast::channel::<String>(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;
     }
 }

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<Arc<ConnectionRegistry>>,
     request: Request<Body>,
@@ -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<ConnectionRegistry>, 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<ConnectionRegistry>, 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<ConnectionRegistry>, 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<ConnectionRegistry>,
     request: Request<Body>,
@@ -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<Arc<ConnectionRegistry>>,
     request: Request<Body>,

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<ConnectionRegistry>,
     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<super::connection::Connection>,
 ) {
     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");
 

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<string>([
   "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<string, AbortController>();
+  const openSessionStreams = new Set<string>();
   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<Uint8Array>) {
+  async function ensureSessionGetStream(sessionId: string): Promise<void> {
+    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<Uint8Array>, 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<AnyMessage>({
     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;