diff --git a/Cargo.lock b/Cargo.lock index 861de1f71b50..46dbd549f512 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2960,6 +2960,7 @@ dependencies = [ "tracing", "tracing-subscriber", "url", + "uuid", "wiremock", ] diff --git a/crates/goose-acp/Cargo.toml b/crates/goose-acp/Cargo.toml index ed2d9abed105..e1a509618666 100644 --- a/crates/goose-acp/Cargo.toml +++ b/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" diff --git a/crates/goose-acp/src/adapters.rs b/crates/goose-acp/src/adapters.rs new file mode 100644 index 000000000000..83813bd0f0b0 --- /dev/null +++ b/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 to AsyncRead +/// Each message is terminated with a newline for JSON-RPC framing +pub(crate) struct ReceiverToAsyncRead { + rx: mpsc::Receiver, + buffer: Vec, + pos: usize, +} + +impl ReceiverToAsyncRead { + pub(crate) fn new(rx: mpsc::Receiver) -> 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> { + 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 to AsyncWrite +/// Splits incoming data on newlines for JSON-RPC framing +pub(crate) struct SenderToAsyncWrite { + tx: mpsc::Sender, + buffer: Vec, +} + +impl SenderToAsyncWrite { + pub(crate) fn new(tx: mpsc::Sender) -> 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> { + 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> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } +} diff --git a/crates/goose-acp/src/bin/server.rs b/crates/goose-acp/src/bin/server.rs index dcfda549c3e8..33ef93e0fff8 100644 --- a/crates/goose-acp/src/bin/server.rs +++ b/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(()) } diff --git a/crates/goose-acp/src/lib.rs b/crates/goose-acp/src/lib.rs index bd21c49fec75..e2830d61237b 100644 --- a/crates/goose-acp/src/lib.rs +++ b/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; diff --git a/crates/goose-acp/src/transport.rs b/crates/goose-acp/src/transport.rs new file mode 100644 index 000000000000..632f540bb8e4 --- /dev/null +++ b/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, + pub from_agent_rx: Arc>>, + pub handle: tokio::task::JoinHandle<()>, +} + +pub(crate) fn accepts_mime_type(request: &Request, 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) -> 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) -> 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) -> Option { + 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, + State(state): State<(Arc, Arc)>, + request: Request, +) -> 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) -> 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) +} diff --git a/crates/goose-acp/src/http.rs b/crates/goose-acp/src/transport/http.rs similarity index 57% rename from crates/goose-acp/src/http.rs rename to crates/goose-acp/src/transport/http.rs index 5583bfa56d36..07555a3fdbc1 100644 --- a/crates/goose-acp/src/http.rs +++ b/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, - from_agent_rx: Arc>>, - handle: tokio::task::JoinHandle<()>, -} - -pub struct HttpState { +pub(crate) struct HttpState { server: Arc, - sessions: RwLock>, + sessions: RwLock>, } 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, - buffer: Vec, - pos: usize, -} - -impl ReceiverToAsyncRead { - fn new(rx: mpsc::Receiver) -> 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> { - 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, - buffer: Vec, -} - -impl SenderToAsyncWrite { - fn new(tx: mpsc::Sender) -> 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> { - 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> { - Poll::Ready(Ok(())) - } - - fn poll_shutdown(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll> { - Poll::Ready(Ok(())) - } -} - -fn accepts_mime_type(request: &Request, 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) -> 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) -> 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) -> Option { - 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>>, cleanup: Option<(Arc, String)>, @@ -365,7 +186,10 @@ async fn handle_notification_or_response( StatusCode::ACCEPTED.into_response() } -async fn handle_post(State(state): State>, request: Request) -> Response { +pub(crate) async fn handle_post( + State(state): State>, + request: Request, +) -> Response { if !accepts_json_and_sse(&request) { return ( StatusCode::NOT_ACCEPTABLE, @@ -409,7 +233,7 @@ async fn handle_post(State(state): State>, request: Request } 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>, request: Request ) .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>, request: Request ) .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>, request: Request) -> Response { +pub(crate) async fn handle_get(state: Arc, request: Request) -> 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>, request: Request) .into_response() } -async fn handle_delete(State(state): State>, request: Request) -> Response { +pub(crate) async fn handle_delete( + State(state): State>, + request: Request, +) -> Response { let session_id = match get_session_id(&request) { Some(id) => id, None => { @@ -497,34 +324,3 @@ async fn handle_delete(State(state): State>, request: Request &'static str { - "ok" -} - -pub fn create_router(state: Arc) -> 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, 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(()) -} diff --git a/crates/goose-acp/src/transport/websocket.rs b/crates/goose-acp/src/transport/websocket.rs new file mode 100644 index 000000000000..a9c6a15edbcd --- /dev/null +++ b/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, + sessions: RwLock>, +} + +impl WsState { + pub fn new(server: Arc) -> Self { + Self { + server, + sessions: RwLock::new(HashMap::new()), + } + } + + async fn create_connection(&self) -> Result { + let (to_agent_tx, to_agent_rx) = mpsc::channel::(256); + let (from_agent_tx, from_agent_rx) = mpsc::channel::(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, 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, 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; +}