diff --git a/md/SUMMARY.md b/md/SUMMARY.md index f5cfb85..6d740c5 100644 --- a/md/SUMMARY.md +++ b/md/SUMMARY.md @@ -25,6 +25,7 @@ # Reference +- [Migrating to v2.0](./migration_v2.0.md) - [Migrating to v0.11](./migration_v0.11.x.md) - [P/ACP Specification](./proxying-acp.md) diff --git a/md/migration_v2.0.md b/md/migration_v2.0.md new file mode 100644 index 0000000..121a817 --- /dev/null +++ b/md/migration_v2.0.md @@ -0,0 +1,26 @@ +# Migrating from `agent-client-protocol` 1.x to 2.0 + +Version 2.0 changes the low-level in-process transport boundary so JSON-RPC frames remain intact +across components and adapters. + +## Raw channels carry frames + +`Channel` is now the single batch-aware in-process transport boundary. Its `rx` and `tx` carry +`TransportFrame`, not `Result`: + +```rust +use agent_client_protocol::{Channel, RawJsonRpcMessage, TransportFrame}; + +# fn send(channel: &Channel, message: RawJsonRpcMessage) { +channel.tx.unbounded_send(TransportFrame::Single(message)).unwrap(); +# } +``` + +`TransportFrame` distinguishes a valid single message, a non-empty `TransportBatch`, and an +explicit malformed wire value. Transport I/O failures are returned by the future from +`ConnectTo::into_channel_and_future`; they are never channel items. Components should preserve a +received frame intact when relaying it so batch response grouping remains correct. + +The separate, hidden `FramedChannel` and `into_framed_channel_and_future` compatibility path no +longer exist. Implement only `ConnectTo::connect_to`, optionally overriding +`into_channel_and_future` for a direct channel adapter. diff --git a/md/transport-architecture.md b/md/transport-architecture.md index 5a500ea..5e36fdd 100644 --- a/md/transport-architecture.md +++ b/md/transport-architecture.md @@ -1,12 +1,13 @@ # Transport Architecture -> **Note**: This document describes internal architecture and uses older terminology (e.g., `JrConnection` instead of the current API). For the user-facing API, see the [Core Library Design](./design.md) and the [`agent-client-protocol` rustdoc](https://docs.rs/agent-client-protocol). +For the broader user-facing API, see the [Core Library Design](./design.md) and +the [`agent-client-protocol` rustdoc](https://docs.rs/agent-client-protocol). This chapter explains how the connection layer separates protocol semantics from transport mechanisms, enabling flexible deployment patterns including in-process message passing. ## Overview -`JrConnection` provides the core JSON-RPC connection abstraction used by all ACP components. Originally designed around byte streams, it has been refactored to support **pluggable transports** that work with different I/O mechanisms while maintaining consistent protocol semantics. +The SDK's connection core provides the JSON-RPC abstraction used by ACP components. It supports **pluggable transports** that work with different I/O mechanisms while maintaining consistent protocol semantics. ## Design Principles @@ -32,7 +33,7 @@ This separation enables: - **Testability**: Mock transports for unit testing - **Clarity**: Clear boundaries between protocol and I/O concerns -### The `RawJsonRpcMessage` Boundary +### The `TransportFrame` Boundary The key insight is that `agent_client_protocol::RawJsonRpcMessage` provides a natural, transport-neutral boundary backed by the JSON-RPC envelope types from @@ -46,20 +47,22 @@ enum RawJsonRpcMessage { } ``` -This type sits between the protocol and transport layers: +`RawJsonRpcMessage` sits inside the transport-neutral public frame type: - **Above**: Protocol layer works with application types (`OutgoingMessage`, `UntypedMessage`) - **Below**: Transport layer parses and serializes `RawJsonRpcMessage` -- **Boundary**: An internal `TransportFrame` carries either one raw message or - the entries retained from one incoming batch - -The public `Channel` API remains a single-message `RawJsonRpcMessage` boundary. +- **Boundary**: `TransportFrame` carries one raw message, a structurally + non-empty batch, or a malformed wire value retained for a relay +- **In-process API**: `Channel::rx` and `Channel::tx` carry `TransportFrame` + directly, so adapters cannot accidentally flatten a batch +- **Failures**: I/O and connection failures are returned by the future driving + a transport; they are not sent as channel entries ## Actor Architecture -### Protocol Actors (Core JrConnection) +### Protocol Actors -These actors live in `JrConnection` and understand JSON-RPC semantics: +These actors live in the protocol connection core and understand JSON-RPC semantics: #### Outgoing Protocol Actor @@ -104,9 +107,9 @@ Manages request/response correlation: Runs user-spawned concurrent tasks via `cx.spawn()`. Unchanged from original design. -### Transport Actors (Provided by Trait) +### Transport Actors -These actors are spawned by `IntoJrConnectionTransport` implementations and have **zero knowledge** of protocol semantics: +These actors are driven by physical transport components and have **zero knowledge** of protocol semantics: #### Transport Outgoing Actor @@ -117,12 +120,12 @@ Output: Writes to I/O (byte stream, channel, socket, etc.) For byte streams: -- Serialize a single `RawJsonRpcMessage` or one response batch to JSON +- Serialize a single `RawJsonRpcMessage` or one non-empty batch to JSON - Write newline-delimited JSON to stream For in-process channels: -- Directly forward `RawJsonRpcMessage` to channel +- Directly forward `TransportFrame` to the channel #### Transport Incoming Actor @@ -142,13 +145,13 @@ For byte streams: For in-process channels: -- Directly forward `RawJsonRpcMessage` from channel +- Directly forward `TransportFrame` from the channel -The public `RawJsonRpcMessage` and `Channel` boundary intentionally remains -single-message. Batch tracking and response aggregation are internal to the -line and byte-stream transports. The SDK continues to initiate requests and -notifications as individual JSON-RPC messages; the only response arrays it -writes are correlated replies to batch calls received from the peer. +The public `Channel` boundary preserves complete frames. The SDK continues to +initiate requests and notifications as individual JSON-RPC messages; response +arrays are correlated replies to batch calls received from the peer. Relays and +instrumentation must forward frames intact so they do not change those wire +semantics. ## Message Flow @@ -199,63 +202,31 @@ When the conductor forwards messages between components, it must preserve send o Without this serialization, responses could overtake notifications when both are forwarded through proxy chains, causing the client to receive messages out of order. See [Conductor Implementation](./conductor.md#message-ordering-invariant) for details. -## Transport Trait +## Component Boundary -The `IntoJrConnectionTransport` trait defines how to bridge internal channels with I/O: +[`ConnectTo`](https://docs.rs/agent-client-protocol/latest/agent_client_protocol/trait.ConnectTo.html) +is the common component and transport abstraction. `connect_to` joins a +component to its counterpart and drives the connection until completion. +`into_channel_and_future` exposes the canonical low-level boundary as a +`Channel` plus the future that drives the component: -```rust -pub trait IntoJrConnectionTransport { - fn setup_transport( - self, - cx: &JrConnectionCx, - outgoing_rx: mpsc::UnboundedReceiver, - incoming_tx: mpsc::UnboundedSender, - ) -> Result<(), Error>; -} +```rust,ignore +fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<()>>); ``` -Key points: - -- **Consumed** (`self`): Implementations move owned resources into spawned actors -- **Spawns via `cx.spawn()`**: Uses connection context to spawn transport actors -- **Channels only**: No knowledge of `OutgoingMessage` or response correlation -- **Returns quickly**: Just spawns actors, doesn't block +The returned future owns transport failures and lifecycle completion. The +channel carries only `TransportFrame` wire events. Most components implement +only `connect_to`; direct transports override `into_channel_and_future` to avoid +an intermediate copy. ## Transport Implementations ### Byte Stream Transport -The default implementation works with any `AsyncRead` + `AsyncWrite` pair: - -```rust -impl IntoJrConnectionTransport for (OB, IB) { - fn setup_transport(self, cx, outgoing_rx, incoming_tx) -> Result<(), Error> { - let (outgoing_bytes, incoming_bytes) = self; - - // Spawn incoming: read bytes → parse JSON → send Message - cx.spawn(async move { - let mut lines = BufReader::new(incoming_bytes).lines(); - while let Some(line) = lines.next().await { - let message: RawJsonRpcMessage = serde_json::from_str(&line?)?; - incoming_tx.unbounded_send(message)?; - } - Ok(()) - }); - - // Spawn outgoing: receive Message → serialize → write bytes - cx.spawn(async move { - while let Some(message) = outgoing_rx.next().await { - let json = serde_json::to_vec(&message)?; - outgoing_bytes.write_all(&json).await?; - outgoing_bytes.write_all(b"\n").await?; - } - Ok(()) - }); - - Ok(()) - } -} -``` +`ByteStreams` works with `futures::io::AsyncWrite` and +`AsyncRead`. It adapts them to `Lines`, parses each incoming JSON value into one +frame, and serializes each outgoing frame to one newline-delimited JSON value. +`Stdio` and `AcpAgent` build on this transport. Use cases: @@ -264,37 +235,12 @@ Use cases: - Unix domain sockets - Any stream-based I/O -### In-Process Channel Transport - -For components in the same process, skip serialization entirely: +### In-Process Channel -```rust -pub struct ChannelTransport { - outgoing: mpsc::UnboundedSender, - incoming: mpsc::UnboundedReceiver, -} - -impl IntoJrConnectionTransport for ChannelTransport { - fn setup_transport(self, cx, outgoing_rx, incoming_tx) -> Result<(), Error> { - // Just forward messages, no serialization - cx.spawn(async move { - while let Some(message) = self.incoming.next().await { - incoming_tx.unbounded_send(message)?; - } - Ok(()) - }); - - cx.spawn(async move { - while let Some(message) = outgoing_rx.next().await { - self.outgoing.unbounded_send(message)?; - } - Ok(()) - }); - - Ok(()) - } -} -``` +For components in the same process, `Channel::duplex()` creates paired +endpoints and skips serialization entirely. Relays forward each received +`TransportFrame` without unpacking it; this preserves batch boundaries and the +original representation of malformed wire input. Benefits: @@ -302,122 +248,27 @@ Benefits: - **Same-process efficiency**: Ideal for conductor with in-process proxies - **Full type safety**: No parsing errors possible -## Construction API - -### Flexible Construction - -The refactored API separates handler setup from transport selection: - -```rust -// Build connection with handlers -let connection = JrConnection::new() - .name("my-component") - .on_receive_request(|req: InitializeRequest, cx| { - cx.respond(InitializeResponse::make()) - }) - .on_receive_notification(|notif: SessionNotification, _cx| { - Ok(()) - }); - -// Provide transport at the end -connection.serve_with(transport).await?; -``` - -### Byte Stream Convenience - -For the common case of byte streams, use the convenience constructor: - -```rust -JrConnection::from_streams(stdout, stdin) - .on_receive_request(...) - .serve() - .await?; -``` - -This is equivalent to: - -```rust -JrConnection::new() - .on_receive_request(...) - .serve_with((stdout, stdin)) - .await?; -``` - ## Use Cases ### 1. Standard Agent (Stdio) -Traditional subprocess agent with stdio communication: - -```rust -JrConnection::from_streams( - tokio::io::stdout().compat_write(), - tokio::io::stdin().compat() -) - .name("my-agent") - .on_receive_request(handle_prompt) - .serve() - .await?; -``` +Use `Stdio` for the current process or `AcpAgent` for a child process. Both use +the same frame-aware line transport underneath. ### 2. In-Process Proxy Chain -Conductor with proxies in the same process for maximum efficiency: - -```rust -// Create paired channel transports -let (transport_a, transport_b) = create_paired_transports(); - -// Spawn proxy in background -tokio::spawn(async move { - JrConnection::new() - .on_receive_message(proxy_handler) - .serve_with(transport_a) - .await -}); - -// Connect to proxy -JrConnection::new() - .on_receive_request(agent_handler) - .serve_with(transport_b) - .await?; -``` - -No serialization overhead between components! +Connect builders, proxies, and conductor components directly. Their default +`ConnectTo` adapter uses `Channel`, so complete frames cross each wrapper with +no serialization. ### 3. Network-Based Components -TCP socket connections between components: - -```rust -let stream = TcpStream::connect("localhost:8080").await?; -let (read, write) = stream.split(); - -JrConnection::new() - .on_receive_request(handler) - .serve_with((write.compat_write(), read.compat())) - .await?; -``` +Split the socket and pass compatible read/write halves to `ByteStreams::new`. ### 4. Testing with Mock Transport -Unit tests without real I/O: - -```rust -let (transport, mock) = create_mock_transport(); - -tokio::spawn(async move { - JrConnection::new() - .on_receive_request(my_handler) - .serve_with(transport) - .await -}); - -// Test by sending messages directly -mock.send_request("initialize", params).await?; -let response = mock.receive_response().await?; -assert_eq!(response.method, "initialized"); -``` +Use `Channel::duplex()` to inject and inspect `TransportFrame` values without +real I/O. ## Benefits @@ -445,14 +296,6 @@ assert_eq!(response.method, "initialized"); - **Focused implementations**: Each layer has single responsibility - **Maintainability**: Changes to transport don't affect protocol logic -## Implementation Status - -- ✅ **Phase 1**: Documentation complete -- 🚧 **Phase 2**: Actor splitting in progress -- 📋 **Phase 3**: Trait introduction planned -- 📋 **Phase 4**: In-process transport planned -- 📋 **Phase 5**: Conductor integration planned - ## Related Documentation - [Core Library Design](./design.md) - High-level crate architecture diff --git a/src/agent-client-protocol-conductor/src/snoop.rs b/src/agent-client-protocol-conductor/src/snoop.rs index 1779d1d..bc19465 100644 --- a/src/agent-client-protocol-conductor/src/snoop.rs +++ b/src/agent-client-protocol-conductor/src/snoop.rs @@ -1,4 +1,4 @@ -use agent_client_protocol::{ConnectTo, DynConnectTo, FramedChannel, RawJsonRpcMessage, Role}; +use agent_client_protocol::{Channel, ConnectTo, DynConnectTo, RawJsonRpcMessage, Role}; use futures_concurrency::future::TryJoin; pub struct SnooperComponent { @@ -36,9 +36,9 @@ impl ConnectTo for SnooperComponent { self, client: impl ConnectTo, ) -> Result<(), agent_client_protocol::Error> { - let (client_channel, client_future) = client.into_framed_channel_and_future(); - let (base_channel, base_future) = self.base_component.into_framed_channel_and_future(); - let snoop = FramedChannel::bridge_with_inspection( + let (client_channel, client_future) = client.into_channel_and_future(); + let (base_channel, base_future) = self.base_component.into_channel_and_future(); + let snoop = Channel::bridge_with_inspection( client_channel, base_channel, self.incoming_message, diff --git a/src/agent-client-protocol-http/CHANGELOG.md b/src/agent-client-protocol-http/CHANGELOG.md index cd5ea15..3ff45a3 100644 --- a/src/agent-client-protocol-http/CHANGELOG.md +++ b/src/agent-client-protocol-http/CHANGELOG.md @@ -2,6 +2,11 @@ ## [Unreleased] +### Fixed + +- Preserve incoming JSON-RPC batch frames and grouped responses across HTTP and WebSocket + transports, including session-aware HTTP routing. + ## [1.3.0](https://github.com/agentclientprotocol/rust-sdk/compare/agent-client-protocol-http-v1.2.0...agent-client-protocol-http-v1.3.0) - 2026-07-20 ### Fixed diff --git a/src/agent-client-protocol-http/src/client.rs b/src/agent-client-protocol-http/src/client.rs index 2c19922..d23bc53 100644 --- a/src/agent-client-protocol-http/src/client.rs +++ b/src/agent-client-protocol-http/src/client.rs @@ -4,7 +4,7 @@ use std::{ }; use agent_client_protocol::{ - Agent, Channel, Client, ConnectTo, Error as AcpError, RawJsonRpcMessage, + Agent, Channel, Client, ConnectTo, Error as AcpError, RawJsonRpcMessage, TransportFrame, schema::v1::{RequestId, Response as RpcResponse}, }; use async_tungstenite::tungstenite::Message as WsMessage; @@ -185,24 +185,23 @@ async fn run(client: HttpClient, channel: Channel) -> Result<(), AcpError> { } }; - let msg = match event { - HttpLoopEvent::Outgoing(msg) => match msg { - Some(Ok(msg)) => msg, - Some(Err(e)) => { - error!("upstream channel produced error: {e}"); - break Err(e); - } - None => { + let frame = match event { + HttpLoopEvent::Outgoing(msg) => { + let Some(frame) = msg else { outgoing_closed = true; continue; - } - }, + }; + frame + } HttpLoopEvent::SseEvent(event) => { let Some(event) = event else { continue; }; - let open_session_id = state.session_to_open_for_response(&event.message); - state.deliver(event.message); + let open_session_id = match &event.frame { + TransportFrame::Single(message) => state.session_to_open_for_response(message), + TransportFrame::Malformed { .. } | TransportFrame::Batch(_) => None, + }; + state.deliver_frame(event.frame); if let Some(session_id) = open_session_id { lifecycle.start_sse(Some(session_id), sse_event_tx.clone()); } @@ -228,6 +227,24 @@ async fn run(client: HttpClient, channel: Channel) -> Result<(), AcpError> { } }; + let msg = match frame { + TransportFrame::Single(message) => message, + frame @ (TransportFrame::Malformed { .. } | TransportFrame::Batch(_)) => { + if state.connection.connection_id().is_none() { + break Err(AcpError::invalid_request() + .data("ACP HTTP transport: first message must be `initialize`")); + } + match state.prepare_frame_post(frame) { + Ok(post) => response_posts.push(post), + Err(error) => { + error!("POST failed: {error}"); + break Err(AcpError::internal_error().data(format!("POST: {error}"))); + } + } + continue; + } + }; + if state.connection.connection_id().is_none() { if !is_initialize_request(&msg) { break Err(AcpError::invalid_request() @@ -270,7 +287,7 @@ async fn run(client: HttpClient, channel: Channel) -> Result<(), AcpError> { } enum HttpLoopEvent { - Outgoing(Option>), + Outgoing(Option), SseEvent(Option), SseFailure(SseFailure), Post(CompletedPost), @@ -284,7 +301,7 @@ struct SseFailure { #[derive(Debug)] struct SseMessage { - message: RawJsonRpcMessage, + frame: TransportFrame, } #[derive(Clone, Debug)] @@ -454,7 +471,7 @@ struct ClientState { connection: HttpConnection, open_session_streams: HashSet, pending_requests: HashMap, - incoming: futures::channel::mpsc::UnboundedSender>, + incoming: futures::channel::mpsc::UnboundedSender, } struct PendingPost { @@ -554,10 +571,17 @@ impl ClientState { return Err(format!("HTTP {status}: {body}")); } - let message = response - .json::() - .await - .map_err(|e| e.to_string())?; + let body = response.text().await.map_err(|error| error.to_string())?; + let message = match TransportFrame::parse_json(&body) { + Some(TransportFrame::Single(message)) => message, + Some(TransportFrame::Malformed { error, .. }) => { + return Err(format!("invalid initialize response: {error}")); + } + Some(TransportFrame::Batch(_)) => { + return Err("initialize response must not be a JSON-RPC batch".to_string()); + } + None => return Err("initialize response was ignored as malformed".to_string()), + }; if matches!( message, @@ -619,6 +643,34 @@ impl ClientState { }) } + fn prepare_frame_post(&self, frame: TransportFrame) -> Result { + let connection_id = self + .connection + .connection_id() + .ok_or_else(|| "POST attempted before initialize".to_string())?; + let body = frame.to_json().map_err(|error| error.to_string())?; + let request = self + .connection + .post() + .header("Content-Type", "application/json") + .header("Accept", "application/json") + .header(HEADER_CONNECTION_ID, connection_id) + .body(body); + let response = async move { + let response = request.send().await.map_err(|error| error.to_string())?; + if response.status().as_u16() != 202 && !response.status().is_success() { + let status = response.status(); + let body = response.text().await.unwrap_or_default(); + return Err(format!("HTTP {status}: {body}")); + } + Ok(()) + }; + Ok(PendingPost { + pending_request: None, + response: response.boxed(), + }) + } + fn remove_pending_request(&mut self, pending_request: Option<&(RequestId, String)>) { if let Some((id, _)) = pending_request { self.pending_requests.remove(id); @@ -651,7 +703,11 @@ impl ClientState { } fn deliver(&self, msg: RawJsonRpcMessage) { - if self.incoming.unbounded_send(Ok(msg)).is_err() { + self.deliver_frame(TransportFrame::Single(msg)); + } + + fn deliver_frame(&self, frame: TransportFrame) { + if self.incoming.unbounded_send(frame).is_err() { debug!("upstream channel closed; dropping inbound message"); } } @@ -690,13 +746,12 @@ async fn read_sse( if payload.is_empty() { continue; } - let msg = serde_json::from_str::(&payload) - .map_err(|e| format!("malformed JSON-RPC payload: {e}"))?; + let Some(frame) = TransportFrame::parse_json(&payload) else { + debug!("ignoring malformed response-shaped SSE payload"); + continue; + }; - if event_tx - .unbounded_send(SseMessage { message: msg }) - .is_err() - { + if event_tx.unbounded_send(SseMessage { frame }).is_err() { return Err("upstream channel closed".to_string()); } } @@ -765,25 +820,17 @@ where tx: incoming, } = channel; let writer = async move { - while let Some(msg) = outgoing.next().await { - match msg { - Ok(msg) => { - let text = match serde_json::to_string(&msg) { - Ok(t) => t, - Err(e) => { - error!("failed to serialize outbound message: {e}"); - return Err(AcpError::internal_error().data(format!("serialize: {e}"))); - } - }; - if let Err(e) = ws_tx.send(WsMessage::Text(text.into())).await { - error!("WebSocket send failed: {e}"); - return Err(AcpError::internal_error().data(format!("ws send: {e}"))); - } - } - Err(e) => { - error!("upstream channel produced error: {e}"); - return Err(e); + while let Some(frame) = outgoing.next().await { + let text = match frame.to_json() { + Ok(text) => text, + Err(error) => { + error!("failed to serialize outbound frame: {error}"); + return Err(AcpError::internal_error().data(format!("serialize: {error}"))); } + }; + if let Err(error) = ws_tx.send(WsMessage::Text(text.into())).await { + error!("WebSocket send failed: {error}"); + return Err(AcpError::internal_error().data(format!("ws send: {error}"))); } } @@ -799,28 +846,15 @@ where if discard_incoming { continue; } - match serde_json::from_str::(text.as_str()) { - Ok(parsed) => { - if incoming.unbounded_send(Ok(parsed)).is_err() { - debug!( - "upstream channel closed; discarding WS input while draining output" - ); - discard_incoming = true; - } - } - Err(e) => { - let message = format!("malformed JSON-RPC payload: {e}"); - warn!("WS: {message}"); - if incoming - .unbounded_send(Err(AcpError::parse_error().data(message))) - .is_err() - { - debug!( - "upstream channel closed; discarding WS input while draining output" - ); - discard_incoming = true; - } - } + let Some(frame) = TransportFrame::parse_json(text.as_str()) else { + debug!("ignoring malformed response-shaped WebSocket payload"); + continue; + }; + if incoming.unbounded_send(frame).is_err() { + debug!( + "upstream channel closed; discarding WS input while draining output" + ); + discard_incoming = true; } } Some(Ok(WsMessage::Binary(_))) => { @@ -861,7 +895,7 @@ mod tests { time::Duration, }; - use agent_client_protocol::schema::v1::RequestId; + use agent_client_protocol::{TransportBatch, schema::v1::RequestId}; use axum::{ Json, Router, extract::{WebSocketUpgrade, ws::Message as AxumWsMessage}, @@ -882,13 +916,13 @@ mod tests { finish: Arc, finished: Arc, escaped_tx: futures::channel::oneshot::Sender< - futures::channel::mpsc::UnboundedSender>, + futures::channel::mpsc::UnboundedSender, >, } struct QueueOutgoingThenText { text: Option, - outgoing: Option>>, + outgoing: Option>, } struct RecordingWsSink(mpsc::UnboundedSender); @@ -904,6 +938,30 @@ mod tests { release: Option>, } + fn single_frame(message: RawJsonRpcMessage) -> TransportFrame { + TransportFrame::Single(message) + } + + fn into_single_message(frame: TransportFrame) -> Result { + match frame { + TransportFrame::Single(message) => Ok(message), + TransportFrame::Malformed { error, .. } => Err(error), + TransportFrame::Batch(_) => { + Err(AcpError::internal_error().data("expected one JSON-RPC message")) + } + } + } + + trait TransportFrameTestExt { + fn unwrap(self) -> RawJsonRpcMessage; + } + + impl TransportFrameTestExt for TransportFrame { + fn unwrap(self) -> RawJsonRpcMessage { + into_single_message(self).unwrap() + } + } + impl WsSink for RecordingWsSink { async fn send(&mut self, message: WsMessage) -> Result<(), String> { self.0 @@ -942,11 +1000,9 @@ mod tests { if let Some(outgoing) = self.outgoing.take() { for method in ["custom/first", "custom/second"] { outgoing - .unbounded_send(Ok(RawJsonRpcMessage::notification( - method.to_string(), - json!({}), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::notification(method.to_string(), json!({})).unwrap(), + )) .unwrap(); } } @@ -988,27 +1044,27 @@ mod tests { })?; channel .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .map_err(|e| { AcpError::internal_error().data(format!("send initialize: {e}")) })?; - channel.rx.next().await.ok_or_else(|| { + into_single_message(channel.rx.next().await.ok_or_else(|| { AcpError::internal_error().data("initialize response channel closed") - })??; + })?)?; for method in ["custom/first", "custom/second"] { channel .tx - .unbounded_send(Ok(RawJsonRpcMessage::notification( - method.to_string(), - json!({}), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::notification(method.to_string(), json!({})).unwrap(), + )) .map_err(|e| { AcpError::internal_error().data(format!("send {method}: {e}")) })?; @@ -1108,12 +1164,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .unwrap(); let init_response = timeout(Duration::from_secs(1), caller.rx.next()) .await @@ -1124,14 +1182,16 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::notification( - "$/cancel_request".to_string(), - json!({ - "requestId": 2, - "sessionId": "session-1" - }), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::notification( + "$/cancel_request".to_string(), + json!({ + "requestId": 2, + "sessionId": "session-1" + }), + ) + .unwrap(), + )) .unwrap(); let (session_header, message) = timeout(Duration::from_secs(1), capture_rx.recv()) @@ -1155,6 +1215,141 @@ mod tests { server.abort(); } + #[tokio::test] + async fn http_preserves_batch_frames_across_post_and_sse() { + let (post_tx, mut post_rx) = tokio::sync::mpsc::unbounded_channel(); + let post_count = Arc::new(AtomicUsize::new(0)); + let emit_sse = Arc::new(Notify::new()); + let inbound_batch = json!([ + { + "jsonrpc": "2.0", + "method": "custom/inbound-one", + "params": {} + }, + { + "jsonrpc": "2.0", + "method": "custom/inbound-two", + "params": {} + } + ]); + let app = Router::new().route( + "/acp", + post({ + let post_count = post_count.clone(); + move |body: String| { + let post_count = post_count.clone(); + let post_tx = post_tx.clone(); + async move { + if post_count.fetch_add(1, Ordering::SeqCst) == 0 { + return initialize_response().await.into_response(); + } + + post_tx + .send(serde_json::from_str::(&body).unwrap()) + .unwrap(); + StatusCode::ACCEPTED.into_response() + } + } + }) + .get({ + let emit_sse = emit_sse.clone(); + let inbound_batch = inbound_batch.clone(); + move || { + let emit_sse = emit_sse.clone(); + let inbound_batch = inbound_batch.clone(); + async move { + let stream = async_stream::stream! { + emit_sse.notified().await; + yield Ok::<_, Infallible>( + Event::default().data(inbound_batch.to_string()), + ); + futures::future::pending::<()>().await; + }; + Sse::new(stream) + } + } + }) + .delete(|| async { StatusCode::ACCEPTED }), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let client = HttpClient::new(format!("http://{addr}")).unwrap(); + let (mut caller, transport) = Channel::duplex(); + let transport = tokio::spawn(run(client, transport)); + + caller + .tx + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) + .unwrap(); + let init_response = timeout(Duration::from_secs(1), caller.rx.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert!(matches!(init_response, RawJsonRpcMessage::Response(_))); + + let outbound_batch = json!([ + { + "jsonrpc": "2.0", + "method": "custom/outbound-one", + "params": {} + }, + { + "jsonrpc": "2.0", + "method": "custom/outbound-two", + "params": {} + } + ]); + caller + .tx + .unbounded_send(TransportFrame::Batch( + TransportBatch::from_messages([ + RawJsonRpcMessage::notification("custom/outbound-one".to_string(), json!({})) + .unwrap(), + RawJsonRpcMessage::notification("custom/outbound-two".to_string(), json!({})) + .unwrap(), + ]) + .unwrap(), + )) + .unwrap(); + + let posted = timeout(Duration::from_secs(1), post_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(posted, outbound_batch); + + emit_sse.notify_one(); + let inbound = timeout(Duration::from_secs(1), caller.rx.next()) + .await + .unwrap() + .unwrap(); + assert!(matches!(&inbound, TransportFrame::Batch(_))); + assert_eq!( + serde_json::from_str::(&inbound.to_json().unwrap()).unwrap(), + inbound_batch + ); + + drop(caller); + timeout(Duration::from_secs(1), transport) + .await + .unwrap() + .unwrap() + .unwrap(); + + server.abort(); + } + #[tokio::test] async fn custom_response_with_session_id_does_not_open_session_sse() { let (get_tx, mut get_rx) = tokio::sync::mpsc::unbounded_channel(); @@ -1218,12 +1413,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .unwrap(); let init_response = timeout(Duration::from_secs(1), caller.rx.next()) .await @@ -1240,12 +1437,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "custom/sessionish".to_string(), - json!({}), - RequestId::Number(2), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "custom/sessionish".to_string(), + json!({}), + RequestId::Number(2), + ) + .unwrap(), + )) .unwrap(); let response = timeout(Duration::from_secs(1), caller.rx.next()) .await @@ -1343,12 +1542,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .unwrap(); let init_response = timeout(Duration::from_secs(1), caller.rx.next()) .await @@ -1365,12 +1566,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "session/fork".to_string(), - json!({ "sessionId": "source-session" }), - RequestId::Number(2), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "session/fork".to_string(), + json!({ "sessionId": "source-session" }), + RequestId::Number(2), + ) + .unwrap(), + )) .unwrap(); let source_sse_header = timeout(Duration::from_secs(1), get_rx.recv()) .await @@ -1485,11 +1688,10 @@ mod tests { ); assert!( escaped - .unbounded_send(Ok(RawJsonRpcMessage::notification( - "custom/too-late".to_string(), - json!({}), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::notification("custom/too-late".to_string(), json!({}),) + .unwrap() + )) .is_err(), "escaped client sender remained open after client completion" ); @@ -1593,12 +1795,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .unwrap(); let init_response = timeout(Duration::from_secs(1), caller.rx.next()) .await @@ -1612,12 +1816,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "custom/slow".to_string(), - json!({}), - RequestId::Number(2), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "custom/slow".to_string(), + json!({}), + RequestId::Number(2), + ) + .unwrap(), + )) .unwrap(); let callback = timeout(Duration::from_secs(1), caller.rx.next()) @@ -1634,7 +1840,7 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::response( + .unbounded_send(single_frame(RawJsonRpcMessage::response( RequestId::Number(99), Ok(json!({})), ))) @@ -1686,12 +1892,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .unwrap(); let init_response = timeout(Duration::from_secs(1), caller.rx.next()) .await @@ -1702,12 +1910,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "session/prompt".to_string(), - json!({}), - RequestId::Number(2), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "session/prompt".to_string(), + json!({}), + RequestId::Number(2), + ) + .unwrap(), + )) .unwrap(); let error = timeout(Duration::from_secs(1), transport) .await @@ -1746,12 +1956,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .unwrap(); let init_response = timeout(Duration::from_secs(1), caller.rx.next()) .await @@ -1773,7 +1985,7 @@ mod tests { } #[tokio::test] - async fn malformed_sse_json_fails_transport() { + async fn malformed_sse_json_is_delivered_and_transport_continues() { let delete_count = Arc::new(AtomicUsize::new(0)); let delete_count_for_handler = delete_count.clone(); let app = Router::new().route( @@ -1799,12 +2011,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .unwrap(); let init_response = timeout(Duration::from_secs(1), caller.rx.next()) .await @@ -1813,13 +2027,22 @@ mod tests { .unwrap(); assert!(matches!(init_response, RawJsonRpcMessage::Response(_))); - let error = timeout(Duration::from_secs(1), transport) + let frame = timeout(Duration::from_secs(1), caller.rx.next()) .await .unwrap() - .unwrap() - .unwrap_err(); + .unwrap(); - assert!(error.to_string().contains("malformed JSON-RPC payload")); + let TransportFrame::Malformed { raw, error } = frame else { + panic!("expected malformed frame, got {frame:?}"); + }; + assert_eq!(raw, "{not json"); + assert_eq!(error.code, AcpError::parse_error().code); + drop(caller); + timeout(Duration::from_secs(1), transport) + .await + .unwrap() + .unwrap() + .unwrap(); assert_eq!(delete_count.load(Ordering::SeqCst), 1); server.abort(); @@ -1837,12 +2060,15 @@ mod tests { let (mut caller, transport) = Channel::duplex(); let transport = tokio::spawn(run(client, transport)); - let error = timeout(Duration::from_secs(1), caller.rx.next()) + let frame = timeout(Duration::from_secs(1), caller.rx.next()) .await .unwrap() - .unwrap() - .unwrap_err(); - assert!(error.to_string().contains("malformed JSON-RPC payload")); + .unwrap(); + let TransportFrame::Malformed { raw, error } = frame else { + panic!("expected malformed frame, got {frame:?}"); + }; + assert_eq!(raw, "{not json"); + assert_eq!(error.code, AcpError::parse_error().code); let message = timeout(Duration::from_secs(1), caller.rx.next()) .await @@ -1861,6 +2087,52 @@ mod tests { server.abort(); } + #[tokio::test] + async fn websocket_serializes_batch_as_one_text_frame() { + let (caller, transport) = Channel::duplex(); + let Channel { + tx: outgoing, + rx: incoming, + } = caller; + drop(incoming); + outgoing + .unbounded_send(TransportFrame::Batch( + TransportBatch::from_messages([ + RawJsonRpcMessage::notification("custom/first".to_string(), json!({})).unwrap(), + RawJsonRpcMessage::notification("custom/second".to_string(), json!({})) + .unwrap(), + ]) + .unwrap(), + )) + .unwrap(); + drop(outgoing); + + let (ws_output_tx, mut ws_output) = mpsc::unbounded(); + timeout( + Duration::from_secs(1), + drive_ws( + RecordingWsSink(ws_output_tx), + futures::stream::pending::>(), + transport, + ), + ) + .await + .unwrap() + .unwrap(); + let frames = ws_output.by_ref().collect::>().await; + + let WsMessage::Text(text) = &frames[0] else { + panic!("batch was not sent as WebSocket text"); + }; + let batch = serde_json::from_str::(text.as_str()).unwrap(); + let entries = batch.as_array().expect("batch should remain an array"); + assert_eq!(entries.len(), 2); + assert_eq!(entries[0]["method"], "custom/first"); + assert_eq!(entries[1]["method"], "custom/second"); + assert!(matches!(frames.get(1), Some(WsMessage::Close(None)))); + assert_eq!(frames.len(), 2); + } + #[tokio::test] async fn websocket_drain_discards_incoming_after_receiver_closes() { let (caller, transport) = Channel::duplex(); @@ -1916,11 +2188,9 @@ mod tests { } = caller; drop(incoming); outgoing - .unbounded_send(Ok(RawJsonRpcMessage::notification( - "custom/queued".to_string(), - json!({}), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::notification("custom/queued".to_string(), json!({})).unwrap(), + )) .unwrap(); drop(outgoing); @@ -1999,12 +2269,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .unwrap(); let init_response = timeout(Duration::from_secs(1), async { tokio::select! { @@ -2055,12 +2327,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .unwrap(); let init_response = timeout(Duration::from_secs(1), caller.rx.next()) .await @@ -2104,12 +2378,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .unwrap(); let init_response = timeout(Duration::from_secs(1), caller.rx.next()) .await @@ -2161,12 +2437,14 @@ mod tests { caller .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( - "initialize".to_string(), - json!({}), - RequestId::Number(1), - ) - .unwrap())) + .unbounded_send(single_frame( + RawJsonRpcMessage::request( + "initialize".to_string(), + json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .unwrap(); let error = timeout(Duration::from_secs(1), transport) .await diff --git a/src/agent-client-protocol-http/src/connection.rs b/src/agent-client-protocol-http/src/connection.rs index 4023364..7163988 100644 --- a/src/agent-client-protocol-http/src/connection.rs +++ b/src/agent-client-protocol-http/src/connection.rs @@ -3,14 +3,17 @@ use std::{ sync::{Arc, Weak}, }; -use agent_client_protocol::{Channel, RawJsonRpcMessage, schema::v1::RequestId}; +use agent_client_protocol::{ + Channel, RawJsonRpcMessage, TransportBatch, TransportBatchEntry, TransportFrame, + schema::v1::RequestId, +}; use futures::{SinkExt, StreamExt}; use tokio::sync::{Mutex, RwLock, mpsc, watch}; use tracing::{debug, error, trace}; use crate::protocol::session_id_from_message; -#[derive(Clone, Debug)] +#[derive(Clone, Debug, PartialEq, Eq)] pub(crate) enum ResponseRoute { Connection, Session(String), @@ -24,7 +27,7 @@ enum OutboundTransport { struct HttpOutbound { connection_stream: Arc, session_streams: RwLock>>, - pending_routes: Mutex>, + pending_routes: Mutex>>, } struct WebSocketOutbound { @@ -85,8 +88,8 @@ impl OutboundStream { pub(crate) const OUTBOUND_STREAM_CAPACITY: usize = 1024; pub(crate) struct Connection { - inbound_tx: mpsc::UnboundedSender>, - outbound_rx: Mutex>>, + inbound_tx: mpsc::UnboundedSender, + outbound_rx: Mutex>>, agent_handle: Mutex>>, router_handle: Mutex>>, closed_tx: watch::Sender, @@ -94,9 +97,9 @@ pub(crate) struct Connection { } impl Connection { - pub(crate) fn send_to_agent(&self, msg: RawJsonRpcMessage) -> Result<(), &'static str> { + pub(crate) fn send_frame_to_agent(&self, frame: TransportFrame) -> Result<(), &'static str> { self.inbound_tx - .send(Ok(msg)) + .send(frame) .map_err(|_| "agent channel closed") } @@ -153,14 +156,24 @@ impl Connection { })); } - async fn route_outbound(&self, msg: RawJsonRpcMessage) { - self.outbound_transport.route_outbound(msg).await; + async fn route_outbound(&self, frame: TransportFrame) { + self.outbound_transport.route_outbound(frame).await; } pub(crate) async fn recv_initial(&self) -> Option { let mut guard = self.outbound_rx.lock().await; let rx = guard.as_mut()?; - rx.recv().await + match rx.recv().await? { + TransportFrame::Single(message) => Some(message), + TransportFrame::Malformed { error, .. } => { + error!(?error, "agent emitted malformed initial JSON-RPC input"); + None + } + TransportFrame::Batch(_) => { + error!("agent emitted an invalid batched initial JSON-RPC message"); + None + } + } } pub(crate) async fn shutdown(&self) { @@ -236,18 +249,38 @@ impl OutboundTransport { http.connection_stream.push(msg).await; } - async fn route_outbound(&self, msg: RawJsonRpcMessage) { - let serialized = match serde_json::to_string(&msg) { - Ok(s) => s, - Err(e) => { - error!("failed to serialize outbound JSON-RPC message: {e}"); - return; + async fn route_outbound(&self, frame: TransportFrame) { + match frame { + TransportFrame::Single(message) => { + let serialized = match serde_json::to_string(&message) { + Ok(serialized) => serialized, + Err(error) => { + error!("failed to serialize outbound JSON-RPC message: {error}"); + return; + } + }; + match self { + Self::Http(http) => http.route_outbound(&message, serialized).await, + Self::WebSocket(websocket) => websocket.all_outbound.push(serialized).await, + } + } + TransportFrame::Malformed { raw, .. } => match self { + Self::Http(http) => http.connection_stream.push(raw).await, + Self::WebSocket(websocket) => websocket.all_outbound.push(raw).await, + }, + TransportFrame::Batch(batch) => { + let serialized = match serde_json::to_string(&batch) { + Ok(serialized) => serialized, + Err(error) => { + error!("failed to serialize outbound JSON-RPC batch: {error}"); + return; + } + }; + match self { + Self::Http(http) => http.route_outbound_batch(&batch, serialized).await, + Self::WebSocket(websocket) => websocket.all_outbound.push(serialized).await, + } } - }; - - match self { - Self::Http(http) => http.route_outbound(&msg, serialized).await, - Self::WebSocket(websocket) => websocket.all_outbound.push(serialized).await, } } } @@ -263,7 +296,12 @@ impl HttpOutbound { async fn record_pending_route(&self, id: RequestId, route: ResponseRoute) { if let Some(key) = pending_route_key(&id) { - self.pending_routes.lock().await.insert(key, route); + self.pending_routes + .lock() + .await + .entry(key) + .or_default() + .push_back(route); } } @@ -292,7 +330,10 @@ impl HttpOutbound { } RawJsonRpcMessage::Response(_) => { let route = match msg.response_id().and_then(pending_route_key) { - Some(key) => self.pending_routes.lock().await.remove(&key), + Some(key) => { + let mut pending_routes = self.pending_routes.lock().await; + take_pending_route(&mut pending_routes, &key) + } None => None, }; route.unwrap_or(ResponseRoute::Connection) @@ -310,6 +351,47 @@ impl HttpOutbound { } } } + + async fn route_outbound_batch(&self, batch: &TransportBatch, serialized: String) { + let mut pending_routes = self.pending_routes.lock().await; + let mut common_route = None; + let mut routes_disagree = false; + for entry in batch.entries() { + let route = match entry { + TransportBatchEntry::Message(message) => message + .response_id() + .and_then(pending_route_key) + .and_then(|key| take_pending_route(&mut pending_routes, &key)) + .unwrap_or(ResponseRoute::Connection), + TransportBatchEntry::Malformed { .. } => ResponseRoute::Connection, + }; + match &common_route { + None => common_route = Some(route), + Some(common_route) if common_route == &route => {} + Some(_) => routes_disagree = true, + } + } + drop(pending_routes); + + let route = if routes_disagree { + ResponseRoute::Connection + } else { + common_route.unwrap_or(ResponseRoute::Connection) + }; + match route { + ResponseRoute::Connection => { + trace!(target = "connection", "→ connection-scoped batch stream"); + self.connection_stream.push(serialized).await; + } + ResponseRoute::Session(session_id) => { + trace!(target = %session_id, "→ session-scoped batch stream"); + self.session_stream(&session_id) + .await + .push(serialized) + .await; + } + } + } } impl WebSocketOutbound { @@ -391,9 +473,8 @@ impl ConnectionRegistry { outbound_transport: OutboundTransport, ) -> Arc { let (channel, agent_future) = self.factory.spawn_agent(); - let (inbound_tx, mut inbound_rx) = - mpsc::unbounded_channel::>(); - let (outbound_tx, outbound_rx) = mpsc::unbounded_channel::(); + let (inbound_tx, mut inbound_rx) = mpsc::unbounded_channel::(); + let (outbound_tx, outbound_rx) = mpsc::unbounded_channel::(); let (closed_tx, _) = watch::channel(false); let Channel { @@ -413,17 +494,9 @@ impl ConnectionRegistry { let inbound_abort_for_outbound = inbound_abort.clone(); let outbound = async move { while let Some(msg) = agent_rx.next().await { - match msg { - Ok(m) => { - if outbound_tx.send(m).is_err() { - break; - } - } - Err(e) => { - error!("agent emitted error: {e}"); - inbound_abort_for_outbound.abort(); - break; - } + if outbound_tx.send(msg).is_err() { + inbound_abort_for_outbound.abort(); + break; } } }; @@ -505,10 +578,24 @@ fn pending_route_key(id: &RequestId) -> Option { } } +fn take_pending_route( + pending_routes: &mut HashMap>, + key: &RequestId, +) -> Option { + let routes = pending_routes.get_mut(key)?; + let route = routes.pop_front(); + let remove_entry = routes.is_empty(); + if remove_entry { + pending_routes.remove(key); + } + route +} + #[cfg(test)] mod tests { use std::sync::Arc; + use agent_client_protocol::TransportBatch; use futures::future::BoxFuture; use tokio::{ sync::Notify, @@ -553,7 +640,7 @@ mod tests { let future = Box::pin(async move { agent .tx - .unbounded_send(Ok(RawJsonRpcMessage::response( + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::response( RequestId::Number(1), Ok(serde_json::json!({ "done": true })), ))) @@ -565,11 +652,11 @@ mod tests { } } - struct ErrorThenWaitAgentFactory { + struct MalformedThenWaitAgentFactory { emit: Arc, } - impl AgentFactory for ErrorThenWaitAgentFactory { + impl AgentFactory for MalformedThenWaitAgentFactory { fn spawn_agent( &self, ) -> ( @@ -582,9 +669,11 @@ mod tests { emit.notified().await; agent .tx - .unbounded_send(Err( - agent_client_protocol::Error::parse_error().data("transport parse error") - )) + .unbounded_send(TransportFrame::Malformed { + raw: "{not json".to_string(), + error: agent_client_protocol::Error::parse_error() + .data("transport parse error"), + }) .unwrap(); std::future::pending::>().await }); @@ -609,7 +698,49 @@ mod tests { let message = self.message.clone(); let exit = self.exit.clone(); let future = Box::pin(async move { - agent.tx.unbounded_send(Ok(message)).unwrap(); + agent + .tx + .unbounded_send(TransportFrame::Single(message)) + .unwrap(); + exit.notified().await; + Ok(()) + }); + + (transport, future) + } + } + + struct BatchThenWaitAgentFactory { + exit: Arc, + } + + impl AgentFactory for BatchThenWaitAgentFactory { + fn spawn_agent( + &self, + ) -> ( + Channel, + BoxFuture<'static, agent_client_protocol::Result<()>>, + ) { + let (agent, transport) = Channel::duplex(); + let exit = self.exit.clone(); + let future = Box::pin(async move { + let batch = TransportBatch::from_messages([ + RawJsonRpcMessage::notification( + "test/first".to_string(), + serde_json::json!({}), + ) + .unwrap(), + RawJsonRpcMessage::notification( + "test/second".to_string(), + serde_json::json!({}), + ) + .unwrap(), + ]) + .expect("test batch is non-empty"); + agent + .tx + .unbounded_send(TransportFrame::Batch(batch)) + .unwrap(); exit.notified().await; Ok(()) }); @@ -643,30 +774,29 @@ mod tests { } #[tokio::test] - async fn agent_error_removes_connection_and_closes_streams() { + async fn malformed_frame_is_relayed_without_closing_connection() { let emit = Arc::new(Notify::new()); - let registry = - ConnectionRegistry::new(Arc::new(ErrorThenWaitAgentFactory { emit: emit.clone() })); + let registry = ConnectionRegistry::new(Arc::new(MalformedThenWaitAgentFactory { + emit: emit.clone(), + })); let (connection_id, connection) = registry.create_connection().await; - let mut closed = connection.subscribe_closed(); + let (_replay, mut outbound) = connection.subscribe_connection_stream().await; assert!(registry.get(&connection_id).await.is_some()); connection.start_router().await; emit.notify_one(); - timeout(Duration::from_secs(1), async { - loop { - if *closed.borrow() { - break; - } - closed.changed().await.unwrap(); - } - }) - .await - .unwrap(); + let raw = timeout(Duration::from_secs(1), outbound.recv()) + .await + .unwrap() + .expect("malformed frame should be relayed"); + assert_eq!(raw, "{not json"); + assert!(registry.get(&connection_id).await.is_some()); + assert!(!*connection.subscribe_closed().borrow()); - assert!(registry.get(&connection_id).await.is_none()); + registry.remove(&connection_id).await; + connection.shutdown().await; } #[tokio::test] @@ -738,6 +868,65 @@ mod tests { connection.shutdown().await; } + #[tokio::test] + async fn batch_is_relayed_as_one_connection_stream_frame() { + let exit = Arc::new(Notify::new()); + let registry = + ConnectionRegistry::new(Arc::new(BatchThenWaitAgentFactory { exit: exit.clone() })); + let (_connection_id, connection) = registry.create_connection().await; + let (_replay, mut connection_rx) = connection.subscribe_connection_stream().await; + + connection.start_router().await; + + let text = timeout(Duration::from_secs(1), connection_rx.recv()) + .await + .unwrap() + .expect("batch should reach the connection stream"); + let batch = serde_json::from_str::(&text).unwrap(); + let entries = batch.as_array().expect("batch should remain an array"); + assert_eq!(entries.len(), 2); + assert_eq!(entries[0]["method"], "test/first"); + assert_eq!(entries[1]["method"], "test/second"); + assert!(connection_rx.try_recv().is_err()); + + exit.notify_one(); + connection.shutdown().await; + } + + #[tokio::test] + async fn duplicate_batch_response_ids_consume_each_pending_route() { + let outbound = HttpOutbound::new(); + let (_connection_replay, mut connection_rx) = outbound.connection_stream.subscribe().await; + let (_session_replay, mut session_rx) = + outbound.session_stream("session-1").await.subscribe().await; + + let id = RequestId::Number(21); + let route = ResponseRoute::Session("session-1".to_string()); + outbound + .record_pending_route(id.clone(), route.clone()) + .await; + outbound.record_pending_route(id.clone(), route).await; + + let batch = TransportBatch::from_messages([ + RawJsonRpcMessage::response(id.clone(), Ok(serde_json::json!({ "slot": 1 }))), + RawJsonRpcMessage::response(id, Ok(serde_json::json!({ "slot": 2 }))), + ]) + .expect("duplicate-ID response batch is non-empty"); + let serialized = serde_json::to_string(&batch).unwrap(); + + outbound + .route_outbound_batch(&batch, serialized.clone()) + .await; + + assert_eq!( + timeout(Duration::from_secs(1), session_rx.recv()) + .await + .unwrap(), + Some(serialized) + ); + assert!(connection_rx.try_recv().is_err()); + } + #[tokio::test] async fn http_connection_does_not_retain_all_outbound_replay() { let exit = Arc::new(Notify::new()); diff --git a/src/agent-client-protocol-http/src/http_server.rs b/src/agent-client-protocol-http/src/http_server.rs index bb30ac1..4d4481f 100644 --- a/src/agent-client-protocol-http/src/http_server.rs +++ b/src/agent-client-protocol-http/src/http_server.rs @@ -1,6 +1,8 @@ use std::{convert::Infallible, error::Error as _, sync::Arc, time::Duration}; -use agent_client_protocol::{RawJsonRpcMessage, schema::v1::Response as RpcResponse}; +use agent_client_protocol::{ + RawJsonRpcMessage, TransportBatchEntry, TransportFrame, schema::v1::Response as RpcResponse, +}; use axum::{ body::Body, extract::State, @@ -54,22 +56,36 @@ pub(crate) async fn handle_post( } }; - if matches!(body.first(), Some(&b'[')) { - return StatusCode::NOT_IMPLEMENTED.into_response(); - } - - let mut message = match serde_json::from_slice::(&body) { - Ok(message) => message, - Err(e) => { - return (StatusCode::BAD_REQUEST, format!("Invalid JSON-RPC: {e}")).into_response(); + let body = match std::str::from_utf8(&body) { + Ok(body) => body, + Err(error) => { + return ( + StatusCode::BAD_REQUEST, + format!("Invalid JSON-RPC: {error}"), + ) + .into_response(); } }; + let Some(mut frame) = TransportFrame::parse_json(body) else { + // Response-shaped invalid input is deliberately ignored by the JSON-RPC + // parser because responses must not themselves receive responses. + return StatusCode::ACCEPTED.into_response(); + }; - if is_initialize_request(&message) { + if matches!( + &frame, + TransportFrame::Single(message) if is_initialize_request(message) + ) { + let TransportFrame::Single(message) = frame else { + unreachable!("initialize was matched as a single frame"); + }; let (connection_id, connection) = registry.create_connection().await; let initialize_cleanup = InitializeCleanup::new(registry.clone(), connection_id.clone(), connection.clone()); - if connection.send_to_agent(message).is_err() { + if connection + .send_frame_to_agent(TransportFrame::Single(message)) + .is_err() + { initialize_cleanup.cleanup().await; return StatusCode::INTERNAL_SERVER_ERROR.into_response(); } @@ -113,38 +129,82 @@ pub(crate) async fn handle_post( return StatusCode::NOT_FOUND.into_response(); }; - if let Some(session_id) = &session_id - && method_for_message(&message).is_some() - && let Err(error) = apply_session_header_to_message(&mut message, session_id) + let mut session_routes = Vec::new(); + let mut pending_routes = Vec::new(); + match &mut frame { + TransportFrame::Single(message) => { + let route = match prepare_message_route(message, session_id.as_deref()) { + Ok(route) => route, + Err(error) => return (StatusCode::BAD_REQUEST, error).into_response(), + }; + collect_route(message, route, &mut session_routes, &mut pending_routes); + trace!(connection_id = %connection_id, ?message, "POST → agent"); + } + TransportFrame::Batch(batch) => { + for entry in batch.entries_mut() { + let TransportBatchEntry::Message(message) = entry else { + continue; + }; + let route = match prepare_message_route(message, session_id.as_deref()) { + Ok(route) => route, + Err(error) => return (StatusCode::BAD_REQUEST, error).into_response(), + }; + collect_route(message, route, &mut session_routes, &mut pending_routes); + } + trace!(connection_id = %connection_id, ?frame, "POST batch → agent"); + } + TransportFrame::Malformed { .. } => { + trace!(connection_id = %connection_id, ?frame, "POST malformed frame → agent"); + } + } + + for session_id in session_routes { + connection.ensure_session(&session_id).await; + } + for (request_id, route) in pending_routes { + connection.record_pending_route(request_id, route).await; + } + + if connection.send_frame_to_agent(frame).is_err() { + return StatusCode::INTERNAL_SERVER_ERROR.into_response(); + } + StatusCode::ACCEPTED.into_response() +} + +fn prepare_message_route( + message: &mut RawJsonRpcMessage, + session_id: Option<&str>, +) -> Result, &'static str> { + if let Some(session_id) = session_id + && method_for_message(message).is_some() { - return (StatusCode::BAD_REQUEST, error).into_response(); + apply_session_header_to_message(message, session_id)?; } - let route = match method_for_message(&message) { - Some(method) => match session_id_from_message(&message) { + Ok(match method_for_message(message) { + Some(method) => match session_id_from_message(message) { Some(session_id) => Some(ResponseRoute::Session(session_id)), None if method_requires_session_header(method) => { - return (StatusCode::BAD_REQUEST, "Acp-Session-Id header required").into_response(); + return Err("Acp-Session-Id header required"); } None => Some(ResponseRoute::Connection), }, None => None, - }; + }) +} +fn collect_route( + message: &RawJsonRpcMessage, + route: Option, + session_routes: &mut Vec, + pending_routes: &mut Vec<(agent_client_protocol::schema::v1::RequestId, ResponseRoute)>, +) { if let Some(ResponseRoute::Session(session_id)) = &route { - connection.ensure_session(session_id).await; + session_routes.push(session_id.clone()); } - if let (RawJsonRpcMessage::Request(req), Some(route)) = (&message, route) { - connection.record_pending_route(req.id.clone(), route).await; - trace!(connection_id = %connection_id, method = %req.method, "POST → agent"); - } else { - trace!(connection_id = %connection_id, ?message, "POST → agent"); + if let (RawJsonRpcMessage::Request(request), Some(route)) = (message, route) { + pending_routes.push((request.id.clone(), route)); } - - if connection.send_to_agent(message).is_err() { - return StatusCode::INTERNAL_SERVER_ERROR.into_response(); - } - StatusCode::ACCEPTED.into_response() } struct InitializeCleanup { @@ -336,7 +396,10 @@ fn post_body_too_large_response() -> Response { mod tests { use std::sync::Arc; - use agent_client_protocol::{Channel, RawJsonRpcMessage, schema::v1::RequestId}; + use agent_client_protocol::{ + Channel, RawJsonRpcMessage, TransportBatch, TransportBatchEntry, TransportFrame, + schema::v1::RequestId, + }; use futures::{StreamExt, future::BoxFuture}; use serde_json::json; use tokio::{ @@ -365,8 +428,11 @@ mod tests { rx: mut incoming, tx: _, } = agent; - while let Some(message) = incoming.next().await { - if forwarded.send(message?).is_err() { + while let Some(frame) = incoming.next().await { + let TransportFrame::Single(message) = frame else { + panic!("expected a single JSON-RPC frame"); + }; + if forwarded.send(message).is_err() { break; } } @@ -388,10 +454,12 @@ mod tests { ) { let (mut agent, transport) = Channel::duplex(); let future = Box::pin(async move { - if let Some(Ok(RawJsonRpcMessage::Request(request))) = agent.rx.next().await { + if let Some(TransportFrame::Single(RawJsonRpcMessage::Request(request))) = + agent.rx.next().await + { agent .tx - .unbounded_send(Ok(RawJsonRpcMessage::response( + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::response( request.id, Err(agent_client_protocol::Error::invalid_request() .data("initialize rejected")), @@ -428,6 +496,56 @@ mod tests { } } + struct BatchAgentFactory { + forwarded: mpsc::UnboundedSender)>>, + } + + impl AgentFactory for BatchAgentFactory { + fn spawn_agent( + &self, + ) -> ( + Channel, + BoxFuture<'static, agent_client_protocol::Result<()>>, + ) { + let (mut agent, transport) = Channel::duplex(); + let forwarded = self.forwarded.clone(); + let future = Box::pin(async move { + let Some(TransportFrame::Batch(batch)) = agent.rx.next().await else { + panic!("expected one batch frame"); + }; + let mut methods = Vec::new(); + let mut responses = Vec::new(); + for entry in batch.entries() { + let TransportBatchEntry::Message(RawJsonRpcMessage::Request(request)) = entry + else { + panic!("expected a request batch entry"); + }; + methods.push(( + request.method.to_string(), + request + .params + .as_ref() + .and_then(crate::protocol::session_id_from_params), + )); + responses.push(RawJsonRpcMessage::response( + request.id.clone(), + Ok(json!({ "ok": true })), + )); + } + forwarded.send(methods).unwrap(); + let responses = + TransportBatch::from_messages(responses).expect("responses are non-empty"); + agent + .tx + .unbounded_send(TransportFrame::Batch(responses)) + .unwrap(); + std::future::pending::>().await + }); + + (transport, future) + } + } + #[tokio::test] async fn post_rejects_declared_body_larger_than_limit() { let (forwarded_tx, _forwarded_rx) = mpsc::unbounded_channel(); @@ -450,6 +568,142 @@ mod tests { assert_eq!(response.status(), StatusCode::PAYLOAD_TOO_LARGE); } + #[tokio::test] + async fn post_applies_session_header_to_batch_and_routes_grouped_response() { + let (forwarded_tx, mut forwarded_rx) = mpsc::unbounded_channel(); + let registry = Arc::new(ConnectionRegistry::new(Arc::new(BatchAgentFactory { + forwarded: forwarded_tx, + }))); + let (connection_id, connection) = registry.create_connection().await; + let (_connection_replay, mut connection_outbound) = + connection.subscribe_connection_stream().await; + let (_session_replay, mut session_outbound) = + connection.subscribe_session_stream("session-1").await; + connection.start_router().await; + + let body = json!([ + { + "jsonrpc": "2.0", + "id": 1, + "method": "custom/first", + "params": {} + }, + { + "jsonrpc": "2.0", + "id": 2, + "method": "custom/second", + "params": {} + } + ]) + .to_string(); + let request = Request::builder() + .method("POST") + .uri("/acp") + .header(header::CONTENT_TYPE, JSON_MIME_TYPE) + .header(HEADER_CONNECTION_ID, connection_id.as_str()) + .header(HEADER_SESSION_ID, "session-1") + .body(Body::from(body)) + .unwrap(); + + let response = handle_post(State(registry.clone()), request).await; + + assert_eq!(response.status(), StatusCode::ACCEPTED); + let methods = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!( + methods, + [ + ("custom/first".to_string(), Some("session-1".to_string())), + ("custom/second".to_string(), Some("session-1".to_string())), + ] + ); + let response = timeout(Duration::from_secs(1), session_outbound.recv()) + .await + .unwrap() + .expect("batch response should be emitted"); + let response = serde_json::from_str::(&response).unwrap(); + let entries = response.as_array().expect("response should remain a batch"); + assert_eq!(entries.len(), 2); + assert_eq!(entries[0]["id"], 1); + assert_eq!(entries[1]["id"], 2); + assert!(session_outbound.try_recv().is_err()); + assert!(connection_outbound.try_recv().is_err()); + + registry.remove(&connection_id).await; + connection.shutdown().await; + } + + #[tokio::test] + async fn mixed_batch_routes_fall_back_to_connection_stream() { + let (forwarded_tx, mut forwarded_rx) = mpsc::unbounded_channel(); + let registry = Arc::new(ConnectionRegistry::new(Arc::new(BatchAgentFactory { + forwarded: forwarded_tx, + }))); + let (connection_id, connection) = registry.create_connection().await; + let (_connection_replay, mut connection_outbound) = + connection.subscribe_connection_stream().await; + let (_session_replay, mut session_outbound) = + connection.subscribe_session_stream("session-1").await; + connection.start_router().await; + + let body = json!([ + { + "jsonrpc": "2.0", + "id": 1, + "method": "custom/session", + "params": {} + }, + { + "jsonrpc": "2.0", + "id": 2, + "method": "$/connection", + "params": {} + } + ]) + .to_string(); + let request = Request::builder() + .method("POST") + .uri("/acp") + .header(header::CONTENT_TYPE, JSON_MIME_TYPE) + .header(HEADER_CONNECTION_ID, connection_id.as_str()) + .header(HEADER_SESSION_ID, "session-1") + .body(Body::from(body)) + .unwrap(); + + let response = handle_post(State(registry.clone()), request).await; + + assert_eq!(response.status(), StatusCode::ACCEPTED); + let methods = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!( + methods, + [ + ("custom/session".to_string(), Some("session-1".to_string())), + ("$/connection".to_string(), None), + ] + ); + let response = timeout(Duration::from_secs(1), connection_outbound.recv()) + .await + .unwrap() + .expect("mixed-route batch should use the connection stream"); + assert_eq!( + serde_json::from_str::(&response) + .unwrap() + .as_array() + .unwrap() + .len(), + 2 + ); + assert!(session_outbound.try_recv().is_err()); + + registry.remove(&connection_id).await; + connection.shutdown().await; + } + #[tokio::test] async fn initialize_error_response_rejects_connection() { let registry = Arc::new(ConnectionRegistry::new(Arc::new( @@ -687,4 +941,45 @@ mod tests { } connection.shutdown().await; } + + #[tokio::test] + async fn post_rejects_batch_session_scoped_method_without_session_id() { + let (forwarded_tx, mut forwarded_rx) = mpsc::unbounded_channel(); + let registry = Arc::new(ConnectionRegistry::new(Arc::new(CapturingAgentFactory { + forwarded: forwarded_tx, + }))); + let (connection_id, connection) = registry.create_connection().await; + let body = json!([ + { + "jsonrpc": "2.0", + "id": 1, + "method": "custom/valid", + "params": {} + }, + { + "jsonrpc": "2.0", + "id": 2, + "method": "session/delete", + "params": {} + } + ]) + .to_string(); + let request = Request::builder() + .method("POST") + .uri("/acp") + .header(header::CONTENT_TYPE, JSON_MIME_TYPE) + .header(HEADER_CONNECTION_ID, connection_id.as_str()) + .body(Body::from(body)) + .unwrap(); + + let response = handle_post(State(registry), request).await; + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + let body = axum::body::to_bytes(response.into_body(), 1024) + .await + .unwrap(); + assert_eq!(body.as_ref(), b"Acp-Session-Id header required"); + assert!(forwarded_rx.try_recv().is_err()); + connection.shutdown().await; + } } diff --git a/src/agent-client-protocol-http/src/websocket_server.rs b/src/agent-client-protocol-http/src/websocket_server.rs index 65a7fa2..658ace1 100644 --- a/src/agent-client-protocol-http/src/websocket_server.rs +++ b/src/agent-client-protocol-http/src/websocket_server.rs @@ -1,6 +1,6 @@ use std::sync::Arc; -use agent_client_protocol::{Error as AcpError, RawJsonRpcMessage, schema::v1::RequestId}; +use agent_client_protocol::{RawJsonRpcMessage, TransportFrame}; use axum::{ extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade}, http::HeaderValue, @@ -75,35 +75,20 @@ async fn run_ws( Some(Ok(WsMessage::Text(text))) => { let text_str = text.to_string(); trace!(connection_id = %connection_id, payload = %text_str, "Client → Agent: {} bytes", text_str.len()); - match serde_json::from_str::(&text_str) { - Ok(parsed) => { - if let Some(sid) = session_id_from_message(&parsed) - && let RawJsonRpcMessage::Request(req) = &parsed { + match TransportFrame::parse_json(&text_str) { + Some(frame) => { + if let TransportFrame::Single(parsed) = &frame + && let Some(sid) = session_id_from_message(parsed) + && let RawJsonRpcMessage::Request(req) = parsed { trace!(connection_id = %connection_id, session_id = %sid, request_id = ?req.id, "Client → Agent (session)"); } - if connection.send_to_agent(parsed).is_err() { + if connection.send_frame_to_agent(frame).is_err() { error!(connection_id = %connection_id, "Agent channel closed"); break; } } - Err(e) => { - let message = format!("malformed JSON-RPC payload: {e}"); - warn!(connection_id = %connection_id, "Returning parse error for malformed JSON-RPC frame: {e}"); - let response = RawJsonRpcMessage::response( - RequestId::Null, - Err(AcpError::parse_error().data(message)), - ); - let text = match serde_json::to_string(&response) { - Ok(text) => text, - Err(e) => { - error!(connection_id = %connection_id, "Failed to serialize parse error response: {e}"); - break; - } - }; - if ws_tx.send(WsMessage::Text(text.into())).await.is_err() { - error!(connection_id = %connection_id, "WebSocket send failed"); - break; - } + None => { + trace!(connection_id = %connection_id, "Ignoring malformed response-shaped JSON-RPC frame"); } } } @@ -153,7 +138,7 @@ async fn run_ws( #[cfg(test)] mod tests { use agent_client_protocol::{ - Channel, + Channel, TransportBatch, TransportBatchEntry, TransportFrame, schema::v1::{RequestId, Response as RpcResponse}, }; use async_tungstenite::{tokio::connect_async, tungstenite::Message as ClientWsMessage}; @@ -186,11 +171,23 @@ mod tests { let future = Box::pin(async move { let Channel { rx: mut incoming, - tx: _, + tx: outgoing, } = agent; - while let Some(message) = incoming.next().await { - if forwarded.send(message?).is_err() { - break; + while let Some(frame) = incoming.next().await { + match frame { + TransportFrame::Single(message) => { + if forwarded.send(message).is_err() { + break; + } + } + TransportFrame::Malformed { error, .. } => { + outgoing + .unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::response(RequestId::Null, Err(error)), + )) + .unwrap(); + } + TransportFrame::Batch(_) => panic!("expected a single JSON-RPC frame"), } } Ok(()) @@ -200,6 +197,50 @@ mod tests { } } + struct BatchAgentFactory { + forwarded: mpsc::UnboundedSender>, + } + + impl AgentFactory for BatchAgentFactory { + fn spawn_agent( + &self, + ) -> ( + Channel, + BoxFuture<'static, agent_client_protocol::Result<()>>, + ) { + let (mut agent, transport) = Channel::duplex(); + let forwarded = self.forwarded.clone(); + let future = Box::pin(async move { + let Some(TransportFrame::Batch(batch)) = agent.rx.next().await else { + panic!("expected one batch frame"); + }; + let mut methods = Vec::new(); + let mut responses = Vec::new(); + for entry in batch.entries() { + let TransportBatchEntry::Message(RawJsonRpcMessage::Request(request)) = entry + else { + panic!("expected a request batch entry"); + }; + methods.push(request.method.to_string()); + responses.push(RawJsonRpcMessage::response( + request.id.clone(), + Ok(json!({ "ok": true })), + )); + } + forwarded.send(methods).unwrap(); + let responses = + TransportBatch::from_messages(responses).expect("responses are non-empty"); + agent + .tx + .unbounded_send(TransportFrame::Batch(responses)) + .unwrap(); + std::future::pending::>().await + }); + + (transport, future) + } + } + #[tokio::test] async fn malformed_ws_frame_returns_parse_error_response_and_continues() { let (forwarded_tx, mut forwarded_rx) = mpsc::unbounded_channel(); @@ -239,12 +280,7 @@ mod tests { let value: serde_json::Value = serde_json::from_str(&text).unwrap(); assert_eq!(value["id"], serde_json::Value::Null); assert_eq!(value["error"]["code"], -32700); - assert!( - value["error"]["data"] - .as_str() - .unwrap() - .contains("malformed JSON-RPC payload") - ); + assert_eq!(value["error"]["data"]["line"], "{not json"); let parsed = serde_json::from_value::(value).unwrap(); assert!(matches!( @@ -276,4 +312,68 @@ mod tests { server.abort(); } + + #[tokio::test] + async fn websocket_forwards_batch_as_one_frame_and_emits_grouped_response() { + let (forwarded_tx, mut forwarded_rx) = mpsc::unbounded_channel(); + let registry = Arc::new(ConnectionRegistry::new(Arc::new(BatchAgentFactory { + forwarded: forwarded_tx, + }))); + let app = Router::new().route( + "/acp", + get({ + let registry = registry.clone(); + move |ws: WebSocketUpgrade| { + let registry = registry.clone(); + async move { handle_ws_upgrade(registry, ws) } + } + }), + ); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let server = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + let (mut client, _) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let batch = json!([ + { + "jsonrpc": "2.0", + "id": 1, + "method": "custom/first", + "params": {} + }, + { + "jsonrpc": "2.0", + "id": 2, + "method": "custom/second", + "params": {} + } + ]); + + client + .send(ClientWsMessage::Text(batch.to_string().into())) + .await + .unwrap(); + + let methods = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(methods, ["custom/first", "custom/second"]); + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + let ClientWsMessage::Text(text) = frame else { + panic!("expected text frame: {frame:?}"); + }; + let response = serde_json::from_str::(&text).unwrap(); + let entries = response.as_array().expect("response should remain a batch"); + assert_eq!(entries.len(), 2); + assert_eq!(entries[0]["id"], 1); + assert_eq!(entries[1]["id"], 2); + + server.abort(); + } } diff --git a/src/agent-client-protocol-polyfill/CHANGELOG.md b/src/agent-client-protocol-polyfill/CHANGELOG.md index a722d01..ec82789 100644 --- a/src/agent-client-protocol-polyfill/CHANGELOG.md +++ b/src/agent-client-protocol-polyfill/CHANGELOG.md @@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Fixed + +- Preserve JSON-RPC batch frames through the MCP-over-ACP HTTP bridge and keep + malformed POST errors correlated with their originating HTTP requests. + ## [1.0.1](https://github.com/agentclientprotocol/rust-sdk/compare/agent-client-protocol-polyfill-v1.0.0...agent-client-protocol-polyfill-v1.0.1) - 2026-06-29 ### Other diff --git a/src/agent-client-protocol-polyfill/src/mcp_over_acp/http.rs b/src/agent-client-protocol-polyfill/src/mcp_over_acp/http.rs index 88a8fea..a60bb44 100644 --- a/src/agent-client-protocol-polyfill/src/mcp_over_acp/http.rs +++ b/src/agent-client-protocol-polyfill/src/mcp_over_acp/http.rs @@ -1,7 +1,8 @@ //! HTTP-based MCP bridge transport. use agent_client_protocol::{ - BoxFuture, Channel, ConnectTo, RawJsonRpcMessage, RawJsonRpcParams, + BoxFuture, Channel, ConnectTo, RawJsonRpcMessage, RawJsonRpcParams, TransportBatchEntry, + TransportFrame, role::mcp, schema::v1::{ Notification as RpcNotification, Request as RpcRequest, RequestId, Response as RpcResponse, @@ -153,7 +154,7 @@ enum HttpMessage { Request { http_request_id: uuid::Uuid, request: RpcRequest, - response_tx: mpsc::UnboundedSender, + response_tx: mpsc::UnboundedSender, }, /// A JSON-RPC notification (no id, no response expected). Notification { @@ -165,23 +166,32 @@ enum HttpMessage { http_request_id: uuid::Uuid, response: RpcResponse, }, + /// A batch retained as one transport frame. + Frame { + http_request_id: uuid::Uuid, + frame: TransportFrame, + request_ids: Vec, + response_tx: Option>, + }, /// A GET request to open an SSE stream for server-initiated messages. Get { http_request_id: uuid::Uuid, - response_tx: mpsc::UnboundedSender, + response_tx: mpsc::UnboundedSender, }, } struct RunningServer { waiting_sessions: FxHashMap, + waiting_batch_sessions: Vec, general_sessions: Vec, - message_deque: VecDeque, + message_deque: VecDeque, } impl RunningServer { fn new() -> Self { RunningServer { waiting_sessions: HashMap::default(), + waiting_batch_sessions: Vec::new(), general_sessions: Vec::default(), message_deque: VecDeque::with_capacity(32), } @@ -196,7 +206,7 @@ impl RunningServer { #[derive(Debug)] enum MultiplexMessage { FromHttpToChannel(HttpMessage), - FromChannelToHttp(Result), + FromChannelToHttp(TransportFrame), } let mut merged_stream = http_rx @@ -211,9 +221,6 @@ impl RunningServer { self.handle_http_message(http_message, &mut channel.tx)?; } MultiplexMessage::FromChannelToHttp(message) => { - let message = message.unwrap_or_else(|err| { - RawJsonRpcMessage::response(RequestId::Null, Err(err)) - }); self.message_deque.push_back(message); } } @@ -228,9 +235,7 @@ impl RunningServer { fn handle_http_message( &mut self, message: HttpMessage, - channel_tx: &mut mpsc::UnboundedSender< - Result, - >, + channel_tx: &mut mpsc::UnboundedSender, ) -> Result<(), agent_client_protocol::Error> { match message { HttpMessage::Request { @@ -241,7 +246,7 @@ impl RunningServer { tracing::debug!(%http_request_id, ?request, "handling request"); let request_id = request.id.clone(); channel_tx - .unbounded_send(Ok(RawJsonRpcMessage::Request(request))) + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::Request(request))) .map_err(agent_client_protocol::util::internal_error)?; let session = RegisteredSession::new(response_tx); self.waiting_sessions.insert(request_id, session); @@ -251,7 +256,9 @@ impl RunningServer { request, } => { channel_tx - .unbounded_send(Ok(RawJsonRpcMessage::Notification(request))) + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::Notification( + request, + ))) .map_err(agent_client_protocol::util::internal_error)?; } HttpMessage::Response { @@ -259,7 +266,34 @@ impl RunningServer { response, } => { channel_tx - .unbounded_send(Ok(RawJsonRpcMessage::Response(response))) + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::Response( + response, + ))) + .map_err(agent_client_protocol::util::internal_error)?; + } + HttpMessage::Frame { + http_request_id, + frame, + request_ids, + response_tx, + } => { + tracing::debug!(%http_request_id, ?frame, "handling retained frame"); + if let Some(response_tx) = response_tx { + let session = RegisteredSession::new(response_tx); + match &frame { + TransportFrame::Batch(_) => { + self.waiting_batch_sessions.push(WaitingBatchSession { + request_ids, + session, + }); + } + TransportFrame::Single(_) | TransportFrame::Malformed { .. } => { + unreachable!("only batches use the retained frame variant") + } + } + } + channel_tx + .unbounded_send(frame) .map_err(agent_client_protocol::util::internal_error)?; } HttpMessage::Get { @@ -285,9 +319,53 @@ impl RunningServer { fn try_dispatch_jsonrpc_message( &mut self, - mut message: RawJsonRpcMessage, - ) -> Option { - let message_id = message.response_id().cloned(); + mut message: TransportFrame, + ) -> Option { + if matches!(message, TransportFrame::Malformed { .. }) { + // Malformed frames emitted by a relay are wire data, not protocol + // responses, so they are delivered through a general stream. + } else if matches!(message, TransportFrame::Batch(_)) { + let response_ids: Vec<_> = match &message { + TransportFrame::Batch(batch) => batch + .entries() + .filter_map(|entry| match entry { + TransportBatchEntry::Message(message) => message.response_id(), + TransportBatchEntry::Malformed { .. } => None, + }) + .collect(), + _ => unreachable!(), + }; + let correlated = self.waiting_batch_sessions.iter().position(|waiting| { + !waiting.request_ids.is_empty() + && waiting + .request_ids + .iter() + .any(|id| response_ids.contains(&id)) + }); + let fallback = self + .waiting_batch_sessions + .iter() + .position(|waiting| waiting.request_ids.is_empty()); + if let Some(index) = correlated.or(fallback) { + let session = self.waiting_batch_sessions.remove(index).session; + match session.outgoing_tx.unbounded_send(message) { + Ok(()) => return None, + Err(error) => { + assert!(error.is_disconnected()); + message = error.into_inner(); + } + } + } + } + + let message_id = match &message { + TransportFrame::Single(message) => message.response_id().cloned(), + TransportFrame::Malformed { .. } => None, + TransportFrame::Batch(batch) => batch.entries().find_map(|entry| match entry { + TransportBatchEntry::Message(message) => message.response_id().cloned(), + TransportBatchEntry::Malformed { .. } => None, + }), + }; if let Some(ref message_id) = message_id && let Some(session) = self.waiting_sessions.remove(message_id) @@ -305,7 +383,12 @@ impl RunningServer { let all_sessions = self .general_sessions .iter_mut() - .chain(self.waiting_sessions.values_mut()); + .chain(self.waiting_sessions.values_mut()) + .chain( + self.waiting_batch_sessions + .iter_mut() + .map(|waiting| &mut waiting.session), + ); for session in all_sessions { match session.outgoing_tx.unbounded_send(message) { Ok(()) => return None, @@ -324,17 +407,24 @@ impl RunningServer { .retain(|session| !session.outgoing_tx.is_closed()); self.waiting_sessions .retain(|_, session| !session.outgoing_tx.is_closed()); + self.waiting_batch_sessions + .retain(|waiting| !waiting.session.outgoing_tx.is_closed()); } } +struct WaitingBatchSession { + request_ids: Vec, + session: RegisteredSession, +} + struct RegisteredSession { #[allow(dead_code)] id: uuid::Uuid, - outgoing_tx: mpsc::UnboundedSender, + outgoing_tx: mpsc::UnboundedSender, } impl RegisteredSession { - fn new(outgoing_tx: mpsc::UnboundedSender) -> Self { + fn new(outgoing_tx: mpsc::UnboundedSender) -> Self { Self { id: uuid::Uuid::new_v4(), outgoing_tx, @@ -342,58 +432,113 @@ impl RegisteredSession { } } -/// Accept a POST request carrying a JSON-RPC message from an MCP client. -/// For requests (messages with id), we return an SSE stream. -/// For notifications/responses (messages without id), we return 202 Accepted. +/// Accept a POST request carrying a JSON-RPC frame from an MCP client. +/// For response-bearing calls and batches, we return an SSE stream. For +/// notification/response-only frames, we return 202 Accepted. async fn handle_post( State(state): State>, body: String, ) -> Result { let http_request_id = uuid::Uuid::new_v4(); - let message: RawJsonRpcMessage = - serde_json::from_str(&body).map_err(agent_client_protocol::util::parse_error)?; - - match message { - RawJsonRpcMessage::Request(request) => { - let (tx, mut rx) = mpsc::unbounded(); - state - .registration_tx - .unbounded_send(HttpMessage::Request { - http_request_id, - request, - response_tx: tx, - }) - .map_err(agent_client_protocol::util::internal_error)?; + let Some(frame) = TransportFrame::parse_json(&body) else { + return Ok(StatusCode::ACCEPTED.into_response()); + }; - let stream = async_stream::stream! { - while let Some(message) = rx.next().await { - match axum::response::sse::Event::default().json_data(message) { - Ok(v) => yield Ok(v), - Err(e) => yield Err(HttpError::from(e)), + match frame { + TransportFrame::Single(message) => match message { + RawJsonRpcMessage::Request(request) => { + let (tx, rx) = mpsc::unbounded(); + state + .registration_tx + .unbounded_send(HttpMessage::Request { + http_request_id, + request, + response_tx: tx, + }) + .map_err(agent_client_protocol::util::internal_error)?; + + Ok(sse_response(rx)) + } + RawJsonRpcMessage::Notification(request) => { + state + .registration_tx + .unbounded_send(HttpMessage::Notification { + http_request_id, + request, + }) + .map_err(agent_client_protocol::util::internal_error)?; + Ok(StatusCode::ACCEPTED.into_response()) + } + RawJsonRpcMessage::Response(response) => { + state + .registration_tx + .unbounded_send(HttpMessage::Response { + http_request_id, + response, + }) + .map_err(agent_client_protocol::util::internal_error)?; + Ok(StatusCode::ACCEPTED.into_response()) + } + }, + TransportFrame::Malformed { error, .. } => Ok(immediate_sse_response( + TransportFrame::Single(RawJsonRpcMessage::response(RequestId::Null, Err(error))), + )), + TransportFrame::Batch(batch) => { + if batch + .entries() + .all(|entry| matches!(entry, TransportBatchEntry::Malformed { .. })) + { + let responses = agent_client_protocol::TransportBatch::from_messages( + batch.entries().map(|entry| { + let TransportBatchEntry::Malformed { error, .. } = entry else { + unreachable!("all batch entries were checked as malformed") + }; + RawJsonRpcMessage::response(RequestId::Null, Err(error.clone())) + }), + ) + .expect("a TransportBatch is non-empty"); + return Ok(immediate_sse_response(TransportFrame::Batch(responses))); + } + + let mut request_ids = Vec::new(); + let mut expects_response = false; + for entry in batch.entries() { + match entry { + TransportBatchEntry::Message(RawJsonRpcMessage::Request(request)) => { + request_ids.push(request.id.clone()); + expects_response = true; } + TransportBatchEntry::Malformed { .. } => expects_response = true, + TransportBatchEntry::Message( + RawJsonRpcMessage::Notification(_) | RawJsonRpcMessage::Response(_), + ) => {} } - }; - Ok(Sse::new(stream).into_response()) - } - RawJsonRpcMessage::Notification(request) => { - state - .registration_tx - .unbounded_send(HttpMessage::Notification { - http_request_id, - request, - }) - .map_err(agent_client_protocol::util::internal_error)?; - Ok(StatusCode::ACCEPTED.into_response()) - } - RawJsonRpcMessage::Response(response) => { - state - .registration_tx - .unbounded_send(HttpMessage::Response { - http_request_id, - response, - }) - .map_err(agent_client_protocol::util::internal_error)?; - Ok(StatusCode::ACCEPTED.into_response()) + } + let frame = TransportFrame::Batch(batch); + if expects_response { + let (tx, rx) = mpsc::unbounded(); + state + .registration_tx + .unbounded_send(HttpMessage::Frame { + http_request_id, + frame, + request_ids, + response_tx: Some(tx), + }) + .map_err(agent_client_protocol::util::internal_error)?; + Ok(sse_response(rx)) + } else { + state + .registration_tx + .unbounded_send(HttpMessage::Frame { + http_request_id, + frame, + request_ids, + response_tx: None, + }) + .map_err(agent_client_protocol::util::internal_error)?; + Ok(StatusCode::ACCEPTED.into_response()) + } } } } @@ -415,12 +560,329 @@ async fn handle_get( let stream = async_stream::stream! { while let Some(message) = rx.next().await { - match axum::response::sse::Event::default().json_data(message) { - Ok(v) => yield Ok(v), - Err(e) => yield Err(HttpError::from(e)), - } + yield sse_event(message); } }; Ok(Sse::new(stream)) } + +fn sse_event(frame: TransportFrame) -> Result { + Ok(axum::response::sse::Event::default().data(frame.to_json()?)) +} + +fn sse_response(mut rx: mpsc::UnboundedReceiver) -> Response { + let stream = async_stream::stream! { + while let Some(message) = rx.next().await { + yield sse_event(message); + } + }; + Sse::new(stream).into_response() +} + +fn immediate_sse_response(frame: TransportFrame) -> Response { + Sse::new(futures::stream::once(async move { sse_event(frame) })).into_response() +} + +#[cfg(test)] +mod tests { + use super::*; + + async fn single_sse_payload(response: Response) -> serde_json::Value { + let body = axum::body::to_bytes(response.into_body(), 64 * 1024) + .await + .expect("SSE response body"); + let body = std::str::from_utf8(&body).expect("UTF-8 SSE response"); + let payload = body + .lines() + .find_map(|line| line.strip_prefix("data:").map(str::trim_start)) + .expect("one SSE data event"); + serde_json::from_str(payload).expect("JSON-RPC SSE payload") + } + + async fn single_sse_message(response: Response) -> RawJsonRpcMessage { + serde_json::from_value(single_sse_payload(response).await) + .expect("single JSON-RPC SSE message") + } + + #[test] + fn malformed_post_cannot_steal_a_valid_null_id_response() { + futures::executor::block_on(async { + let (registration_tx, mut registration_rx) = mpsc::unbounded(); + let state = Arc::new(BridgeState { registration_tx }); + + let valid_http_response = handle_post( + State(state.clone()), + serde_json::json!({ + "jsonrpc": "2.0", + "id": null, + "method": "example", + "params": {} + }) + .to_string(), + ) + .await + .expect("valid null-ID POST"); + let malformed_http_response = handle_post(State(state), "{not json".to_owned()) + .await + .expect("malformed POST receives a JSON-RPC error"); + + let valid_request = registration_rx + .next() + .await + .expect("valid request is forwarded"); + assert!( + registration_rx.try_recv().is_err(), + "malformed input must be answered by its own HTTP request" + ); + + let mut server = RunningServer::new(); + let (mut channel_tx, mut channel_rx) = mpsc::unbounded(); + server + .handle_http_message(valid_request, &mut channel_tx) + .expect("forward valid request"); + assert!(matches!( + channel_rx.next().await, + Some(TransportFrame::Single(RawJsonRpcMessage::Request(_))) + )); + + let valid_response = TransportFrame::Single(RawJsonRpcMessage::response( + RequestId::Null, + Ok(serde_json::json!({ "source": "valid" })), + )); + assert!( + server + .try_dispatch_jsonrpc_message(valid_response) + .is_none() + ); + + assert!(matches!( + single_sse_message(valid_http_response).await, + RawJsonRpcMessage::Response(RpcResponse::Result { + id: RequestId::Null, + result, + .. + }) if result == serde_json::json!({ "source": "valid" }) + )); + assert!(matches!( + single_sse_message(malformed_http_response).await, + RawJsonRpcMessage::Response(RpcResponse::Error { + id: RequestId::Null, + error, + .. + }) if error.code == agent_client_protocol::ErrorCode::ParseError + )); + }); + } + + #[test] + fn forwards_batch_and_routes_grouped_response_without_flattening() { + futures::executor::block_on(async { + let mut server = RunningServer::new(); + let (mut channel_tx, mut channel_rx) = mpsc::unbounded(); + let (response_tx, mut response_rx) = mpsc::unbounded(); + let incoming = TransportFrame::Batch( + agent_client_protocol::TransportBatch::from_messages([RawJsonRpcMessage::request( + "example".into(), + serde_json::json!({}), + RequestId::Number(7), + ) + .unwrap()]) + .unwrap(), + ); + + server + .handle_http_message( + HttpMessage::Frame { + http_request_id: uuid::Uuid::new_v4(), + frame: incoming, + request_ids: vec![RequestId::Number(7)], + response_tx: Some(response_tx), + }, + &mut channel_tx, + ) + .unwrap(); + assert!(matches!( + channel_rx.next().await, + Some(TransportFrame::Batch(_)) + )); + + let callback = TransportFrame::Single( + RawJsonRpcMessage::notification("callback".into(), serde_json::json!({})).unwrap(), + ); + assert!(server.try_dispatch_jsonrpc_message(callback).is_none()); + assert!(matches!( + response_rx.next().await, + Some(TransportFrame::Single(RawJsonRpcMessage::Notification(_))) + )); + + let frame = TransportFrame::Batch( + agent_client_protocol::TransportBatch::from_messages([ + RawJsonRpcMessage::response( + RequestId::Number(7), + Ok(serde_json::json!({ "ok": true })), + ), + ]) + .unwrap(), + ); + let expected = frame.to_json().unwrap(); + + assert!(server.try_dispatch_jsonrpc_message(frame).is_none()); + let received = response_rx + .next() + .await + .expect("waiting HTTP request stays open"); + assert_eq!(received.to_json().unwrap(), expected); + assert!(matches!(received, TransportFrame::Batch(_))); + }); + } + + #[test] + fn batch_post_round_trips_as_one_grouped_sse_response() { + futures::executor::block_on(async { + let (registration_tx, mut registration_rx) = mpsc::unbounded(); + let state = Arc::new(BridgeState { registration_tx }); + let incoming = serde_json::json!([ + { + "jsonrpc": "2.0", + "id": 7, + "method": "example/first", + "params": {} + }, + { + "jsonrpc": "2.0", + "id": 8, + "method": "example/second", + "params": {} + } + ]); + + let http_response = handle_post(State(state), incoming.to_string()) + .await + .expect("batch POST should open an SSE response"); + let registration = registration_rx + .next() + .await + .expect("batch POST should register with the bridge"); + + let mut server = RunningServer::new(); + let (mut channel_tx, mut channel_rx) = mpsc::unbounded(); + server + .handle_http_message(registration, &mut channel_tx) + .expect("batch POST should be forwarded to the channel"); + let forwarded = channel_rx + .next() + .await + .expect("channel should receive the batch frame"); + assert!(matches!(&forwarded, TransportFrame::Batch(_))); + assert_eq!( + serde_json::from_str::(&forwarded.to_json().unwrap()).unwrap(), + incoming + ); + + let response = TransportFrame::Batch( + agent_client_protocol::TransportBatch::from_messages([ + RawJsonRpcMessage::response( + RequestId::Number(7), + Ok(serde_json::json!({ "source": "first" })), + ), + RawJsonRpcMessage::response( + RequestId::Number(8), + Ok(serde_json::json!({ "source": "second" })), + ), + ]) + .expect("grouped response should be non-empty"), + ); + assert!(server.try_dispatch_jsonrpc_message(response).is_none()); + + let payload = single_sse_payload(http_response).await; + let entries = payload + .as_array() + .expect("SSE payload should remain one JSON-RPC array"); + assert_eq!(entries.len(), 2); + assert_eq!(entries[0]["id"], 7); + assert_eq!(entries[0]["result"]["source"], "first"); + assert_eq!(entries[1]["id"], 8); + assert_eq!(entries[1]["result"]["source"], "second"); + }); + } + + #[test] + fn malformed_batch_cannot_steal_a_valid_null_id_batch_response() { + futures::executor::block_on(async { + let (registration_tx, mut registration_rx) = mpsc::unbounded(); + let state = Arc::new(BridgeState { registration_tx }); + + let valid_http_response = handle_post( + State(state.clone()), + serde_json::json!([{ + "jsonrpc": "2.0", + "id": null, + "method": "example", + "params": {} + }]) + .to_string(), + ) + .await + .expect("valid null-ID batch POST"); + let malformed_http_response = handle_post(State(state), "[17,false]".to_owned()) + .await + .expect("malformed batch receives its own JSON-RPC error array"); + + let valid_batch = registration_rx + .next() + .await + .expect("valid batch is forwarded"); + assert!( + registration_rx.try_recv().is_err(), + "malformed-only batch must not register a bridge waiter" + ); + + let mut server = RunningServer::new(); + let (mut channel_tx, mut channel_rx) = mpsc::unbounded(); + server + .handle_http_message(valid_batch, &mut channel_tx) + .expect("forward valid null-ID batch"); + assert!(matches!( + channel_rx.next().await, + Some(TransportFrame::Batch(_)) + )); + + let valid_response = TransportFrame::Batch( + agent_client_protocol::TransportBatch::from_messages([ + RawJsonRpcMessage::response( + RequestId::Null, + Ok(serde_json::json!({ "source": "valid" })), + ), + ]) + .expect("valid response batch is non-empty"), + ); + assert!( + server + .try_dispatch_jsonrpc_message(valid_response) + .is_none() + ); + + let valid_payload = single_sse_payload(valid_http_response).await; + let valid_entries = valid_payload + .as_array() + .expect("valid response should remain a batch"); + assert_eq!(valid_entries.len(), 1); + assert_eq!(valid_entries[0]["id"], serde_json::Value::Null); + assert_eq!(valid_entries[0]["result"]["source"], "valid"); + + let malformed_payload = single_sse_payload(malformed_http_response).await; + let malformed_entries = malformed_payload + .as_array() + .expect("malformed response should be an error batch"); + assert_eq!(malformed_entries.len(), 2); + for entry in malformed_entries { + assert_eq!(entry["id"], serde_json::Value::Null); + assert_eq!( + entry["error"]["code"], + i32::from(agent_client_protocol::ErrorCode::InvalidRequest) + ); + } + }); + } +} diff --git a/src/agent-client-protocol/CHANGELOG.md b/src/agent-client-protocol/CHANGELOG.md index 860719f..36ab39b 100644 --- a/src/agent-client-protocol/CHANGELOG.md +++ b/src/agent-client-protocol/CHANGELOG.md @@ -2,8 +2,16 @@ ## [Unreleased] +### Changed + +- **Breaking:** Make `Channel` the batch-aware `TransportFrame` boundary and remove the hidden + `FramedChannel` compatibility path. See the + [2.0 migration guide](../../md/migration_v2.0.md). + ### Fixed +- Preserve JSON-RPC batch framing across component adapters, and suppress replies + to malformed response-shaped input. - Preserve `MatchDispatchFrom` retry state across chained matchers. ## [1.3.0](https://github.com/agentclientprotocol/rust-sdk/compare/v1.2.0...v1.3.0) - 2026-07-20 diff --git a/src/agent-client-protocol/src/acp_agent.rs b/src/agent-client-protocol/src/acp_agent.rs index 803c434..b12b9c3 100644 --- a/src/agent-client-protocol/src/acp_agent.rs +++ b/src/agent-client-protocol/src/acp_agent.rs @@ -729,13 +729,13 @@ impl crate::ConnectTo for Acp } } - fn into_framed_channel_and_future( + fn into_channel_and_future( self, ) -> ( - crate::FramedChannel, + crate::Channel, crate::BoxFuture<'static, Result<(), crate::Error>>, ) { - let (channel_for_caller, channel_for_agent) = crate::FramedChannel::duplex(); + let (channel_for_caller, channel_for_agent) = crate::Channel::duplex(); let future = Box::pin(crate::ConnectTo::::connect_to( self, channel_for_agent, @@ -1259,7 +1259,7 @@ mod tests { Ok(serde_json::json!({ "payload": "x".repeat(4 * 1024 * 1024) })), ); outgoing - .unbounded_send(Ok(response)) + .unbounded_send(crate::TransportFrame::Single(response)) .expect("response should be accepted before the connection starts"); outgoing.close_channel(); diff --git a/src/agent-client-protocol/src/component.rs b/src/agent-client-protocol/src/component.rs index 713fd31..f03785a 100644 --- a/src/agent-client-protocol/src/component.rs +++ b/src/agent-client-protocol/src/component.rs @@ -29,14 +29,10 @@ //! } //! ``` -use futures::{ - FutureExt as _, StreamExt as _, - channel::{mpsc, oneshot}, - future::{self, BoxFuture, Shared}, -}; +use futures::future::BoxFuture; use std::{fmt::Debug, future::Future, marker::PhantomData}; -use crate::{Channel, FramedChannel, Result, role::Role}; +use crate::{Channel, Result, role::Role}; /// A component that can exchange JSON-RPC messages to an endpoint playing the role `R` /// (e.g., an ACP [`Agent`](`crate::role::acp::Agent`) or an MCP [`Server`](`crate::role::mcp::Server`)). @@ -137,6 +133,10 @@ pub trait ConnectTo: Send + 'static { /// Convert this component into a channel endpoint and server future. /// + /// The returned [`Channel`] is the canonical frame-aware boundary. It carries + /// complete [`TransportFrame`](crate::TransportFrame) values so default + /// adapters preserve batch grouping. + /// /// This method returns: /// - A `Channel` that can be used to communicate with this component /// - A `BoxFuture` that runs the component's server logic @@ -158,82 +158,6 @@ pub trait ConnectTo: Send + 'static { let future = Box::pin(self.connect_to(channel_b)); (channel_a, future) } - - /// Convert this component into the SDK's batch-aware transport channel. - /// - /// This is an internal extension point used by built-in transports to retain - /// JSON-RPC batch boundaries. Implementations normally do not need to - /// override it. - #[doc(hidden)] - fn into_framed_channel_and_future(self) -> (FramedChannel, BoxFuture<'static, Result<()>>) - where - Self: Sized, - { - let (channel_for_caller, channel_for_component) = FramedChannel::duplex(); - let (bridge_tx, mut bridge_rx) = mpsc::unbounded(); - let (component_done_tx, component_done_rx) = oneshot::channel(); - let component_done = component_done_rx.map(|_| ()).boxed().shared(); - let component_channel = DefaultFramedChannel { - channel: channel_for_component, - bridge_tx, - component_done, - }; - - let future = Box::pin(async move { - let component = async move { - let result = self.connect_to(component_channel).await; - let _ = component_done_tx.send(()); - result - }; - let bridges = async move { - while let Some(bridge) = bridge_rx.next().await { - bridge.await?; - } - Ok::<(), crate::Error>(()) - }; - - futures::try_join!(component, bridges)?; - Ok(()) - }); - (channel_for_caller, future) - } -} - -/// A negotiating endpoint for the default component adapter. -/// -/// Transparent components receive its framed channel directly, while legacy -/// components can still request a [`Channel`]. Legacy bridge work is driven by -/// the outer component future so their transport future retains the immediate -/// completion behavior of an in-process `Channel`. -struct DefaultFramedChannel { - channel: FramedChannel, - bridge_tx: mpsc::UnboundedSender>>, - component_done: Shared>, -} - -impl ConnectTo for DefaultFramedChannel { - async fn connect_to(self, client: impl ConnectTo) -> Result<()> { - let Self { channel, .. } = self; - ConnectTo::::connect_to(channel, client).await - } - - fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<()>>) { - let Self { - channel, - bridge_tx, - component_done, - } = self; - let (channel, bridge) = channel.into_legacy_channel_until(component_done); - let registered = bridge_tx - .unbounded_send(bridge) - .map_err(crate::Error::into_internal_error); - (channel, Box::pin(future::ready(registered))) - } - - fn into_framed_channel_and_future(self) -> (FramedChannel, BoxFuture<'static, Result<()>>) { - let Self { channel, .. } = self; - ConnectTo::::into_framed_channel_and_future(channel) - } } /// Type-erased connect trait for object-safe dynamic dispatch. @@ -251,10 +175,6 @@ trait ErasedConnectTo: Send { fn into_channel_and_future_erased(self: Box) -> (Channel, BoxFuture<'static, Result<()>>); - - fn into_framed_channel_and_future_erased( - self: Box, - ) -> (FramedChannel, BoxFuture<'static, Result<()>>); } /// Blanket implementation: any `Serve` can be type-erased. @@ -282,12 +202,6 @@ impl, R: Role> ErasedConnectTo for C { ) -> (Channel, BoxFuture<'static, Result<()>>) { (*self).into_channel_and_future() } - - fn into_framed_channel_and_future_erased( - self: Box, - ) -> (FramedChannel, BoxFuture<'static, Result<()>>) { - (*self).into_framed_channel_and_future() - } } /// A dynamically-typed component for heterogeneous collections. @@ -340,10 +254,6 @@ impl ConnectTo for DynConnectTo { fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<()>>) { self.inner.into_channel_and_future_erased() } - - fn into_framed_channel_and_future(self) -> (FramedChannel, BoxFuture<'static, Result<()>>) { - self.inner.into_framed_channel_and_future_erased() - } } impl Debug for DynConnectTo { diff --git a/src/agent-client-protocol/src/jsonrpc.rs b/src/agent-client-protocol/src/jsonrpc.rs index 28c16af..7c9bd56 100644 --- a/src/agent-client-protocol/src/jsonrpc.rs +++ b/src/agent-client-protocol/src/jsonrpc.rs @@ -64,64 +64,79 @@ pub enum RawJsonRpcMessage { Response(RpcResponse), } -/// A batch-aware packet exchanged between the SDK's protocol and transport actors. +/// A JSON-RPC frame exchanged between protocol components and transports. /// -/// This is an internal extension point. Use [`Channel`] when exchanging individual -/// raw JSON-RPC messages with application code. -#[doc(hidden)] +/// A frame preserves the boundary between a single JSON-RPC value and a batch. +/// Malformed wire input is represented explicitly; transport failures are +/// reported by the future that drives the transport rather than sent through a +/// [`Channel`]. #[derive(Debug)] -pub(crate) enum TransportFrame { - /// One JSON-RPC message or an opaque component error. - Single(Result), - /// One malformed or invalid wire value retained for framed relays. - InvalidSingle { raw: String, error: crate::Error }, +pub enum TransportFrame { + /// One valid JSON-RPC message. + Single(RawJsonRpcMessage), + /// One malformed or invalid wire value retained for relays. + Malformed { + /// The original wire representation. + raw: String, + /// The JSON-RPC error associated with the malformed value. + error: crate::Error, + }, /// Entries retained from one non-empty JSON-RPC batch, kept in source order. Batch(TransportBatch), } /// A structurally non-empty JSON-RPC batch retained across framed relays. #[derive(Debug)] -pub(crate) struct TransportBatch { +pub struct TransportBatch { first: TransportBatchEntry, rest: Vec, } +/// One entry in a [`TransportBatch`]. #[derive(Debug)] -enum TransportBatchEntry { +pub enum TransportBatchEntry { + /// A valid JSON-RPC message. Message(RawJsonRpcMessage), - Invalid { + /// A malformed or invalid JSON-RPC value retained for relays. + Malformed { + /// The original JSON value. raw: serde_json::Value, + /// The JSON-RPC error associated with the malformed value. error: crate::Error, }, } impl TransportBatchEntry { - fn message(message: RawJsonRpcMessage) -> Self { + /// Create a valid batch entry. + #[must_use] + pub fn message(message: RawJsonRpcMessage) -> Self { Self::Message(message) } - fn invalid(raw: serde_json::Value, error: crate::Error) -> Self { - Self::Invalid { raw, error } + /// Create a malformed batch entry. + #[must_use] + pub fn malformed(raw: serde_json::Value, error: crate::Error) -> Self { + Self::Malformed { raw, error } } fn as_result(&self) -> Result<&RawJsonRpcMessage, &crate::Error> { match self { Self::Message(message) => Ok(message), - Self::Invalid { error, .. } => Err(error), + Self::Malformed { error, .. } => Err(error), } } fn into_result(self) -> Result { match self { Self::Message(message) => Ok(message), - Self::Invalid { error, .. } => Err(error), + Self::Malformed { error, .. } => Err(error), } } fn message_ref(&self) -> Option<&RawJsonRpcMessage> { match self { Self::Message(message) => Some(message), - Self::Invalid { .. } => None, + Self::Malformed { .. } => None, } } } @@ -133,13 +148,16 @@ impl Serialize for TransportBatchEntry { { match self { Self::Message(message) => message.serialize(serializer), - Self::Invalid { raw, .. } => raw.serialize(serializer), + Self::Malformed { raw, .. } => raw.serialize(serializer), } } } impl TransportBatch { - fn from_entries(entries: Vec) -> Option { + /// Create a non-empty batch from entries. + /// + /// Returns `None` when the iterator is empty. + pub fn from_entries(entries: impl IntoIterator) -> Option { let mut entries = entries.into_iter(); Some(Self { first: entries.next()?, @@ -147,21 +165,38 @@ impl TransportBatch { }) } - pub(crate) fn from_messages( - messages: impl IntoIterator, - ) -> Option { - Self::from_entries( - messages - .into_iter() - .map(TransportBatchEntry::message) - .collect(), - ) + /// Create a non-empty batch from valid messages. + /// + /// Returns `None` when the iterator is empty. + pub fn from_messages(messages: impl IntoIterator) -> Option { + Self::from_entries(messages.into_iter().map(TransportBatchEntry::message)) } - fn entries(&self) -> impl Iterator { + /// Iterate over entries in source order. + pub fn entries(&self) -> impl Iterator { std::iter::once(&self.first).chain(&self.rest) } + /// Iterate mutably over entries in source order. + pub fn entries_mut(&mut self) -> impl Iterator { + std::iter::once(&mut self.first).chain(&mut self.rest) + } + + /// Return the number of entries in this non-empty batch. + #[must_use] + pub fn len(&self) -> usize { + 1 + self.rest.len() + } + + /// Return whether this batch is empty. + /// + /// A `TransportBatch` is structurally non-empty, so this always returns + /// `false`. + #[must_use] + pub const fn is_empty(&self) -> bool { + false + } + #[cfg(any(feature = "unstable_protocol_v2", test))] pub(crate) fn iter_results( &self, @@ -181,18 +216,13 @@ impl TransportBatch { pub(crate) fn first_result_mut(&mut self) -> Result<&mut RawJsonRpcMessage, crate::Error> { match &mut self.first { TransportBatchEntry::Message(message) => Ok(message), - TransportBatchEntry::Invalid { error, .. } => Err(error.clone()), + TransportBatchEntry::Malformed { error, .. } => Err(error.clone()), } } fn messages(&self) -> impl Iterator { self.entries().filter_map(TransportBatchEntry::message_ref) } - - #[cfg(test)] - fn len(&self) -> usize { - 1 + self.rest.len() - } } impl Serialize for TransportBatch { @@ -210,21 +240,13 @@ impl Serialize for TransportBatch { } impl TransportFrame { - fn into_messages(self) -> impl Iterator> { - match self { - Self::Single(message) => vec![message].into_iter(), - Self::InvalidSingle { error, .. } => vec![Err(error)].into_iter(), - Self::Batch(batch) => batch.into_results().collect::>().into_iter(), - } - } - fn inspect_messages( &self, observer: &mut impl FnMut(&RawJsonRpcMessage) -> Result<(), crate::Error>, ) -> Result<(), crate::Error> { match self { - Self::Single(Ok(message)) => observer(message), - Self::Single(Err(_)) | Self::InvalidSingle { .. } => Ok(()), + Self::Single(message) => observer(message), + Self::Malformed { .. } => Ok(()), Self::Batch(batch) => { for message in batch.messages() { observer(message)?; @@ -1682,8 +1704,7 @@ impl< // Convert transport into server - this returns a channel for us to use // and a future that runs the transport. let transport_component = crate::DynConnectTo::new(transport); - let (transport_channel, transport_future) = - transport_component.into_framed_channel_and_future(); + let (transport_channel, transport_future) = transport_component.into_channel_and_future(); let (transport_completion_tx, transport_completion_rx) = oneshot::channel(); let transport_completion = transport_completion_rx .map(|result| { @@ -1711,7 +1732,7 @@ impl< }); // Destructure the channel endpoints - let FramedChannel { + let Channel { rx: transport_incoming_rx, tx: transport_outgoing_tx, } = transport_channel; @@ -1784,10 +1805,8 @@ where Builder::connect_to(self, client).await } - fn into_framed_channel_and_future( - self, - ) -> (FramedChannel, BoxFuture<'static, Result<(), crate::Error>>) { - let (channel_for_caller, channel_for_builder) = FramedChannel::duplex(); + fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { + let (channel_for_caller, channel_for_builder) = Channel::duplex(); let future = Box::pin(Builder::connect_to(self, channel_for_builder)); (channel_for_caller, future) } @@ -2393,7 +2412,7 @@ impl ResponseDestination { fn complete(self, response: RawJsonRpcMessage) -> Option { match self { - Self::Individual => Some(TransportFrame::Single(Ok(response))), + Self::Individual => Some(TransportFrame::Single(response)), Self::Batch(slot) => slot.complete(response).map(batch_response_frame), } } @@ -5082,14 +5101,12 @@ where Self { outgoing, incoming } } - fn into_framed_transport( - self, - ) -> (FramedChannel, BoxFuture<'static, Result<(), crate::Error>>) { + fn into_channel_transport(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { let Self { outgoing, incoming } = self; - let (channel_for_caller, channel_for_lines) = FramedChannel::duplex(); + let (channel_for_caller, channel_for_lines) = Channel::duplex(); let server_future = Box::pin(async move { - let FramedChannel { rx, tx } = channel_for_lines; + let Channel { rx, tx } = channel_for_lines; let outgoing_future = transport_actor::transport_outgoing_lines_actor(rx, outgoing); let incoming_future = transport_actor::transport_incoming_lines_actor(incoming, tx); futures::try_join!(outgoing_future, incoming_future)?; @@ -5107,7 +5124,13 @@ where { async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { let Self { outgoing, incoming } = self; - let (FramedChannel { rx, tx }, client_future) = client.into_framed_channel_and_future(); + let (Channel { rx, tx }, client_channel) = Channel::duplex(); + let close_client_output = client_channel.tx.clone(); + let client_future = Box::pin(async move { + let result = client.connect_to(client_channel).await; + close_client_output.close_channel(); + result + }); // Once the client completes successfully, its incoming channel is // gone. Keep consuming successful messages from the physical read @@ -5160,25 +5183,7 @@ where } fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { - let Self { outgoing, incoming } = self; - let (channel_for_caller, channel_for_lines) = Channel::duplex(); - let Channel { rx, tx } = channel_for_lines; - - let future = Box::pin(async move { - futures::try_join!( - transport_actor::transport_outgoing_legacy_lines_actor(rx, outgoing), - transport_actor::transport_incoming_legacy_lines_actor(incoming, tx), - )?; - Ok(()) - }); - - (channel_for_caller, future) - } - - fn into_framed_channel_and_future( - self, - ) -> (FramedChannel, BoxFuture<'static, Result<(), crate::Error>>) { - self.into_framed_transport() + self.into_channel_transport() } } @@ -5283,28 +5288,42 @@ where fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { ConnectTo::::into_channel_and_future(self.into_lines()) } - - fn into_framed_channel_and_future( - self, - ) -> (FramedChannel, BoxFuture<'static, Result<(), crate::Error>>) { - ConnectTo::::into_framed_channel_and_future(self.into_lines()) - } } -/// A batch-aware channel used internally between protocol and transport components. +/// A channel endpoint representing one side of a bidirectional JSON-RPC transport. /// -/// This internal extension point preserves whether messages arrived in one JSON-RPC -/// batch. Application code should normally use [`Channel`]. -#[doc(hidden)] +/// A channel carries complete TransportFrame values, preserving batch boundaries +/// across in-process components and transport adapters. Malformed wire input is an +/// explicit frame; failures while driving a physical transport are returned by that +/// transport's future. +/// +/// # Example +/// +/// ```no_run +/// # use agent_client_protocol::UntypedRole; +/// # use agent_client_protocol::Channel; +/// # async fn example() -> Result<(), agent_client_protocol::Error> { +/// let (channel_a, _channel_b) = Channel::duplex(); +/// +/// UntypedRole.builder() +/// .name("connection-a") +/// .connect_to(channel_a) +/// .await?; +/// # Ok(()) +/// # } +/// ``` #[derive(Debug)] -pub struct FramedChannel { - pub(crate) rx: mpsc::UnboundedReceiver, - pub(crate) tx: mpsc::UnboundedSender, +pub struct Channel { + /// Receives frames from the counterpart. + pub rx: mpsc::UnboundedReceiver, + /// Sends frames to the counterpart. + pub tx: mpsc::UnboundedSender, } -impl FramedChannel { - /// Create a connected pair of batch-aware channel endpoints. - #[doc(hidden)] +impl Channel { + /// Create a pair of connected channel endpoints. + /// + /// Frames sent through either endpoint are received by the other endpoint. #[must_use] pub fn duplex() -> (Self, Self) { let (a_tx, b_rx) = mpsc::unbounded(); @@ -5313,124 +5332,12 @@ impl FramedChannel { (Self { rx: a_rx, tx: a_tx }, Self { rx: b_rx, tx: b_tx }) } - pub(crate) fn from_legacy_channel( - channel: Channel, - component_future: Option>>, - ) -> (Self, BoxFuture<'static, Result<(), crate::Error>>) { - let (channel_for_caller, channel_for_bridge) = Self::duplex(); - let Channel { - rx: legacy_rx, - tx: legacy_tx, - } = channel; - let FramedChannel { - rx: framed_rx, - tx: framed_tx, - } = channel_for_bridge; - let future: BoxFuture<'static, Result<(), crate::Error>> = - if let Some(component_future) = component_future { - Box::pin(async move { - let (component_done_tx, component_done_rx) = oneshot::channel(); - let component_done = component_done_rx.map(|_| ()).boxed().shared(); - let component = async move { - let result = component_future.await; - let _ = component_done_tx.send(()); - result - }; - let component_to_caller = { - let component_done = component_done.clone(); - async move { - let mut legacy_rx = legacy_rx; - loop { - match future::select( - Box::pin(legacy_rx.next()), - Box::pin(component_done.clone()), - ) - .await - { - Either::Left((Some(message), _)) => { - if framed_tx - .unbounded_send(TransportFrame::Single(message)) - .is_err() - { - return Ok(()); - } - } - Either::Left((None, _)) => return Ok(()), - Either::Right(((), _)) => { - legacy_rx.close(); - while let Some(message) = legacy_rx.next().await { - if framed_tx - .unbounded_send(TransportFrame::Single(message)) - .is_err() - { - break; - } - } - return Ok(()); - } - } - } - } - }; - let caller_to_component = async move { - let mut framed_rx = framed_rx; - loop { - match future::select( - Box::pin(framed_rx.next()), - Box::pin(component_done.clone()), - ) - .await - { - Either::Left((Some(frame), _)) => { - for message in frame.into_messages() { - if legacy_tx.unbounded_send(message).is_err() { - return Ok(()); - } - } - } - Either::Left((None, _)) | Either::Right(((), _)) => return Ok(()), - } - } - }; - - futures::try_join!(component, component_to_caller, caller_to_component)?; - Ok::<(), crate::Error>(()) - }) - } else { - Box::pin(async move { - let component_to_caller = async move { - let mut legacy_rx = legacy_rx; - while let Some(message) = legacy_rx.next().await { - if framed_tx - .unbounded_send(TransportFrame::Single(message)) - .is_err() - { - break; - } - } - Ok::<(), crate::Error>(()) - }; - let caller_to_component = async move { - let mut framed_rx = framed_rx; - while let Some(frame) = framed_rx.next().await { - for message in frame.into_messages() { - if legacy_tx.unbounded_send(message).is_err() { - return Ok::<(), crate::Error>(()); - } - } - } - Ok::<(), crate::Error>(()) - }; - - futures::try_join!(component_to_caller, caller_to_component)?; - Ok::<(), crate::Error>(()) - }) - }; - - (channel_for_caller, future) - } - - pub(crate) async fn copy(mut self) -> Result<(), crate::Error> { + /// Copy frames from `rx` to `tx` until the input closes. + /// + /// # Errors + /// + /// Returns an error if the receiving endpoint closes before the input. + pub async fn copy(mut self) -> Result<(), crate::Error> { while let Some(frame) = self.rx.next().await { self.tx .unbounded_send(frame) @@ -5439,17 +5346,15 @@ impl FramedChannel { Ok(()) } - /// Bridge two internal framed endpoints while inspecting valid messages. + /// Bridge two endpoints while inspecting every valid message. /// - /// Observers are invoked for each message in source order, including each - /// valid member of a batch. The original frame is then forwarded unchanged, - /// preserving batch boundaries across instrumentation layers. + /// Observers are invoked in source order, including for each valid member of + /// a batch. The original frame is forwarded unchanged after inspection. /// /// # Errors /// - /// Returns an error from either observer or when a destination closes - /// before its source. - #[doc(hidden)] + /// Returns an observer error or an error if a destination closes before its + /// source. pub async fn bridge_with_inspection( left: Self, right: Self, @@ -5487,155 +5392,19 @@ impl FramedChannel { futures::try_join!(left_to_right, right_to_left)?; Ok(()) } - - pub(crate) fn into_legacy_channel_until( - self, - component_done: futures::future::Shared>, - ) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { - let (channel_for_component, channel_for_bridge) = Channel::duplex(); - let FramedChannel { - rx: framed_rx, - tx: framed_tx, - } = self; - let Channel { - rx: legacy_rx, - tx: legacy_tx, - } = channel_for_bridge; - let future = Box::pin(async move { - let component_to_transport = { - let component_done = component_done.clone(); - async move { - let mut legacy_rx = legacy_rx; - loop { - match future::select( - Box::pin(legacy_rx.next()), - Box::pin(component_done.clone()), - ) - .await - { - Either::Left((Some(message), _)) => { - if framed_tx - .unbounded_send(TransportFrame::Single(message)) - .is_err() - { - return Ok::<(), crate::Error>(()); - } - } - Either::Left((None, _)) => return Ok::<(), crate::Error>(()), - Either::Right(((), _)) => { - legacy_rx.close(); - while let Some(message) = legacy_rx.next().await { - if framed_tx - .unbounded_send(TransportFrame::Single(message)) - .is_err() - { - break; - } - } - return Ok::<(), crate::Error>(()); - } - } - } - } - }; - let transport_to_component = async move { - let mut framed_rx = framed_rx; - loop { - match future::select( - Box::pin(framed_rx.next()), - Box::pin(component_done.clone()), - ) - .await - { - Either::Left((Some(frame), _)) => { - for message in frame.into_messages() { - if legacy_tx.unbounded_send(message).is_err() { - return Ok::<(), crate::Error>(()); - } - } - } - Either::Left((None, _)) | Either::Right(((), _)) => { - return Ok::<(), crate::Error>(()); - } - } - } - }; - - futures::try_join!(component_to_transport, transport_to_component)?; - Ok(()) - }); - - (channel_for_component, future) - } - - fn into_legacy_channel(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { - let (channel_for_caller, channel_for_bridge) = Channel::duplex(); - let FramedChannel { - rx: framed_rx, - tx: framed_tx, - } = self; - let Channel { - rx: legacy_rx, - tx: legacy_tx, - } = channel_for_bridge; - let future = Box::pin(async move { - let legacy_to_framed = async move { - let mut legacy_rx = legacy_rx; - while let Some(message) = legacy_rx.next().await { - if framed_tx - .unbounded_send(TransportFrame::Single(message)) - .is_err() - { - break; - } - } - Ok::<(), crate::Error>(()) - }; - let framed_to_legacy = async move { - let mut framed_rx = framed_rx; - while let Some(frame) = framed_rx.next().await { - match frame { - TransportFrame::Single(message) => { - if legacy_tx.unbounded_send(message).is_err() { - return Ok::<(), crate::Error>(()); - } - } - TransportFrame::InvalidSingle { error, .. } => { - if legacy_tx.unbounded_send(Err(error)).is_err() { - return Ok::<(), crate::Error>(()); - } - } - TransportFrame::Batch(batch) => { - for entry in batch.into_results() { - if legacy_tx.unbounded_send(entry).is_err() { - return Ok::<(), crate::Error>(()); - } - } - } - } - } - Ok::<(), crate::Error>(()) - }; - - futures::try_join!(legacy_to_framed, framed_to_legacy)?; - Ok::<(), crate::Error>(()) - }); - - (channel_for_caller, future) - } } -impl ConnectTo for FramedChannel { +impl ConnectTo for Channel { async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { - let (client_channel, client_future) = client.into_framed_channel_and_future(); + let (client_channel, client_future) = client.into_channel_and_future(); let ((), (), ()) = futures::try_join!( - FramedChannel { + Channel { rx: client_channel.rx, tx: self.tx, } .copy(), - FramedChannel { + Channel { rx: self.rx, tx: client_channel.tx, } @@ -5645,133 +5414,9 @@ impl ConnectTo for FramedChannel { Ok(()) } - fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { - self.into_legacy_channel() - } - - fn into_framed_channel_and_future( - self, - ) -> (FramedChannel, BoxFuture<'static, Result<(), crate::Error>>) { - let (channel_for_caller, channel_for_bridge) = FramedChannel::duplex(); - let future = Box::pin(async move { - let ((), ()) = futures::try_join!( - FramedChannel { - rx: channel_for_bridge.rx, - tx: self.tx, - } - .copy(), - FramedChannel { - rx: self.rx, - tx: channel_for_bridge.tx, - } - .copy(), - )?; - Ok(()) - }); - (channel_for_caller, future) - } -} - -/// A channel endpoint representing one side of a bidirectional message channel. -/// -/// `Channel` represents a single endpoint's view of a bidirectional communication channel. -/// Each endpoint has: -/// - `rx`: A receiver for incoming messages (or errors) from the counterpart -/// - `tx`: A sender for outgoing messages (or errors) to the counterpart -/// -/// # Example -/// -/// ```no_run -/// # use agent_client_protocol::UntypedRole; -/// # use agent_client_protocol::{Channel, Builder}; -/// # async fn example() -> Result<(), agent_client_protocol::Error> { -/// // Create a pair of connected channels -/// let (channel_a, channel_b) = Channel::duplex(); -/// -/// // Each channel can be used by a different component -/// UntypedRole.builder() -/// .name("connection-a") -/// .connect_to(channel_a) -/// .await?; -/// # Ok(()) -/// # } -/// ``` -#[derive(Debug)] -pub struct Channel { - /// Receives messages (or errors) from the counterpart. - pub rx: mpsc::UnboundedReceiver>, - /// Sends messages (or errors) to the counterpart. - pub tx: mpsc::UnboundedSender>, -} - -impl Channel { - /// Create a pair of connected channel endpoints. - /// - /// Returns two `Channel` instances that are connected to each other: - /// - Messages sent via `channel_a.tx` are received on `channel_b.rx` - /// - Messages sent via `channel_b.tx` are received on `channel_a.rx` - /// - /// # Returns - /// - /// A tuple `(channel_a, channel_b)` of connected channel endpoints. - #[must_use] - pub fn duplex() -> (Self, Self) { - // Create channels: A sends Result which B receives as Message - let (a_tx, b_rx) = mpsc::unbounded(); - let (b_tx, a_rx) = mpsc::unbounded(); - - let channel_a = Self { rx: a_rx, tx: a_tx }; - let channel_b = Self { rx: b_rx, tx: b_tx }; - - (channel_a, channel_b) - } - - /// Copy messages from `rx` to `tx`. - /// - /// # Returns - /// - /// A `Result` indicating success or failure. - pub async fn copy(mut self) -> Result<(), crate::Error> { - while let Some(msg) = self.rx.next().await { - self.tx - .unbounded_send(msg) - .map_err(crate::util::internal_error)?; - } - Ok(()) - } -} - -impl ConnectTo for Channel { - async fn connect_to(self, client: impl ConnectTo) -> Result<(), crate::Error> { - let (client_channel, client_serve) = client.into_channel_and_future(); - - match futures::try_join!( - Channel { - rx: client_channel.rx, - tx: self.tx - } - .copy(), - Channel { - rx: self.rx, - tx: client_channel.tx - } - .copy(), - client_serve - ) { - Ok(((), (), ())) => Ok(()), - Err(err) => Err(err), - } - } - fn into_channel_and_future(self) -> (Channel, BoxFuture<'static, Result<(), crate::Error>>) { (self, Box::pin(future::ready(Ok(())))) } - - fn into_framed_channel_and_future( - self, - ) -> (FramedChannel, BoxFuture<'static, Result<(), crate::Error>>) { - FramedChannel::from_legacy_channel(self, None) - } } #[cfg(test)] diff --git a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs index ad9d6fb..0b85989 100644 --- a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs @@ -293,11 +293,11 @@ fn frame_entries( ) { let entries: Vec> = match frame { TransportFrame::Single(message) => { - let destination = - message_requires_response(&message).then_some(ResponseDestination::Individual); - return (vec![(message, destination)], None); + let destination = matches!(&message, RawJsonRpcMessage::Request(_)) + .then_some(ResponseDestination::Individual); + return (vec![(Ok(message), destination)], None); } - TransportFrame::InvalidSingle { error, .. } => { + TransportFrame::Malformed { error, .. } => { return ( vec![(Err(error), Some(ResponseDestination::Individual))], None, diff --git a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs index 0a7eb89..2203352 100644 --- a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs @@ -82,8 +82,7 @@ pub(super) async fn outgoing_protocol_actor( continue; } - if let Err(error) = transport_tx.unbounded_send(TransportFrame::Single(Ok(request))) - { + if let Err(error) = transport_tx.unbounded_send(TransportFrame::Single(request)) { let error = crate::Error::into_internal_error(error); if let Some(pending_reply) = pending_replies.remove(&id) { pending_reply.fail(error.clone()); @@ -116,7 +115,7 @@ pub(super) async fn outgoing_protocol_actor( } }; transport_tx - .unbounded_send(TransportFrame::Single(Ok(message))) + .unbounded_send(TransportFrame::Single(message)) .map_err(crate::Error::into_internal_error)?; } continue; diff --git a/src/agent-client-protocol/src/jsonrpc/transport_actor.rs b/src/agent-client-protocol/src/jsonrpc/transport_actor.rs index ae8f275..92b33b3 100644 --- a/src/agent-client-protocol/src/jsonrpc/transport_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/transport_actor.rs @@ -9,8 +9,8 @@ use serde::Deserialize as _; enum ParsedIncomingLine { Ignored, - Single(Result), - InvalidSingle { raw: String, error: crate::Error }, + Single(RawJsonRpcMessage), + Malformed { raw: String, error: crate::Error }, Batch(TransportBatch), } @@ -28,7 +28,7 @@ fn parse_incoming_line(line: &str) -> ParsedIncomingLine { Ok(value) => value, Err(error) => { tracing::debug!(?error, "Failed to parse incoming JSON-RPC JSON"); - return ParsedIncomingLine::InvalidSingle { + return ParsedIncomingLine::Malformed { raw: line.to_owned(), error: crate::Error::parse_error().data(serde_json::json!({ "line": line })), }; @@ -36,12 +36,10 @@ fn parse_incoming_line(line: &str) -> ParsedIncomingLine { }; match value { - serde_json::Value::Array(entries) if entries.is_empty() => { - ParsedIncomingLine::InvalidSingle { - raw: line.to_owned(), - error: crate::Error::invalid_request(), - } - } + serde_json::Value::Array(entries) if entries.is_empty() => ParsedIncomingLine::Malformed { + raw: line.to_owned(), + error: crate::Error::invalid_request(), + }, serde_json::Value::Array(entries) => { let mut has_response_entry = false; let mut has_call_entry = false; @@ -68,7 +66,10 @@ fn parse_incoming_line(line: &str) -> ParsedIncomingLine { // mixed batch. Keep ambiguous call-shaped entries // so invalid requests still receive an error. (!looks_like_response || looks_like_call).then(|| { - TransportBatchEntry::invalid(entry, crate::Error::invalid_request()) + TransportBatchEntry::malformed( + entry, + crate::Error::invalid_request(), + ) }) } } @@ -89,13 +90,13 @@ fn parse_incoming_line(line: &str) -> ParsedIncomingLine { value => { let (looks_like_call, looks_like_response) = message_shape(&value); match serde_json::from_value(value) { - Ok(message) => ParsedIncomingLine::Single(Ok(message)), + Ok(message) => ParsedIncomingLine::Single(message), Err(error) => { tracing::debug!(?error, "Invalid JSON-RPC message"); if looks_like_response && !looks_like_call { ParsedIncomingLine::Ignored } else { - ParsedIncomingLine::InvalidSingle { + ParsedIncomingLine::Malformed { raw: line.to_owned(), error: crate::Error::invalid_request(), } @@ -106,6 +107,41 @@ fn parse_incoming_line(line: &str) -> ParsedIncomingLine { } } +impl TransportFrame { + /// Parse one JSON-RPC wire value while preserving batch boundaries. + /// + /// Malformed calls are returned as explicit malformed frames or batch + /// entries. Malformed response-shaped input is ignored, yielding `None`, + /// because JSON-RPC responses must not themselves receive responses. + #[must_use] + pub fn parse_json(input: &str) -> Option { + match parse_incoming_line(input) { + ParsedIncomingLine::Ignored => None, + ParsedIncomingLine::Single(message) => Some(Self::Single(message)), + ParsedIncomingLine::Malformed { raw, error } => Some(Self::Malformed { raw, error }), + ParsedIncomingLine::Batch(batch) => Some(Self::Batch(batch)), + } + } + + /// Serialize this frame to its JSON-RPC wire representation. + /// + /// # Errors + /// + /// Returns an internal error if a valid message or batch cannot be + /// serialized. Malformed frames return their original wire text unchanged. + pub fn to_json(&self) -> Result { + match self { + Self::Single(message) => { + serde_json::to_string(message).map_err(crate::Error::into_internal_error) + } + Self::Malformed { raw, .. } => Ok(raw.clone()), + Self::Batch(batch) => { + serde_json::to_string(batch).map_err(crate::Error::into_internal_error) + } + } + } +} + /// Transport outgoing actor for line streams: Serializes RawJsonRpcMessage and yields lines. /// /// This is a line-based variant of `transport_outgoing_actor` that works with a Sink @@ -113,7 +149,6 @@ fn parse_incoming_line(line: &str) -> ParsedIncomingLine { /// written to the underlying transport. /// /// This actor handles transport mechanics: -/// - Unwraps Result from the channel /// - Serializes RawJsonRpcMessage to JSON strings /// - Yields newline-terminated strings /// - Handles serialization errors @@ -128,10 +163,11 @@ async fn transport_outgoing_frames_actor( let mut outgoing_lines = pin!(outgoing_lines); while let Some(frame) = transport_rx.next().await { - let message_result = match frame { + let json_rpc_message = match frame { TransportFrame::Single(message) => message, - TransportFrame::InvalidSingle { raw, .. } => { - tracing::trace!(message = %raw, "Relaying invalid JSON-RPC value"); + TransportFrame::Malformed { raw, .. } => { + let raw = malformed_line_value(raw)?; + tracing::trace!(message = ?raw, "Relaying invalid JSON-RPC value"); outgoing_lines .send(raw) .await @@ -149,7 +185,6 @@ async fn transport_outgoing_frames_actor( continue; } }; - let json_rpc_message = message_result?; match serde_json::to_string(&json_rpc_message) { Ok(line) => { tracing::trace!(message = %line, "Sending JSON-RPC message"); @@ -199,6 +234,18 @@ async fn transport_outgoing_frames_actor( Ok(()) } +fn malformed_line_value(raw: String) -> Result { + if !raw.contains('\r') && !raw.contains('\n') { + return Ok(raw); + } + + match serde_json::from_str::(&raw) { + Ok(value) => serde_json::to_string(&value), + Err(_) => serde_json::to_string(&raw), + } + .map_err(crate::Error::into_internal_error) +} + pub(super) async fn transport_outgoing_lines_actor( transport_rx: mpsc::UnboundedReceiver, outgoing_lines: impl futures::Sink, @@ -206,13 +253,6 @@ pub(super) async fn transport_outgoing_lines_actor( transport_outgoing_frames_actor(transport_rx, outgoing_lines).await } -pub(super) async fn transport_outgoing_legacy_lines_actor( - transport_rx: mpsc::UnboundedReceiver>, - outgoing_lines: impl futures::Sink, -) -> Result<(), crate::Error> { - transport_outgoing_frames_actor(transport_rx.map(TransportFrame::Single), outgoing_lines).await -} - /// Transport incoming actor for line streams: Parses lines into RawJsonRpcMessage values. /// /// This is a line-based variant of `transport_incoming_actor` that works with a @@ -241,9 +281,9 @@ pub(super) async fn transport_incoming_lines_actor( .unbounded_send(TransportFrame::Single(message)) .map_err(crate::Error::into_internal_error)?; } - ParsedIncomingLine::InvalidSingle { raw, error } => { + ParsedIncomingLine::Malformed { raw, error } => { transport_tx - .unbounded_send(TransportFrame::InvalidSingle { raw, error }) + .unbounded_send(TransportFrame::Malformed { raw, error }) .map_err(crate::Error::into_internal_error)?; } ParsedIncomingLine::Batch(entries) => { @@ -256,32 +296,10 @@ pub(super) async fn transport_incoming_lines_actor( Ok(()) } -pub(super) async fn transport_incoming_legacy_lines_actor( - incoming_lines: impl futures::Stream>, - transport_tx: mpsc::UnboundedSender>, -) -> Result<(), crate::Error> { - let mut incoming_lines = pin!(incoming_lines); - while let Some(line_result) = incoming_lines.next().await { - let line = line_result.map_err(crate::Error::into_internal_error)?; - tracing::trace!(message = %line, "Received JSON-RPC message"); - - let entries = match parse_incoming_line(&line) { - ParsedIncomingLine::Ignored => Vec::new(), - ParsedIncomingLine::Single(message) => vec![message], - ParsedIncomingLine::InvalidSingle { error, .. } => vec![Err(error)], - ParsedIncomingLine::Batch(batch) => batch.into_results().collect(), - }; - for entry in entries { - transport_tx - .unbounded_send(entry) - .map_err(crate::Error::into_internal_error)?; - } - } - Ok(()) -} - #[cfg(test)] mod tests { + use std::sync::{Arc, Mutex}; + use super::*; use crate::ErrorCode; @@ -382,7 +400,7 @@ mod tests { #[test] fn preserves_malformed_call_shaped_standalone_message() { - let ParsedIncomingLine::InvalidSingle { error, .. } = + let ParsedIncomingLine::Malformed { error, .. } = parse_incoming_line(r#"{"jsonrpc":"2.0","id":1,"method":"one","result":null}"#) else { panic!("expected one invalid-request error"); @@ -395,7 +413,7 @@ mod tests { fn parses_valid_standalone_response() { assert!(matches!( parse_incoming_line(r#"{"jsonrpc":"2.0","id":1,"result":{"ok":true}}"#), - ParsedIncomingLine::Single(Ok(RawJsonRpcMessage::Response(_))) + ParsedIncomingLine::Single(RawJsonRpcMessage::Response(_)) )); } @@ -415,7 +433,7 @@ mod tests { #[test] fn empty_batch_is_an_invalid_request() { - let ParsedIncomingLine::InvalidSingle { raw, error } = parse_incoming_line("[]") else { + let ParsedIncomingLine::Malformed { raw, error } = parse_incoming_line("[]") else { panic!("expected one invalid-request error"); }; @@ -425,7 +443,7 @@ mod tests { #[test] fn malformed_json_is_a_parse_error() { - let ParsedIncomingLine::InvalidSingle { raw, error } = parse_incoming_line("[") else { + let ParsedIncomingLine::Malformed { raw, error } = parse_incoming_line("[") else { panic!("expected one parse error"); }; @@ -435,11 +453,49 @@ mod tests { #[test] fn valid_json_with_an_invalid_envelope_is_an_invalid_request() { - let ParsedIncomingLine::InvalidSingle { raw, error } = parse_incoming_line("17") else { + let ParsedIncomingLine::Malformed { raw, error } = parse_incoming_line("17") else { panic!("expected one invalid-request error"); }; assert_eq!(raw, "17"); assert_eq!(error.code, ErrorCode::InvalidRequest); } + + #[tokio::test] + async fn multiline_malformed_frame_is_written_as_one_line_value() { + let raw = "not json\r\n{\"jsonrpc\":\"2.0\",\"method\":\"injected\"}".to_string(); + let captured = Arc::new(Mutex::new(Vec::new())); + let outgoing = futures::sink::unfold(captured.clone(), |captured, line| async move { + captured.lock().unwrap().push(line); + Ok::<_, std::io::Error>(captured) + }); + + transport_outgoing_frames_actor( + futures::stream::iter([TransportFrame::Malformed { + raw: raw.clone(), + error: crate::Error::parse_error(), + }]), + outgoing, + ) + .await + .unwrap(); + + let lines = captured.lock().unwrap(); + assert_eq!(lines.len(), 1); + assert!(!lines[0].contains('\r') && !lines[0].contains('\n')); + assert_eq!(serde_json::from_str::(&lines[0]).unwrap(), raw); + } + + #[test] + fn multiline_invalid_json_rpc_value_is_compacted_without_changing_value() { + let raw = "{\n \"jsonrpc\": \"2.0\",\n \"method\": 1\n}".to_string(); + let expected = serde_json::from_str::(&raw).unwrap(); + let line = malformed_line_value(raw).unwrap(); + + assert!(!line.contains('\r') && !line.contains('\n')); + assert_eq!( + serde_json::from_str::(&line).unwrap(), + expected + ); + } } diff --git a/src/agent-client-protocol/src/lib.rs b/src/agent-client-protocol/src/lib.rs index 2b88ba3..ddff711 100644 --- a/src/agent-client-protocol/src/lib.rs +++ b/src/agent-client-protocol/src/lib.rs @@ -94,14 +94,12 @@ pub mod util; pub use capabilities::*; -#[doc(hidden)] -pub use jsonrpc::FramedChannel; pub use jsonrpc::{ Builder, ByteStreams, Channel, ConnectionTo, Dispatch, HandleConnectionClose, HandleDispatchFrom, Handled, INCOMING_TRANSPORT_CLOSED_REASON, IntoHandled, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Lines, NullClose, NullHandler, - RawJsonRpcMessage, RawJsonRpcParams, Responder, ResponseRouter, SentRequest, UntypedMessage, - is_incoming_transport_closed, + RawJsonRpcMessage, RawJsonRpcParams, Responder, ResponseRouter, SentRequest, TransportBatch, + TransportBatchEntry, TransportFrame, UntypedMessage, is_incoming_transport_closed, run::{ChainRun, NullRun, RunWithConnectionTo}, }; pub use jsonrpc::{RequestCancellation, is_cancel_request_notification}; diff --git a/src/agent-client-protocol/src/role/acp.rs b/src/agent-client-protocol/src/role/acp.rs index ade263a..b057beb 100644 --- a/src/agent-client-protocol/src/role/acp.rs +++ b/src/agent-client-protocol/src/role/acp.rs @@ -18,9 +18,9 @@ use crate::schema::{InitializeProxyRequest, METHOD_INITIALIZE_PROXY}; #[cfg(feature = "unstable_protocol_v2")] use crate::schema::{ProtocolVersion, v2}; use crate::util::MatchDispatchFrom; -use crate::{ConnectTo, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, Role, RoleId}; #[cfg(feature = "unstable_protocol_v2")] -use crate::{FramedChannel, RawJsonRpcMessage, RawJsonRpcParams}; +use crate::{Channel, RawJsonRpcMessage, RawJsonRpcParams}; +use crate::{ConnectTo, ConnectionTo, Dispatch, HandleDispatchFrom, Handled, Role, RoleId}; /// The client role - typically an IDE or CLI that controls an agent. /// @@ -397,13 +397,10 @@ impl ConnectTo for AgentProtocolRouter { pipe_protocol_peers_until_closed(client, agent).await } - fn into_framed_channel_and_future( + fn into_channel_and_future( self, - ) -> ( - FramedChannel, - crate::BoxFuture<'static, Result<(), crate::Error>>, - ) { - let (channel_for_caller, channel_for_router) = FramedChannel::duplex(); + ) -> (Channel, crate::BoxFuture<'static, Result<(), crate::Error>>) { + let (channel_for_caller, channel_for_router) = Channel::duplex(); let future = Box::pin(ConnectTo::::connect_to(self, channel_for_router)); (channel_for_caller, future) } @@ -635,14 +632,14 @@ fn send_initialize_error( let response = match frame { TransportFrame::Single(entry) => { - let Some(response) = response_for_entry(entry.as_ref(), &error) else { + let Some(response) = response_for_entry(Ok(entry), &error) else { return Ok(()); }; - TransportFrame::Single(Ok(response)) + TransportFrame::Single(response) } - TransportFrame::InvalidSingle { error, .. } => TransportFrame::Single(Ok( + TransportFrame::Malformed { error, .. } => TransportFrame::Single( RawJsonRpcMessage::response(RequestId::Null, Err(error.clone())), - )), + ), TransportFrame::Batch(batch) => { let responses = batch .iter_results() @@ -672,10 +669,8 @@ async fn reject_initialize( let drain_incoming = async move { while let Some(frame) = rx.next().await { match frame { - TransportFrame::Single(message) => { - message?; - } - TransportFrame::InvalidSingle { error, .. } => { + TransportFrame::Single(_) => {} + TransportFrame::Malformed { error, .. } => { return Err(error); } TransportFrame::Batch(batch) => { @@ -702,7 +697,7 @@ struct RunningProtocolPeer { #[cfg(feature = "unstable_protocol_v2")] impl RunningProtocolPeer { fn new(component: impl ConnectTo) -> Self { - let (FramedChannel { rx, tx }, future) = component.into_framed_channel_and_future(); + let (Channel { rx, tx }, future) = component.into_channel_and_future(); Self { rx, tx, future } } @@ -742,7 +737,7 @@ impl RunningProtocolPeer { } fn send(&self, message: RawJsonRpcMessage) -> Result<(), crate::Error> { - self.send_frame(TransportFrame::Single(Ok(message))) + self.send_frame(TransportFrame::Single(message)) } fn send_frame(&self, frame: TransportFrame) -> Result<(), crate::Error> { @@ -755,8 +750,8 @@ impl RunningProtocolPeer { #[cfg(feature = "unstable_protocol_v2")] fn initialize_message(frame: TransportFrame) -> Result { match frame { - TransportFrame::Single(message) => message, - TransportFrame::InvalidSingle { error, .. } => Err(error), + TransportFrame::Single(message) => Ok(message), + TransportFrame::Malformed { error, .. } => Err(error), TransportFrame::Batch(_) => Err(crate::Error::invalid_request() .data("ACP initialize request and response messages must be sent individually")), } @@ -768,10 +763,10 @@ fn initialize_message_mut( ) -> Result<&mut RawJsonRpcMessage, crate::Error> { let entry = match frame { TransportFrame::Single(entry) => entry, - TransportFrame::InvalidSingle { error, .. } => return Err(error.clone()), + TransportFrame::Malformed { error, .. } => return Err(error.clone()), TransportFrame::Batch(batch) => return batch.first_result_mut(), }; - entry.as_mut().map_err(|error| error.clone()) + Ok(entry) } #[cfg(feature = "unstable_protocol_v2")] @@ -782,12 +777,12 @@ async fn pipe_protocol_peers_until_closed( let ((), (), (), ()) = futures::try_join!( left.future, right.future, - FramedChannel { + Channel { rx: left.rx, tx: right.tx, } .copy(), - FramedChannel { + Channel { rx: right.rx, tx: left.tx, } @@ -804,12 +799,12 @@ async fn pipe_protocol_peers_until_done( ) -> Result<(), crate::Error> { let bridge = Box::pin(async move { let ((), ()) = futures::try_join!( - FramedChannel { + Channel { rx: left.rx, tx: right.tx, } .copy(), - FramedChannel { + Channel { rx: right.rx, tx: left.tx, } diff --git a/src/agent-client-protocol/src/stdio.rs b/src/agent-client-protocol/src/stdio.rs index ff0ae23..12d3748 100644 --- a/src/agent-client-protocol/src/stdio.rs +++ b/src/agent-client-protocol/src/stdio.rs @@ -84,13 +84,13 @@ impl ConnectTo for Stdio { } } - fn into_framed_channel_and_future( + fn into_channel_and_future( self, ) -> ( - crate::FramedChannel, + crate::Channel, crate::BoxFuture<'static, Result<(), crate::Error>>, ) { - let (channel_for_caller, channel_for_stdio) = crate::FramedChannel::duplex(); + let (channel_for_caller, channel_for_stdio) = crate::Channel::duplex(); let future = Box::pin(ConnectTo::::connect_to( self, channel_for_stdio, diff --git a/src/agent-client-protocol/tests/jsonrpc_batch.rs b/src/agent-client-protocol/tests/jsonrpc_batch.rs index a7ca442..8af57b6 100644 --- a/src/agent-client-protocol/tests/jsonrpc_batch.rs +++ b/src/agent-client-protocol/tests/jsonrpc_batch.rs @@ -1,9 +1,9 @@ //! Wire-level regressions for receiving JSON-RPC batch arrays. //! -//! Batch parsing is a transport concern: the public channel boundary continues -//! to carry individual JSON-RPC messages. Calls received in one batch are -//! answered with one consolidated response array, while requests and -//! notifications initiated by the SDK remain individual messages. +//! Batch parsing is a transport concern: the public channel boundary carries +//! batch-aware transport frames. Calls received in one batch are answered with +//! one consolidated response array, while requests and notifications initiated +//! by the SDK remain individual messages. use std::sync::{ Arc, diff --git a/src/agent-client-protocol/tests/jsonrpc_transport_close.rs b/src/agent-client-protocol/tests/jsonrpc_transport_close.rs index 0b87143..0dd8fdf 100644 --- a/src/agent-client-protocol/tests/jsonrpc_transport_close.rs +++ b/src/agent-client-protocol/tests/jsonrpc_transport_close.rs @@ -12,7 +12,8 @@ use std::{ use agent_client_protocol::{ ByteStreams, Channel, ConnectTo, ConnectionTo, Dispatch, Error, Handled, JsonRpcMessage, - JsonRpcRequest, Lines, RawJsonRpcMessage, UntypedMessage, is_incoming_transport_closed, + JsonRpcRequest, Lines, RawJsonRpcMessage, TransportFrame, UntypedMessage, + is_incoming_transport_closed, role::{Role, UntypedRole}, schema::v1::{RequestId, Response}, }; @@ -88,9 +89,8 @@ impl ConnectTo for PendingTransport { struct QueuedClient { started: futures::channel::oneshot::Sender<()>, - escaped: futures::channel::oneshot::Sender< - futures::channel::mpsc::UnboundedSender>, - >, + escaped: + futures::channel::oneshot::Sender>, } impl ConnectTo for QueuedClient { @@ -102,7 +102,7 @@ impl ConnectTo for QueuedClient { )?; channel .tx - .unbounded_send(Ok(message)) + .unbounded_send(TransportFrame::Single(message)) .map_err(Error::into_internal_error)?; drop(self.escaped.send(channel.tx.clone())); let _ = self.started.send(()); @@ -177,18 +177,19 @@ async fn receive_requests_then_close(mut peer: Channel, count: usize) { for _ in 0..count { assert!(matches!( peer.rx.next().await, - Some(Ok(RawJsonRpcMessage::Request(_))) + Some(TransportFrame::Single(RawJsonRpcMessage::Request(_))) )); } drop(peer); } async fn respond_then_close(mut peer: Channel) { - let Some(Ok(RawJsonRpcMessage::Request(request))) = peer.rx.next().await else { + let Some(TransportFrame::Single(RawJsonRpcMessage::Request(request))) = peer.rx.next().await + else { panic!("expected outgoing request"); }; peer.tx - .send(Ok(RawJsonRpcMessage::response( + .send(TransportFrame::Single(RawJsonRpcMessage::response( request.id, Ok(serde_json::to_value(MyResponse { status: "received".into(), @@ -321,18 +322,22 @@ async fn channel_peer_receives_final_response_after_write_half_closes() { let peer = async move { let Channel { mut rx, tx } = peer; - tx.unbounded_send(Ok(RawJsonRpcMessage::request( - "myRequest".into(), - serde_json::json!({}), - RequestId::Number(40), - ) - .unwrap())) - .expect("channel should accept the final request"); + tx.unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::request( + "myRequest".into(), + serde_json::json!({}), + RequestId::Number(40), + ) + .unwrap(), + )) + .expect("channel should accept the final request"); tx.close_channel(); drop(tx); - let Some(Ok(RawJsonRpcMessage::Response(Response::Result { id, result }))) = - rx.next().await + let Some(TransportFrame::Single(RawJsonRpcMessage::Response(Response::Result { + id, + result, + }))) = rx.next().await else { panic!("channel read half closed before the final response"); }; @@ -371,13 +376,15 @@ async fn transport_channel_keeps_read_half_open_after_write_half_closes() { let (channel, transport_future) = ConnectTo::::into_channel_and_future(transport); let Channel { mut rx, tx } = channel; - tx.unbounded_send(Ok(RawJsonRpcMessage::request( - "myRequest".into(), - serde_json::json!({}), - RequestId::Number(41), - ) - .unwrap())) - .expect("transport channel should accept the request"); + tx.unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::request( + "myRequest".into(), + serde_json::json!({}), + RequestId::Number(41), + ) + .unwrap(), + )) + .expect("transport channel should accept the request"); tx.close_channel(); drop(tx); @@ -408,8 +415,10 @@ async fn transport_channel_keeps_read_half_open_after_write_half_closes() { }; let receive_response = async move { - let Some(Ok(RawJsonRpcMessage::Response(Response::Result { id, result }))) = - rx.next().await + let Some(TransportFrame::Single(RawJsonRpcMessage::Response(Response::Result { + id, + result, + }))) = rx.next().await else { panic!("read half closed before delivering the peer's final response"); }; @@ -518,11 +527,9 @@ async fn outgoing_drain_keeps_the_full_duplex_read_half_moving() { assert!( escaped - .unbounded_send(Ok(RawJsonRpcMessage::notification( - "too-late".into(), - serde_json::json!({}), - ) - .unwrap())) + .unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::notification("too-late".into(), serde_json::json!({}),).unwrap() + )) .is_err(), "escaped sender accepted a message after client completion" ); @@ -1066,12 +1073,14 @@ async fn request_finishing_conversion_after_eof_keeps_the_eof_cause() { let connection = tokio::spawn(connection); peer.tx - .send(Ok(RawJsonRpcMessage::request( - "myRequest".into(), - serde_json::json!({}), - RequestId::Number(1), - ) - .unwrap())) + .send(TransportFrame::Single( + RawJsonRpcMessage::request( + "myRequest".into(), + serde_json::json!({}), + RequestId::Number(1), + ) + .unwrap(), + )) .await .expect("send request that exposes the connection handle"); let cx = tokio::time::timeout(TIMEOUT, connection_rx) @@ -1082,7 +1091,7 @@ async fn request_finishing_conversion_after_eof_keeps_the_eof_cause() { tokio::time::timeout(TIMEOUT, peer.rx.next()) .await .expect("handler response was not sent"), - Some(Ok(RawJsonRpcMessage::Response(_))) + Some(TransportFrame::Single(RawJsonRpcMessage::Response(_))) )); let (entered_tx, entered_rx) = futures::channel::oneshot::channel(); @@ -1219,11 +1228,13 @@ async fn response_buffered_before_eof_is_delivered() { Ok(()) }); let respond_then_close = async move { - let Some(Ok(RawJsonRpcMessage::Request(request))) = peer.rx.next().await else { + let Some(TransportFrame::Single(RawJsonRpcMessage::Request(request))) = + peer.rx.next().await + else { panic!("expected outgoing request"); }; peer.tx - .send(Ok(RawJsonRpcMessage::response( + .send(TransportFrame::Single(RawJsonRpcMessage::response( request.id, Ok(serde_json::to_value(MyResponse { status: "received".into(), diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index 6136073..ee78381 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -12,7 +12,7 @@ use agent_client_protocol::schema::{ProtocolVersion, v1, v2}; use agent_client_protocol::{ Agent, AgentProtocolRouter, Builder, ByteStreams, Client, ClientProtocolConnector, ConnectTo, Error, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, NullHandler, RawJsonRpcMessage, Role, - UntypedRole, + TransportFrame, UntypedRole, }; use agent_client_protocol_test::testy::Testy; use futures::StreamExt as _; @@ -341,7 +341,7 @@ async fn assert_malformed_initialize_rejected(params: Map) -> Res channel .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::request( "initialize".into(), Value::Object(params), v1::RequestId::Number(1), @@ -349,7 +349,9 @@ async fn assert_malformed_initialize_rejected(params: Map) -> Res .map_err(Error::into_internal_error)?; while let Some(message) = channel.rx.next().await { - let message = message?; + let TransportFrame::Single(message) = message else { + continue; + }; let RawJsonRpcMessage::Response(response) = message else { continue; }; @@ -1264,7 +1266,7 @@ async fn protocol_router_v2_only_rejects_v1_client() -> Result<(), Error> { channel .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::request( "initialize".into(), json_value(v1_initialize_request(ProtocolVersion::V1))?, v1::RequestId::Number(1), @@ -1272,7 +1274,9 @@ async fn protocol_router_v2_only_rejects_v1_client() -> Result<(), Error> { .map_err(Error::into_internal_error)?; while let Some(message) = channel.rx.next().await { - let message = message?; + let TransportFrame::Single(message) = message else { + continue; + }; let RawJsonRpcMessage::Response(v1::Response::Error { error, .. }) = message else { continue; }; @@ -1709,7 +1713,7 @@ async fn protocol_router_routes_future_protocol_version_to_v2() -> Result<(), Er channel .tx - .unbounded_send(Ok(RawJsonRpcMessage::request( + .unbounded_send(TransportFrame::Single(RawJsonRpcMessage::request( "initialize".into(), json_value(v2_initialize_request(ProtocolVersion::from(3_u16)))?, v1::RequestId::Number(1), @@ -1717,7 +1721,9 @@ async fn protocol_router_routes_future_protocol_version_to_v2() -> Result<(), Er .map_err(Error::into_internal_error)?; while let Some(message) = channel.rx.next().await { - let message = message?; + let TransportFrame::Single(message) = message else { + continue; + }; let RawJsonRpcMessage::Response(v1::Response::Result { result, .. }) = message else { continue; };