feat: WebSocket transport for goose-acp (#6895)

Cargo.lock

@@ -2960,6 +2960,7 @@ dependencies = [
  "tracing",
  "tracing-subscriber",
  "url",
+ "uuid",
  "wiremock",
 ]
 

crates/goose-acp/Cargo.toml

@@ -29,7 +29,7 @@ fs-err = "3"
 url = { workspace = true }
 
 # HTTP server dependencies
-axum = "0.8"
+axum = { version = "0.8", features = ["ws"] }
 clap = { version = "4", features = ["derive"] }
 serde = { version = "1.0", features = ["derive"] }
 tower-http = { version = "0.6", features = ["cors"] }
@@ -37,6 +37,7 @@ tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
 async-stream = "0.3.6"
 bytes = "1.11.0"
 http-body-util = "0.1.3"
+uuid = { version = "1.11", features = ["v7"] }
 
 [dev-dependencies]
 assert-json-diff = "2.0.2"

crates/goose-acp/src/adapters.rs

@@ -0,0 +1,123 @@
+//! Shared adapter classes for converting mpsc channels to AsyncRead/AsyncWrite streams
+//! Used by both HTTP and WebSocket transports
+
+use std::{
+    pin::Pin,
+    task::{Context, Poll},
+};
+use tokio::sync::mpsc;
+use tracing::error;
+
+/// Converts an mpsc::Receiver<String> to AsyncRead
+/// Each message is terminated with a newline for JSON-RPC framing
+pub(crate) struct ReceiverToAsyncRead {
+    rx: mpsc::Receiver<String>,
+    buffer: Vec<u8>,
+    pos: usize,
+}
+
+impl ReceiverToAsyncRead {
+    pub(crate) fn new(rx: mpsc::Receiver<String>) -> Self {
+        Self {
+            rx,
+            buffer: Vec::new(),
+            pos: 0,
+        }
+    }
+}
+
+impl tokio::io::AsyncRead for ReceiverToAsyncRead {
+    fn poll_read(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        buf: &mut tokio::io::ReadBuf<'_>,
+    ) -> Poll<std::io::Result<()>> {
+        if self.pos < self.buffer.len() {
+            let remaining = &self.buffer[self.pos..];
+            let to_copy = remaining.len().min(buf.remaining());
+            buf.put_slice(&remaining[..to_copy]);
+            self.pos += to_copy;
+            if self.pos >= self.buffer.len() {
+                self.buffer.clear();
+                self.pos = 0;
+            }
+            return Poll::Ready(Ok(()));
+        }
+
+        match Pin::new(&mut self.rx).poll_recv(cx) {
+            Poll::Ready(Some(msg)) => {
+                let bytes = format!("{}\n", msg).into_bytes();
+                let to_copy = bytes.len().min(buf.remaining());
+                buf.put_slice(&bytes[..to_copy]);
+                if to_copy < bytes.len() {
+                    self.buffer = bytes[to_copy..].to_vec();
+                    self.pos = 0;
+                }
+                Poll::Ready(Ok(()))
+            }
+            Poll::Ready(None) => Poll::Ready(Ok(())),
+            Poll::Pending => Poll::Pending,
+        }
+    }
+}
+
+/// Converts an mpsc::Sender<String> to AsyncWrite
+/// Splits incoming data on newlines for JSON-RPC framing
+pub(crate) struct SenderToAsyncWrite {
+    tx: mpsc::Sender<String>,
+    buffer: Vec<u8>,
+}
+
+impl SenderToAsyncWrite {
+    pub(crate) fn new(tx: mpsc::Sender<String>) -> Self {
+        Self {
+            tx,
+            buffer: Vec::new(),
+        }
+    }
+}
+
+impl tokio::io::AsyncWrite for SenderToAsyncWrite {
+    fn poll_write(
+        mut self: Pin<&mut Self>,
+        _cx: &mut Context<'_>,
+        buf: &[u8],
+    ) -> Poll<std::io::Result<usize>> {
+        self.buffer.extend_from_slice(buf);
+
+        while let Some(pos) = self.buffer.iter().position(|&b| b == b'\n') {
+            let line = String::from_utf8_lossy(&self.buffer[..pos]).to_string();
+            self.buffer.drain(..=pos);
+
+            if !line.is_empty() {
+                if let Err(e) = self.tx.try_send(line.clone()) {
+                    match e {
+                        mpsc::error::TrySendError::Full(_) => {
+                            let truncated: String = line.chars().take(100).collect();
+                            error!(
+                                "Channel full, dropping message (backpressure): {}",
+                                truncated
+                            );
+                        }
+                        mpsc::error::TrySendError::Closed(_) => {
+                            return Poll::Ready(Err(std::io::Error::new(
+                                std::io::ErrorKind::BrokenPipe,
+                                "Channel closed",
+                            )));
+                        }
+                    }
+                }
+            }
+        }
+
+        Poll::Ready(Ok(buf.len()))
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+        Poll::Ready(Ok(()))
+    }
+}

crates/goose-acp/src/bin/server.rs

@@ -1,9 +1,6 @@
 use anyhow::Result;
 use clap::Parser;
-use goose_acp::{
-    http::{self, HttpState},
-    server_factory::{AcpServer, AcpServerFactoryConfig},
-};
+use goose_acp::server_factory::{AcpServer, AcpServerFactoryConfig};
 use std::net::SocketAddr;
 use std::sync::Arc;
 use tracing::info;
@@ -11,7 +8,7 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilte
 
 #[derive(Parser)]
 #[command(name = "goose-acp-server")]
-#[command(about = "ACP server for goose over streamable HTTP")]
+#[command(about = "ACP server for goose over HTTP and WebSocket")]
 struct Cli {
     #[arg(long, default_value = "127.0.0.1")]
     host: String,
@@ -45,12 +42,13 @@ async fn main() -> Result<()> {
     };
 
     let server = Arc::new(AcpServer::new(config));
-    let state = Arc::new(HttpState::new(server));
+    let router = goose_acp::transport::create_router(server);
 
     let addr: SocketAddr = format!("{}:{}", cli.host, cli.port).parse()?;
     info!("Starting goose-acp-server on {}", addr);
 
-    http::serve(state, addr).await?;
+    let listener = tokio::net::TcpListener::bind(addr).await?;
+    axum::serve(listener, router).await?;
 
     Ok(())
 }

crates/goose-acp/src/lib.rs

@@ -1,5 +1,6 @@
 #![recursion_limit = "256"]
 
-pub mod http;
+mod adapters;
 pub mod server;
 pub mod server_factory;
+pub mod transport;

crates/goose-acp/src/transport.rs

@@ -0,0 +1,127 @@
+pub mod http;
+pub mod websocket;
+
+use std::sync::Arc;
+
+use axum::{
+    body::Body,
+    extract::{
+        ws::{rejection::WebSocketUpgradeRejection, WebSocketUpgrade},
+        State,
+    },
+    http::{header, Method, Request},
+    response::Response,
+    routing::{delete, get, post},
+    Router,
+};
+use serde_json::Value;
+use tokio::sync::{mpsc, Mutex};
+use tower_http::cors::{Any, CorsLayer};
+
+use crate::server_factory::AcpServer;
+
+pub(crate) const HEADER_SESSION_ID: &str = "Acp-Session-Id";
+pub(crate) const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream";
+pub(crate) const JSON_MIME_TYPE: &str = "application/json";
+
+pub(crate) struct TransportSession {
+    pub to_agent_tx: mpsc::Sender<String>,
+    pub from_agent_rx: Arc<Mutex<mpsc::Receiver<String>>>,
+    pub handle: tokio::task::JoinHandle<()>,
+}
+
+pub(crate) fn accepts_mime_type(request: &Request<Body>, mime_type: &str) -> bool {
+    request
+        .headers()
+        .get(axum::http::header::ACCEPT)
+        .and_then(|v| v.to_str().ok())
+        .is_some_and(|accept| accept.contains(mime_type))
+}
+
+pub(crate) fn accepts_json_and_sse(request: &Request<Body>) -> bool {
+    request
+        .headers()
+        .get(axum::http::header::ACCEPT)
+        .and_then(|v| v.to_str().ok())
+        .is_some_and(|accept| {
+            accept.contains(JSON_MIME_TYPE) && accept.contains(EVENT_STREAM_MIME_TYPE)
+        })
+}
+
+pub(crate) fn content_type_is_json(request: &Request<Body>) -> bool {
+    request
+        .headers()
+        .get(axum::http::header::CONTENT_TYPE)
+        .and_then(|v| v.to_str().ok())
+        .is_some_and(|ct| ct.starts_with(JSON_MIME_TYPE))
+}
+
+pub(crate) fn get_session_id(request: &Request<Body>) -> Option<String> {
+    request
+        .headers()
+        .get(HEADER_SESSION_ID)
+        .and_then(|v| v.to_str().ok())
+        .map(|s| s.to_string())
+}
+
+pub(crate) fn is_jsonrpc_request(value: &Value) -> bool {
+    value.get("method").is_some() && value.get("id").is_some()
+}
+
+pub(crate) fn is_jsonrpc_notification(value: &Value) -> bool {
+    value.get("method").is_some() && value.get("id").is_none()
+}
+
+pub(crate) fn is_jsonrpc_response(value: &Value) -> bool {
+    value.get("id").is_some() && (value.get("result").is_some() || value.get("error").is_some())
+}
+
+pub(crate) fn is_initialize_request(value: &Value) -> bool {
+    value.get("method").is_some_and(|m| m == "initialize") && value.get("id").is_some()
+}
+
+async fn handle_get(
+    ws_upgrade: Result<WebSocketUpgrade, WebSocketUpgradeRejection>,
+    State(state): State<(Arc<http::HttpState>, Arc<websocket::WsState>)>,
+    request: Request<Body>,
+) -> Response {
+    match ws_upgrade {
+        Ok(ws) => websocket::handle_get(state.1, ws).await,
+        Err(_) => http::handle_get(state.0, request).await,
+    }
+}
+
+async fn health() -> &'static str {
+    "ok"
+}
+
+pub fn create_router(server: Arc<AcpServer>) -> Router {
+    let http_state = Arc::new(http::HttpState::new(server.clone()));
+    let ws_state = Arc::new(websocket::WsState::new(server));
+
+    let cors = CorsLayer::new()
+        .allow_origin(Any)
+        .allow_methods([Method::GET, Method::POST, Method::DELETE, Method::OPTIONS])
+        .allow_headers([
+            header::CONTENT_TYPE,
+            header::ACCEPT,
+            HEADER_SESSION_ID.parse().unwrap(),
+            header::SEC_WEBSOCKET_VERSION,
+            header::SEC_WEBSOCKET_KEY,
+            header::CONNECTION,
+            header::UPGRADE,
+        ]);
+
+    Router::new()
+        .route("/health", get(health))
+        .route(
+            "/acp",
+            post(http::handle_post).with_state(http_state.clone()),
+        )
+        .route(
+            "/acp",
+            get(handle_get).with_state((http_state.clone(), ws_state)),
+        )
+        .route("/acp", delete(http::handle_delete).with_state(http_state))
+        .layer(cors)
+}

