tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

feat: ACP streamable http spec compliance (#9034)

Browse Source
This commit is contained in:
Alex Hancock
2026-05-06 09:47:50 -04:00
committed by GitHub
parent 8f35a4db7c
commit f9d0c6416d
4 changed files with 420 additions and 167 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/acp/transport/connection.rs
+259 -99
View File
@@ -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
+33 -17
View File
@@ -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
+2 -12
View File
@@ -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");