feat: WebSocket transport for goose-acp

crates/goose-acp/Cargo.toml

@@ -29,14 +29,15 @@ 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"] }
 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,17 +1,14 @@
 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;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter};
 
 #[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)
+}