crates/goose-acp/src/transport/http.rs

@@ -2,42 +2,23 @@ use anyhow::Result;
 use axum::{
     body::Body,
     extract::State,
-    http::{header, Method, Request, StatusCode},
+    http::{Request, StatusCode},
     response::{IntoResponse, Response, Sse},
-    routing::{delete, get, post},
-    Router,
 };
 use http_body_util::BodyExt;
 use serde_json::Value;
-use std::{
-    collections::HashMap,
-    convert::Infallible,
-    pin::Pin,
-    sync::Arc,
-    task::{Context, Poll},
-    time::Duration,
-};
+use std::{collections::HashMap, convert::Infallible, sync::Arc, time::Duration};
 use tokio::sync::{mpsc, Mutex, RwLock};
 use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
-use tower_http::cors::{Any, CorsLayer};
 use tracing::{error, info};
 
+use super::*;
+use crate::adapters::{ReceiverToAsyncRead, SenderToAsyncWrite};
 use crate::server_factory::AcpServer;
 
-// ACP header constants
-const HEADER_SESSION_ID: &str = "Acp-Session-Id";
-const EVENT_STREAM_MIME_TYPE: &str = "text/event-stream";
-const JSON_MIME_TYPE: &str = "application/json";
-
-struct HttpSession {
-    to_agent_tx: mpsc::Sender<String>,
-    from_agent_rx: Arc<Mutex<mpsc::Receiver<String>>>,
-    handle: tokio::task::JoinHandle<()>,
-}
-
-pub struct HttpState {
+pub(crate) struct HttpState {
     server: Arc<AcpServer>,
-    sessions: RwLock<HashMap<String, HttpSession>>,
+    sessions: RwLock<HashMap<String, TransportSession>>,
 }
 
 impl HttpState {
@@ -75,7 +56,7 @@ impl HttpState {
 
         self.sessions.write().await.insert(
             session_id.clone(),
-            HttpSession {
+            TransportSession {
                 to_agent_tx,
                 from_agent_rx: Arc::new(Mutex::new(from_agent_rx)),
                 handle,
@@ -117,166 +98,6 @@ impl HttpState {
     }
 }
 
-struct ReceiverToAsyncRead {
-    rx: mpsc::Receiver<String>,
-    buffer: Vec<u8>,
-    pos: usize,
-}
-
-impl ReceiverToAsyncRead {
-    fn new(rx: mpsc::Receiver<String>) -> Self {
-        Self {
-            rx,
-            buffer: Vec::new(),
-            pos: 0,
-        }
-    }
-}
-
-impl tokio::io::AsyncRead for ReceiverToAsyncRead {
-    fn poll_read(
-        mut self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-        buf: &mut tokio::io::ReadBuf<'_>,
-    ) -> Poll<std::io::Result<()>> {
-        if self.pos < self.buffer.len() {
-            let remaining = &self.buffer[self.pos..];
-            let to_copy = remaining.len().min(buf.remaining());
-            buf.put_slice(&remaining[..to_copy]);
-            self.pos += to_copy;
-            if self.pos >= self.buffer.len() {
-                self.buffer.clear();
-                self.pos = 0;
-            }
-            return Poll::Ready(Ok(()));
-        }
-
-        match Pin::new(&mut self.rx).poll_recv(cx) {
-            Poll::Ready(Some(msg)) => {
-                let bytes = format!("{}\n", msg).into_bytes();
-                let to_copy = bytes.len().min(buf.remaining());
-                buf.put_slice(&bytes[..to_copy]);
-                if to_copy < bytes.len() {
-                    self.buffer = bytes[to_copy..].to_vec();
-                    self.pos = 0;
-                }
-                Poll::Ready(Ok(()))
-            }
-            Poll::Ready(None) => Poll::Ready(Ok(())),
-            Poll::Pending => Poll::Pending,
-        }
-    }
-}
-
-struct SenderToAsyncWrite {
-    tx: mpsc::Sender<String>,
-    buffer: Vec<u8>,
-}
-
-impl SenderToAsyncWrite {
-    fn new(tx: mpsc::Sender<String>) -> Self {
-        Self {
-            tx,
-            buffer: Vec::new(),
-        }
-    }
-}
-
-impl tokio::io::AsyncWrite for SenderToAsyncWrite {
-    fn poll_write(
-        mut self: Pin<&mut Self>,
-        _cx: &mut Context<'_>,
-        buf: &[u8],
-    ) -> Poll<std::io::Result<usize>> {
-        self.buffer.extend_from_slice(buf);
-
-        while let Some(pos) = self.buffer.iter().position(|&b| b == b'\n') {
-            let line = String::from_utf8_lossy(&self.buffer[..pos]).to_string();
-            self.buffer.drain(..=pos);
-
-            if !line.is_empty() {
-                if let Err(e) = self.tx.try_send(line.clone()) {
-                    match e {
-                        mpsc::error::TrySendError::Full(_) => {
-                            let truncated: String = line.chars().take(100).collect();
-                            error!(
-                                "Channel full, dropping message (backpressure): {}",
-                                truncated
-                            );
-                        }
-                        mpsc::error::TrySendError::Closed(_) => {
-                            return Poll::Ready(Err(std::io::Error::new(
-                                std::io::ErrorKind::BrokenPipe,
-                                "Channel closed",
-                            )));
-                        }
-                    }
-                }
-            }
-        }
-
-        Poll::Ready(Ok(buf.len()))
-    }
-
-    fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        Poll::Ready(Ok(()))
-    }
-
-    fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        Poll::Ready(Ok(()))
-    }
-}
-
-fn accepts_mime_type(request: &Request<Body>, mime_type: &str) -> bool {
-    request
-        .headers()
-        .get(header::ACCEPT)
-        .and_then(|v| v.to_str().ok())
-        .is_some_and(|accept| accept.contains(mime_type))
-}
-
-fn accepts_json_and_sse(request: &Request<Body>) -> bool {
-    request
-        .headers()
-        .get(header::ACCEPT)
-        .and_then(|v| v.to_str().ok())
-        .is_some_and(|accept| {
-            accept.contains(JSON_MIME_TYPE) && accept.contains(EVENT_STREAM_MIME_TYPE)
-        })
-}
-
-fn content_type_is_json(request: &Request<Body>) -> bool {
-    request
-        .headers()
-        .get(header::CONTENT_TYPE)
-        .and_then(|v| v.to_str().ok())
-        .is_some_and(|ct| ct.starts_with(JSON_MIME_TYPE))
-}
-
-fn get_session_id(request: &Request<Body>) -> Option<String> {
-    request
-        .headers()
-        .get(HEADER_SESSION_ID)
-        .and_then(|v| v.to_str().ok())
-        .map(|s| s.to_string())
-}
-
-fn is_jsonrpc_request(value: &Value) -> bool {
-    value.get("method").is_some() && value.get("id").is_some()
-}
-
-fn is_jsonrpc_notification(value: &Value) -> bool {
-    value.get("method").is_some() && value.get("id").is_none()
-}
-
-fn is_jsonrpc_response(value: &Value) -> bool {
-    value.get("id").is_some() && (value.get("result").is_some() || value.get("error").is_some())
-}
-
-fn is_initialize_request(value: &Value) -> bool {
-    value.get("method").is_some_and(|m| m == "initialize") && value.get("id").is_some()
-}
-
 fn create_sse_stream(
     receiver: Arc<Mutex<mpsc::Receiver<String>>>,
     cleanup: Option<(Arc<HttpState>, String)>,
@@ -365,7 +186,10 @@ async fn handle_notification_or_response(
     StatusCode::ACCEPTED.into_response()
 }
 
-async fn handle_post(State(state): State<Arc<HttpState>>, request: Request<Body>) -> Response {
+pub(crate) async fn handle_post(
+    State(state): State<Arc<HttpState>>,
+    request: Request<Body>,
+) -> Response {
     if !accepts_json_and_sse(&request) {
         return (
             StatusCode::NOT_ACCEPTABLE,
@@ -409,7 +233,7 @@ async fn handle_post(State(state): State<Arc<HttpState>>, request: Request<Body>
     }
 
     if is_initialize_request(&json_message) {
-        handle_initialize(state, &json_message).await
+        handle_initialize(state.clone(), &json_message).await
     } else if is_jsonrpc_request(&json_message) {
         let Some(id) = session_id else {
             return (
@@ -418,7 +242,7 @@ async fn handle_post(State(state): State<Arc<HttpState>>, request: Request<Body>
             )
                 .into_response();
         };
-        handle_request(state, id, &json_message).await
+        handle_request(state.clone(), id, &json_message).await
     } else if is_jsonrpc_notification(&json_message) || is_jsonrpc_response(&json_message) {
         let Some(id) = session_id else {
             return (
@@ -427,13 +251,13 @@ async fn handle_post(State(state): State<Arc<HttpState>>, request: Request<Body>
             )
                 .into_response();
         };
-        handle_notification_or_response(state, id, &json_message).await
+        handle_notification_or_response(state.clone(), id, &json_message).await
     } else {
         (StatusCode::BAD_REQUEST, "Invalid JSON-RPC message").into_response()
     }
 }
 
-async fn handle_get(State(state): State<Arc<HttpState>>, request: Request<Body>) -> Response {
+pub(crate) async fn handle_get(state: Arc<HttpState>, request: Request<Body>) -> Response {
     if !accepts_mime_type(&request, EVENT_STREAM_MIME_TYPE) {
         return (
             StatusCode::NOT_ACCEPTABLE,
@@ -478,7 +302,10 @@ async fn handle_get(State(state): State<Arc<HttpState>>, request: Request<Body>)
         .into_response()
 }
 
-async fn handle_delete(State(state): State<Arc<HttpState>>, request: Request<Body>) -> Response {
+pub(crate) async fn handle_delete(
+    State(state): State<Arc<HttpState>>,
+    request: Request<Body>,
+) -> Response {
     let session_id = match get_session_id(&request) {
         Some(id) => id,
         None => {
@@ -497,34 +324,3 @@ async fn handle_delete(State(state): State<Arc<HttpState>>, request: Request<Bod
     state.remove_session(&session_id).await;
     StatusCode::ACCEPTED.into_response()
 }
-
-async fn health() -> &'static str {
-    "ok"
-}
-
-pub fn create_router(state: Arc<HttpState>) -> Router {
-    let cors = CorsLayer::new()
-        .allow_origin(Any)
-        .allow_methods([Method::GET, Method::POST, Method::DELETE, Method::OPTIONS])
-        .allow_headers([
-            header::CONTENT_TYPE,
-            header::ACCEPT,
-            HEADER_SESSION_ID.parse().unwrap(),
-        ]);
-
-    Router::new()
-        .route("/health", get(health))
-        .route("/acp", post(handle_post))
-        .route("/acp", get(handle_get))
-        .route("/acp", delete(handle_delete))
-        .layer(cors)
-        .with_state(state)
-}
-
-pub async fn serve(state: Arc<HttpState>, addr: std::net::SocketAddr) -> Result<()> {
-    let router = create_router(state);
-    let listener = tokio::net::TcpListener::bind(addr).await?;
-    info!("ACP HTTP server listening on {}", addr);
-    axum::serve(listener, router).await?;
-    Ok(())
-}

crates/goose-acp/src/transport/websocket.rs

@@ -0,0 +1,160 @@
+use anyhow::Result;
+use axum::{
+    extract::ws::{Message, WebSocket, WebSocketUpgrade},
+    http::StatusCode,
+    response::{IntoResponse, Response},
+};
+use futures::{SinkExt, StreamExt};
+use std::{collections::HashMap, sync::Arc};
+use tokio::sync::{mpsc, Mutex, RwLock};
+use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt};
+use tracing::{debug, error, info, warn};
+
+use super::{TransportSession, HEADER_SESSION_ID};
+use crate::adapters::{ReceiverToAsyncRead, SenderToAsyncWrite};
+use crate::server_factory::AcpServer;
+
+pub(crate) struct WsState {
+    server: Arc<AcpServer>,
+    sessions: RwLock<HashMap<String, TransportSession>>,
+}
+
+impl WsState {
+    pub fn new(server: Arc<AcpServer>) -> Self {
+        Self {
+            server,
+            sessions: RwLock::new(HashMap::new()),
+        }
+    }
+
+    async fn create_connection(&self) -> Result<String> {
+        let (to_agent_tx, to_agent_rx) = mpsc::channel::<String>(256);
+        let (from_agent_tx, from_agent_rx) = mpsc::channel::<String>(256);
+
+        let agent = self.server.create_agent().await?;
+
+        // Create a Goose ACP session (not just the transport connection)
+        let session_id = agent.create_session().await?;
+
+        let handle = tokio::spawn(async move {
+            let read_stream = ReceiverToAsyncRead::new(to_agent_rx);
+            let write_stream = SenderToAsyncWrite::new(from_agent_tx);
+
+            if let Err(e) =
+                crate::server::serve(agent, read_stream.compat(), write_stream.compat_write()).await
+            {
+                error!("ACP WebSocket session error: {}", e);
+            }
+        });
+
+        self.sessions.write().await.insert(
+            session_id.clone(),
+            TransportSession {
+                to_agent_tx,
+                from_agent_rx: Arc::new(Mutex::new(from_agent_rx)),
+                handle,
+            },
+        );
+
+        info!(session_id = %session_id, "WebSocket connection created");
+        Ok(session_id)
+    }
+
+    async fn remove_connection(&self, session_id: &str) {
+        if let Some(session) = self.sessions.write().await.remove(session_id) {
+            session.handle.abort();
+            info!(session_id = %session_id, "WebSocket connection removed");
+        }
+    }
+}
+
+pub(crate) async fn handle_get(state: Arc<WsState>, ws: WebSocketUpgrade) -> Response {
+    let session_id = match state.create_connection().await {
+        Ok(id) => id,
+        Err(e) => {
+            error!("Failed to create WebSocket connection: {}", e);
+            return (
+                StatusCode::INTERNAL_SERVER_ERROR,
+                "Failed to create WebSocket connection",
+            )
+                .into_response();
+        }
+    };
+
+    let mut response = ws.on_upgrade({
+        let session_id = session_id.clone();
+        move |socket| handle_ws(socket, state, session_id)
+    });
+    response
+        .headers_mut()
+        .insert(HEADER_SESSION_ID, session_id.parse().unwrap());
+    response
+}
+
+pub(crate) async fn handle_ws(socket: WebSocket, state: Arc<WsState>, session_id: String) {
+    let (mut ws_tx, mut ws_rx) = socket.split();
+
+    let (to_agent, from_agent) = {
+        let sessions = state.sessions.read().await;
+        match sessions.get(&session_id) {
+            Some(session) => (session.to_agent_tx.clone(), session.from_agent_rx.clone()),
+            None => {
+                error!(session_id = %session_id, "Session not found after creation");
+                return;
+            }
+        }
+    };
+
+    debug!(session_id = %session_id, "Starting bidirectional message loop");
+
+    let mut from_agent_rx = from_agent.lock().await;
+
+    loop {
+        tokio::select! {
+            Some(msg_result) = ws_rx.next() => {
+                match msg_result {
+                    Ok(Message::Text(text)) => {
+                        let text_str = text.to_string();
+                        debug!(session_id = %session_id, "Client → Agent: {} bytes", text_str.len());
+                        if let Err(e) = to_agent.send(text_str).await {
+                            error!(session_id = %session_id, "Failed to send to agent: {}", e);
+                            break;
+                        }
+                    }
+                    Ok(Message::Close(frame)) => {
+                        debug!(session_id = %session_id, "Client closed connection: {:?}", frame);
+                        break;
+                    }
+                    Ok(Message::Ping(_)) | Ok(Message::Pong(_)) => {
+                        // Axum handles ping/pong automatically
+                        continue;
+                    }
+                    Ok(Message::Binary(_)) => {
+                        warn!(session_id = %session_id, "Ignoring binary message (ACP uses text)");
+                        continue;
+                    }
+                    Err(e) => {
+                        error!(session_id = %session_id, "WebSocket error: {}", e);
+                        break;
+                    }
+                }
+            }
+
+            Some(text) = from_agent_rx.recv() => {
+                debug!(session_id = %session_id, "Agent → Client: {} bytes", text.len());
+                if let Err(e) = ws_tx.send(Message::Text(text.into())).await {
+                    error!(session_id = %session_id, "Failed to send to client: {}", e);
+                    break;
+                }
+            }
+
+            else => {
+                debug!(session_id = %session_id, "Both channels closed");
+                break;
+            }
+        }
+    }
+
+    debug!(session_id = %session_id, "Cleaning up connection");
+    state.remove_connection(&session_id).await;
+}