From 3c9502c76b43397cf1300d3a198a8cff8cdf92cd Mon Sep 17 00:00:00 2001 From: Ben Brandt Date: Tue, 21 Jul 2026 13:57:45 +0200 Subject: [PATCH] fix(acp)!: preserve JSON-RPC frames across transport adapters Carry complete single, batch, and malformed frames through component, tracing, HTTP, WebSocket, and polyfill boundaries. Preserve grouped responses and HTTP request provenance, including duplicate and null IDs. BREAKING CHANGE: Channel now carries TransportFrame instead of Result. FramedChannel and into_framed_channel_and_future are removed. --- md/SUMMARY.md | 1 + md/migration_v2.0.md | 26 + md/transport-architecture.md | 261 ++----- .../src/snoop.rs | 8 +- src/agent-client-protocol-http/CHANGELOG.md | 5 + src/agent-client-protocol-http/src/client.rs | 690 ++++++++++++------ .../src/connection.rs | 307 ++++++-- .../src/http_server.rs | 365 ++++++++- .../src/websocket_server.rs | 170 ++++- .../CHANGELOG.md | 5 + .../src/mcp_over_acp/http.rs | 598 +++++++++++++-- src/agent-client-protocol/CHANGELOG.md | 8 + src/agent-client-protocol/src/acp_agent.rs | 8 +- src/agent-client-protocol/src/component.rs | 102 +-- src/agent-client-protocol/src/jsonrpc.rs | 613 ++++------------ .../src/jsonrpc/incoming_actor.rs | 8 +- .../src/jsonrpc/outgoing_actor.rs | 5 +- .../src/jsonrpc/transport_actor.rs | 166 +++-- src/agent-client-protocol/src/lib.rs | 6 +- src/agent-client-protocol/src/role/acp.rs | 47 +- src/agent-client-protocol/src/stdio.rs | 6 +- .../tests/jsonrpc_batch.rs | 8 +- .../tests/jsonrpc_transport_close.rs | 91 ++- .../tests/protocol_v2.rs | 20 +- 24 files changed, 2178 insertions(+), 1346 deletions(-) create mode 100644 md/migration_v2.0.md 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; };