diff --git a/Cargo.lock b/Cargo.lock index 346e6357..b3976db5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -25,6 +25,7 @@ dependencies = [ "tokio", "tokio-util", "tracing", + "tracing-subscriber", "uuid", "windows-sys 0.61.2", ] @@ -91,6 +92,7 @@ dependencies = [ "eventsource-stream", "futures", "reqwest", + "serde", "serde_json", "thiserror", "tokio", @@ -98,6 +100,7 @@ dependencies = [ "tower-http 0.7.0", "tracing", "tracing-subscriber", + "tungstenite 0.29.0", "url", "uuid", ] diff --git a/Cargo.toml b/Cargo.toml index 867564c1..7b27c192 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,6 +45,7 @@ tokio = { version = "1.52", default-features = false } rustix = { version = "1", default-features = false, features = ["std", "process"] } tokio-util = { version = "0.7", features = ["compat"] } async-tungstenite = { version = "0.35.0", default-features = false, features = ["tokio-rustls-webpki-roots"] } +tungstenite = "=0.29.0" # Serialization serde = { version = "1.0", features = ["derive", "rc"] } @@ -67,7 +68,9 @@ rmcp = { version = "2.1.0", features = ["server", "transport-io", "schemars"] } clap = { version = "4.5", features = ["derive"] } # HTTP -axum = "0.8" +# The typed capacity-error downcast requires Axum and this crate to share the +# same Tungstenite 0.29 identity. Update both constraints together. +axum = "=0.8.9" reqwest = { version = "0.13", default-features = false, features = ["rustls", "json"] } eventsource-stream = "0.2" url = "2.5" diff --git a/src/agent-client-protocol-http/Cargo.toml b/src/agent-client-protocol-http/Cargo.toml index 799cc6cd..624cb75b 100644 --- a/src/agent-client-protocol-http/Cargo.toml +++ b/src/agent-client-protocol-http/Cargo.toml @@ -35,6 +35,7 @@ server = [ "dep:async-stream", "dep:axum", "dep:futures", + "dep:serde", "dep:serde_json", "dep:tokio", "tokio/macros", @@ -42,12 +43,14 @@ server = [ "tokio/sync", "dep:tower-http", "dep:tracing", + "dep:tungstenite", "dep:uuid", ] [dependencies] agent-client-protocol = { workspace = true, optional = true } futures = { workspace = true, optional = true } +serde = { workspace = true, optional = true } serde_json = { workspace = true, optional = true } thiserror = { workspace = true, optional = true } tokio = { workspace = true, optional = true } @@ -57,6 +60,7 @@ tracing = { workspace = true, optional = true } async-stream = { workspace = true, optional = true } axum = { workspace = true, features = ["ws", "macros"], optional = true } tower-http = { workspace = true, features = ["cors"], optional = true } +tungstenite = { workspace = true, optional = true } uuid = { workspace = true, optional = true } # Client diff --git a/src/agent-client-protocol-http/src/connection.rs b/src/agent-client-protocol-http/src/connection.rs index 3fdbe0f5..9dd599f4 100644 --- a/src/agent-client-protocol-http/src/connection.rs +++ b/src/agent-client-protocol-http/src/connection.rs @@ -143,6 +143,10 @@ impl Connection { self.outbound_transport.subscribe_all_outbound() } + pub(crate) fn enqueue_websocket_text(&self, text: String) -> Result<(), &'static str> { + self.outbound_transport.enqueue_websocket_text(text) + } + pub(crate) fn subscribe_closed(&self) -> watch::Receiver { self.closed_tx.subscribe() } @@ -154,10 +158,7 @@ impl Connection { #[cfg(test)] pub(crate) fn push_all_outbound_for_test(&self, msg: String) -> Result<(), &'static str> { - let OutboundTransport::WebSocket(websocket) = &self.outbound_transport else { - return Err("not a WebSocket connection"); - }; - websocket.all_outbound.push(msg) + self.enqueue_websocket_text(msg) } pub(crate) async fn start_router(self: &Arc) { @@ -250,6 +251,13 @@ impl OutboundTransport { } } + fn enqueue_websocket_text(&self, text: String) -> Result<(), &'static str> { + let Self::WebSocket(websocket) = self else { + return Err("not a WebSocket connection"); + }; + websocket.all_outbound.push(text) + } + #[cfg(test)] fn push_connection_stream_for_test(&self, msg: String) -> Result<(), &'static str> { let Self::Http(http) = self else { diff --git a/src/agent-client-protocol-http/src/lib.rs b/src/agent-client-protocol-http/src/lib.rs index 854c12a5..d0cbc134 100644 --- a/src/agent-client-protocol-http/src/lib.rs +++ b/src/agent-client-protocol-http/src/lib.rs @@ -16,4 +16,4 @@ mod websocket_server; #[cfg(feature = "client")] pub use client::{HttpClient, HttpClientError}; #[cfg(feature = "server")] -pub use server::{AcpHttpServer, CorsOptions, ServerOptions}; +pub use server::{AcpHttpServer, CorsOptions, ServerOptions, WebSocketLimits}; diff --git a/src/agent-client-protocol-http/src/server.rs b/src/agent-client-protocol-http/src/server.rs index 695b1699..5bcd694c 100644 --- a/src/agent-client-protocol-http/src/server.rs +++ b/src/agent-client-protocol-http/src/server.rs @@ -13,6 +13,73 @@ use tower_http::cors::{AllowOrigin, CorsLayer}; use crate::connection::ConnectionRegistry; +/// Finite resource limits for an ACP WebSocket connection. +/// +/// Frame and message limits are enforced by the WebSocket transport. The lower +/// JSON-RPC request limit applies to one complete WebSocket text value, including +/// an entire JSON-RPC batch. An oversized call, or every response-bearing entry +/// in an oversized mixed batch, is rejected without forwarding any part of that +/// text value. Correlated errors leave the physical connection available for +/// other sessions when the bounded error response fits `max_message_size`; +/// otherwise the connection closes with WebSocket code 1009. +/// Limits apply per connection; concurrent connections and later protocol +/// parsing can multiply total process memory use. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct WebSocketLimits { + max_frame_size: usize, + max_message_size: usize, + max_json_rpc_request_size: usize, +} + +impl WebSocketLimits { + /// Creates finite hard frame/message limits and a lower soft request limit. + /// + /// # Panics + /// + /// Panics if any limit is zero or if `max_json_rpc_request_size` exceeds + /// `max_message_size`. + #[must_use] + pub const fn new( + max_frame_size: usize, + max_message_size: usize, + max_json_rpc_request_size: usize, + ) -> Self { + assert!(max_frame_size > 0, "max_frame_size must be positive"); + assert!(max_message_size > 0, "max_message_size must be positive"); + assert!( + max_json_rpc_request_size > 0, + "max_json_rpc_request_size must be positive" + ); + assert!( + max_json_rpc_request_size <= max_message_size, + "soft request limit must not exceed hard message limit" + ); + Self { + max_frame_size, + max_message_size, + max_json_rpc_request_size, + } + } + + /// Returns the maximum accepted WebSocket frame size in bytes. + #[must_use] + pub const fn max_frame_size(self) -> usize { + self.max_frame_size + } + + /// Returns the maximum accepted reassembled WebSocket message size in bytes. + #[must_use] + pub const fn max_message_size(self) -> usize { + self.max_message_size + } + + /// Returns the maximum text value or aggregate batch size before correlated rejection. + #[must_use] + pub const fn max_json_rpc_request_size(self) -> usize { + self.max_json_rpc_request_size + } +} + #[derive(Debug, Clone)] pub struct ServerOptions { pub path: String, @@ -85,17 +152,20 @@ impl CorsOptions { struct ServerState { registry: Arc, cors: CorsOptions, + websocket_limits: Option, } pub struct AcpHttpServer { registry: Arc, options: ServerOptions, + websocket_limits: Option, } impl std::fmt::Debug for AcpHttpServer { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("AcpHttpServer") .field("options", &self.options) + .field("websocket_limits", &self.websocket_limits) .finish_non_exhaustive() } } @@ -109,6 +179,7 @@ impl AcpHttpServer { Self { registry: Arc::new(ConnectionRegistry::new(Arc::new(factory))), options: ServerOptions::default(), + websocket_limits: None, } } @@ -118,6 +189,18 @@ impl AcpHttpServer { self } + /// Opts into finite WebSocket and JSON-RPC request limits. + /// + /// Without this builder, the server retains its existing WebSocket defaults. + /// For example, call this with + /// `WebSocketLimits::new(64 * 1024 * 1024, 64 * 1024 * 1024, 32 * 1024 * 1024)` + /// to set 64 MiB hard transport limits and a 32 MiB soft request limit. + #[must_use] + pub fn with_websocket_limits(mut self, limits: WebSocketLimits) -> Self { + self.websocket_limits = Some(limits); + self + } + pub fn into_router(self) -> Router { let registry = self.registry.clone(); let path = self.options.path.clone(); @@ -125,6 +208,7 @@ impl AcpHttpServer { let state = ServerState { registry: registry.clone(), cors: cors.clone(), + websocket_limits: self.websocket_limits, }; let mut router = Router::new() @@ -187,7 +271,13 @@ async fn handle_get( { return (StatusCode::FORBIDDEN, "WebSocket origin not allowed").into_response(); } - crate::websocket_server::handle_ws_upgrade(state.registry, ws) + let ws = if let Some(limits) = state.websocket_limits { + ws.max_frame_size(limits.max_frame_size()) + .max_message_size(limits.max_message_size()) + } else { + ws + }; + crate::websocket_server::handle_ws_upgrade(state.registry, ws, state.websocket_limits) } Err(_) => crate::http_server::handle_get(state.registry, request).await, } diff --git a/src/agent-client-protocol-http/src/websocket_server.rs b/src/agent-client-protocol-http/src/websocket_server.rs index 4ed7fac9..d5a90161 100644 --- a/src/agent-client-protocol-http/src/websocket_server.rs +++ b/src/agent-client-protocol-http/src/websocket_server.rs @@ -1,22 +1,354 @@ -use std::sync::Arc; +use std::{ + error::Error as _, + fmt, + io::{self, Write as _}, + sync::Arc, +}; -use agent_client_protocol::{RawJsonRpcMessage, TransportFrame}; +use agent_client_protocol::{ + Error as AcpError, RawJsonRpcMessage, TransportFrame, schema::v1::RequestId, +}; use axum::{ - extract::ws::{Message as WsMessage, WebSocket, WebSocketUpgrade}, + extract::ws::{CloseFrame, Message as WsMessage, WebSocket, WebSocketUpgrade, close_code}, http::HeaderValue, response::Response, }; use futures::{SinkExt, StreamExt}; use tracing::{debug, error, info, trace, warn}; +use tungstenite::error::CapacityError; use crate::{ connection::{ConnectionRegistry, OutboundLease}, - protocol::{HEADER_CONNECTION_ID, session_id_from_message}, + protocol::HEADER_CONNECTION_ID, + server::WebSocketLimits, }; +enum OversizedRequests { + Single(RequestId), + Batch { + response: BoundedSoftLimitResponse, + request_count: usize, + }, +} + +enum JsonRpcElement { + Request(RequestId), + Notification, + Response, + Invalid, +} + +struct BatchSoftLimitSeed { + actual_bytes: usize, + max_request_bytes: usize, + max_response_bytes: usize, +} + +impl<'de> serde::de::DeserializeSeed<'de> for BatchSoftLimitSeed { + type Value = Option; + + fn deserialize(self, deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + deserializer.deserialize_seq(BatchSoftLimitVisitor { + actual_bytes: self.actual_bytes, + max_request_bytes: self.max_request_bytes, + max_response_bytes: self.max_response_bytes, + }) + } +} + +struct BatchSoftLimitVisitor { + actual_bytes: usize, + max_request_bytes: usize, + max_response_bytes: usize, +} + +impl<'de> serde::de::Visitor<'de> for BatchSoftLimitVisitor { + type Value = Option; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a JSON-RPC batch") + } + + fn visit_seq(self, mut sequence: A) -> Result + where + A: serde::de::SeqAccess<'de>, + { + let mut writer = BoundedJsonWriter::new(self.max_response_bytes); + let mut entry_count = 0usize; + let mut request_count = 0usize; + while let Some(raw) = sequence.next_element::<&serde_json::value::RawValue>()? { + entry_count = entry_count.saturating_add(1); + if writer.too_long().is_some() { + continue; + } + let request_id = match classify_json_rpc_element(raw.get()) { + Ok(JsonRpcElement::Request(request_id)) => Some(request_id), + Ok(JsonRpcElement::Notification | JsonRpcElement::Response) => None, + Ok(JsonRpcElement::Invalid) | Err(_) => Some(RequestId::Null), + }; + let Some(request_id) = request_id else { + continue; + }; + + request_count = request_count.saturating_add(1); + let separator = if request_count == 1 { b"[" } else { b"," }; + if writer.write_all(separator).is_err() { + continue; + } + let response = + payload_too_large_response(request_id, self.actual_bytes, self.max_request_bytes); + if let Err(error) = serde_json::to_writer(&mut writer, &response) + && writer.too_long().is_none() + { + return Err(::custom(error)); + } + } + + if entry_count == 0 { + return Ok(Some(OversizedRequests::Single(RequestId::Null))); + } + if request_count == 0 { + return Ok(None); + } + if writer.too_long().is_none() { + drop(writer.write_all(b"]")); + } + let response = match writer.too_long() { + Some(too_long) => too_long, + None => writer.finish(), + }; + Ok(Some(OversizedRequests::Batch { + response, + request_count, + })) + } +} + +struct JsonRpcElementVisitor; + +impl<'de> serde::de::Visitor<'de> for JsonRpcElementVisitor { + type Value = JsonRpcElement; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("a JSON-RPC object") + } + + fn visit_map(self, mut object: A) -> Result + where + A: serde::de::MapAccess<'de>, + { + let mut jsonrpc_valid = false; + let mut method_present = false; + let mut method_valid = false; + let mut id = None; + let mut params_valid = true; + let mut result_present = false; + let mut error_present = false; + + while let Some(key) = object.next_key::()? { + match key.as_str() { + "jsonrpc" => { + let raw = object.next_value::<&serde_json::value::RawValue>()?; + jsonrpc_valid = serde_json::from_str::(raw.get()) + .is_ok_and(|version| version == "2.0"); + } + "method" => { + method_present = true; + let raw = object.next_value::<&serde_json::value::RawValue>()?; + method_valid = first_non_whitespace_byte(raw.get()) == Some(b'"'); + } + "id" => { + let raw = object.next_value::<&serde_json::value::RawValue>()?; + id = Some(serde_json::from_str::(raw.get()).ok()); + } + "params" => { + let raw = object.next_value::<&serde_json::value::RawValue>()?; + params_valid = match first_non_whitespace_byte(raw.get()) { + Some(b'{' | b'[') => true, + Some(b'n') => raw.get().trim() == "null", + _ => false, + }; + } + "result" => { + result_present = true; + object.next_value::()?; + } + "error" => { + error_present = true; + object.next_value::()?; + } + _ => { + object.next_value::()?; + } + } + } + + if !method_present && (result_present || error_present) { + return Ok(JsonRpcElement::Response); + } + if !jsonrpc_valid || !method_valid || result_present || error_present || !params_valid { + return Ok(JsonRpcElement::Invalid); + } + match id { + Some(Some(request_id)) => Ok(JsonRpcElement::Request(request_id)), + Some(None) => Ok(JsonRpcElement::Invalid), + None => Ok(JsonRpcElement::Notification), + } + } +} + +enum BoundedSoftLimitResponse { + Text(String), + TooLong { + actual_bytes: usize, + max_bytes: usize, + }, +} + +struct BoundedJsonWriter { + bytes: Vec, + max_bytes: usize, + overflow_at: Option, +} + +impl BoundedJsonWriter { + fn new(max_bytes: usize) -> Self { + Self { + bytes: Vec::with_capacity(max_bytes.min(1_024)), + max_bytes, + overflow_at: None, + } + } + + fn too_long(&self) -> Option { + self.overflow_at + .map(|actual_bytes| BoundedSoftLimitResponse::TooLong { + actual_bytes, + max_bytes: self.max_bytes, + }) + } + + fn finish(self) -> BoundedSoftLimitResponse { + BoundedSoftLimitResponse::Text( + String::from_utf8(self.bytes).expect("serde_json always emits valid UTF-8"), + ) + } +} + +impl io::Write for BoundedJsonWriter { + fn write(&mut self, buffer: &[u8]) -> io::Result { + let attempted_bytes = self.bytes.len().saturating_add(buffer.len()); + if attempted_bytes > self.max_bytes { + self.overflow_at = Some(attempted_bytes); + return Err(io::Error::other( + "serialized JSON-RPC response exceeds configured limit", + )); + } + self.bytes.extend_from_slice(buffer); + Ok(buffer.len()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } +} + +fn first_non_whitespace_byte(text: &str) -> Option { + text.as_bytes() + .iter() + .copied() + .find(|byte| !byte.is_ascii_whitespace()) +} + +fn classify_json_rpc_element(text: &str) -> Result { + let mut deserializer = serde_json::Deserializer::from_str(text); + let element = serde::Deserializer::deserialize_map(&mut deserializer, JsonRpcElementVisitor)?; + deserializer.end()?; + Ok(element) +} + +fn oversized_requests( + text: &str, + max_request_bytes: usize, + max_response_bytes: usize, +) -> Option { + if text.len() <= max_request_bytes { + return None; + } + + match first_non_whitespace_byte(text)? { + b'{' => match classify_json_rpc_element(text).ok()? { + JsonRpcElement::Request(request_id) => Some(OversizedRequests::Single(request_id)), + JsonRpcElement::Invalid => Some(OversizedRequests::Single(RequestId::Null)), + JsonRpcElement::Notification | JsonRpcElement::Response => None, + }, + b'[' => { + let mut deserializer = serde_json::Deserializer::from_str(text); + let requests = serde::de::DeserializeSeed::deserialize( + BatchSoftLimitSeed { + actual_bytes: text.len(), + max_request_bytes, + max_response_bytes, + }, + &mut deserializer, + ) + .ok()?; + deserializer.end().ok()?; + requests + } + _ => None, + } +} + +fn payload_too_large_response( + request_id: RequestId, + actual_bytes: usize, + max_bytes: usize, +) -> RawJsonRpcMessage { + let error = AcpError::new(-32600, "JSON-RPC request exceeds configured size limit").data( + serde_json::json!({ + "kind": "payload_too_large", + "max_bytes": max_bytes, + "actual_bytes": actual_bytes, + }), + ); + RawJsonRpcMessage::response(request_id, Err(error)) +} + +fn serialize_soft_limit_response( + request_id: RequestId, + actual_bytes: usize, + max_request_bytes: usize, + max_response_bytes: usize, +) -> Result { + let mut writer = BoundedJsonWriter::new(max_response_bytes); + let response = payload_too_large_response(request_id, actual_bytes, max_request_bytes); + if let Err(error) = serde_json::to_writer(&mut writer, &response) { + return match writer.too_long() { + Some(too_long) => Ok(too_long), + None => Err(error), + }; + } + Ok(writer.finish()) +} + +fn message_too_long(error: &axum::Error) -> Option<(usize, usize)> { + let error = error.source()?.downcast_ref::()?; + match error { + tungstenite::Error::Capacity(CapacityError::MessageTooLong { size, max_size }) => { + Some((*size, *max_size)) + } + _ => None, + } +} + pub(crate) fn handle_ws_upgrade( registry: Arc, ws: WebSocketUpgrade, + websocket_limits: Option, ) -> Response { let connection_id = ConnectionRegistry::next_connection_id(); let conn_id_for_handler = connection_id.clone(); @@ -32,6 +364,7 @@ pub(crate) fn handle_ws_upgrade( registry_for_handler, conn_id_for_handler, connection, + websocket_limits, ) .await; }); @@ -47,6 +380,7 @@ async fn run_ws( registry: Arc, connection_id: String, connection: Arc, + websocket_limits: Option, ) { let (mut ws_tx, mut ws_rx) = socket.split(); let Some(mut outbound_rx) = connection.subscribe_all_outbound() else { @@ -67,6 +401,7 @@ async fn run_ws( &mut closed, &connection_id, &connection, + websocket_limits, ) .await; @@ -83,6 +418,7 @@ async fn run_ws_message_loop( closed: &mut tokio::sync::watch::Receiver, connection_id: &str, connection: &crate::connection::Connection, + websocket_limits: Option, ) { loop { if *closed.borrow() { @@ -112,12 +448,13 @@ async fn run_ws_message_loop( match msg_result { Some(Ok(WsMessage::Text(text))) => { if !forward_client_text( - text.to_string(), + text.as_str(), ws_tx, outbound_rx, closed, connection_id, connection, + websocket_limits, ) .await { @@ -133,7 +470,29 @@ async fn run_ws_message_loop( warn!(connection_id = %connection_id, "Ignoring binary message (ACP uses text)"); } Some(Err(e)) => { - error!(connection_id = %connection_id, "WebSocket error: {e}"); + if let Some((actual_bytes, max_bytes)) = message_too_long(&e) { + error!( + connection_id = %connection_id, + actual_bytes, + max_bytes, + error_category = "message_too_long", + "Closing WebSocket because an inbound message exceeded the hard limit" + ); + let close = CloseFrame { + code: close_code::SIZE, + reason: "message exceeds configured WebSocket limit".into(), + }; + if let Err(close_error) = + ws_tx.send(WsMessage::Close(Some(close))).await + { + warn!( + connection_id = %connection_id, + "Failed to send WebSocket close 1009: {close_error}" + ); + } + } else { + error!(connection_id = %connection_id, "WebSocket error: {e}"); + } break; } None => break, @@ -144,23 +503,105 @@ async fn run_ws_message_loop( } async fn forward_client_text( - text: String, + text: &str, ws_tx: &mut S, outbound_rx: &mut OutboundLease, closed: &mut tokio::sync::watch::Receiver, connection_id: &str, connection: &crate::connection::Connection, + websocket_limits: Option, ) -> bool where S: futures::Sink + Unpin, { - trace!(connection_id = %connection_id, payload = %text, "Client → Agent: {} bytes", text.len()); - let frame = TransportFrame::parse_json(&text); - if let TransportFrame::Single(parsed) = &frame - && let Some(sid) = session_id_from_message(parsed) - && let RawJsonRpcMessage::Request(req) = parsed + if let Some(limits) = websocket_limits + && let Some(requests) = oversized_requests( + text, + limits.max_json_rpc_request_size(), + limits.max_message_size(), + ) { - trace!(connection_id = %connection_id, session_id = %sid, request_id = ?req.id, "Client → Agent (session)"); + let actual_bytes = text.len(); + let max_bytes = limits.max_json_rpc_request_size(); + let response = match requests { + OversizedRequests::Single(request_id) => { + warn!( + connection_id = %connection_id, + actual_bytes, + max_bytes, + error_category = "payload_too_large", + "Rejecting oversized JSON-RPC request" + ); + serialize_soft_limit_response( + request_id, + actual_bytes, + max_bytes, + limits.max_message_size(), + ) + } + OversizedRequests::Batch { + response, + request_count, + } => { + if matches!(&response, BoundedSoftLimitResponse::Text(_)) { + warn!( + connection_id = %connection_id, + actual_bytes, + max_bytes, + request_count, + error_category = "payload_too_large", + "Rejecting oversized JSON-RPC request batch" + ); + } else { + warn!( + connection_id = %connection_id, + actual_bytes, + max_bytes, + error_category = "payload_too_large", + "Rejecting oversized JSON-RPC request batch" + ); + } + Ok(response) + } + }; + let response = match response { + Ok(BoundedSoftLimitResponse::Text(response)) => response, + Ok(BoundedSoftLimitResponse::TooLong { + actual_bytes, + max_bytes, + }) => { + error!( + connection_id = %connection_id, + actual_bytes, + max_bytes, + error_category = "soft_limit_response_too_long", + "Closing WebSocket because the bounded soft-limit response exceeded the hard limit" + ); + let close = CloseFrame { + code: close_code::SIZE, + reason: "generated response exceeds configured WebSocket limit".into(), + }; + if ws_tx.send(WsMessage::Close(Some(close))).await.is_err() { + warn!(connection_id = %connection_id, "Failed to send WebSocket close 1009"); + } + return false; + } + Err(e) => { + error!(connection_id = %connection_id, "Failed to serialize payload-too-large response: {e}"); + return false; + } + }; + if connection.enqueue_websocket_text(response).is_err() { + error!(connection_id = %connection_id, "WebSocket outbound mailbox closed"); + return false; + } + return true; + } + + trace!(connection_id = %connection_id, "Client → Agent: {} bytes", text.len()); + let mut frame = TransportFrame::parse_json(text); + if let TransportFrame::Malformed { error, .. } = &mut frame { + error.data = None; } if connection.send_frame_to_agent(frame).is_err() { error!(connection_id = %connection_id, "Agent channel closed"); @@ -223,7 +664,7 @@ async fn send_outbound_text(ws_tx: &mut S, text: String, connection_id: &str) where S: futures::Sink + Unpin, { - trace!(connection_id = %connection_id, payload = %text, "Agent → Client: {} bytes", text.len()); + trace!(connection_id = %connection_id, "Agent → Client: {} bytes", text.len()); if ws_tx.send(WsMessage::Text(text.into())).await.is_err() { error!(connection_id = %connection_id, "WebSocket send failed"); false @@ -234,11 +675,24 @@ where #[cfg(test)] mod tests { + use std::{ + io, + path::PathBuf, + sync::{Mutex, OnceLock}, + }; + use agent_client_protocol::{ - Channel, TransportBatch, TransportBatchEntry, TransportFrame, - schema::v1::{RequestId, Response as RpcResponse}, + Agent, Channel, Dispatch, TransportBatch, TransportBatchEntry, TransportFrame, + UntypedMessage, + schema::v1::{ + ReadTextFileRequest, ReadTextFileResponse, RequestId, Response as RpcResponse, + SessionId, + }, + }; + use async_tungstenite::{ + tokio::connect_async, + tungstenite::{Message as ClientWsMessage, protocol::frame::coding::CloseCode}, }; - use async_tungstenite::{tokio::connect_async, tungstenite::Message as ClientWsMessage}; use axum::{Router, extract::WebSocketUpgrade, routing::get}; use futures::{StreamExt as _, future::BoxFuture}; use serde_json::json; @@ -248,16 +702,39 @@ mod tests { time::{Duration, timeout}, }; - use crate::connection::{AgentFactory, ConnectionRegistry}; + use crate::{ + AcpHttpServer, WebSocketLimits, + connection::{AgentFactory, ConnectionRegistry}, + }; use super::*; const ISSUE_288_BURST: usize = 1_025; + const SENSITIVE_TEXT_CONTENT: &str = "TOP_SECRET_PROMPT_SENTINEL"; + const SENSITIVE_BASE64_CONTENT: &str = "TOP_SECRET_BASE64_SENTINEL"; + const SENSITIVE_REQUEST_ID_CONTENT: &str = "TOP_SECRET_REQUEST_ID_SENTINEL"; + const SENSITIVE_MALFORMED_CONTENT: &str = "TOP_SECRET_MALFORMED_SENTINEL"; + const SENSITIVE_TYPED_INVALID_TEXT_CONTENT: &str = "TOP_SECRET_TYPED_INVALID_PROMPT_SENTINEL"; + const SENSITIVE_TYPED_INVALID_BASE64_CONTENT: &str = "TOP_SECRET_TYPED_INVALID_BASE64_SENTINEL"; + const SENSITIVE_TYPED_RESPONSE_CONTENT: &str = "TOP_SECRET_TYPED_RESPONSE_SENTINEL"; struct CapturingAgentFactory { forwarded: mpsc::UnboundedSender, } + #[derive(Clone)] + struct SharedLogWriter(Arc>>); + + impl io::Write for SharedLogWriter { + fn write(&mut self, buffer: &[u8]) -> io::Result { + io::Write::write(&mut *self.0.lock().unwrap(), buffer) + } + + fn flush(&mut self) -> io::Result<()> { + io::Write::flush(&mut *self.0.lock().unwrap()) + } + } + impl AgentFactory for CapturingAgentFactory { fn spawn_agent( &self, @@ -279,12 +756,14 @@ mod tests { break; } } - TransportFrame::Malformed { error, .. } => { - outgoing - .unbounded_send(TransportFrame::Single( - RawJsonRpcMessage::response(RequestId::Null, Err(error)), - )) - .unwrap(); + TransportFrame::Malformed { raw, error } => { + if !test_raw_is_response_only_shape(&raw) { + outgoing + .unbounded_send(TransportFrame::Single( + RawJsonRpcMessage::response(RequestId::Null, Err(error)), + )) + .unwrap(); + } } TransportFrame::Batch(_) => panic!("expected a single JSON-RPC frame"), } @@ -405,6 +884,235 @@ mod tests { } } + fn text_prompt_request_with_size(id: i64, session_id: &str, target_size: usize) -> String { + let request = |text: String| { + json!({ + "jsonrpc": "2.0", + "id": id, + "method": "session/prompt", + "params": { + "sessionId": session_id, + "prompt": [{"type": "text", "text": text}] + } + }) + }; + let empty_request = serde_json::to_string(&request(String::new())).unwrap(); + assert!(empty_request.len() <= target_size); + let padding_size = target_size - empty_request.len(); + assert!(SENSITIVE_TEXT_CONTENT.len() <= padding_size); + let text = format!( + "{SENSITIVE_TEXT_CONTENT}{}", + "x".repeat(padding_size - SENSITIVE_TEXT_CONTENT.len()) + ); + let request = serde_json::to_string(&request(text)).unwrap(); + assert_eq!(request.len(), target_size); + request + } + + fn null_id_text_prompt_request_with_size(session_id: &str, target_size: usize) -> String { + let request = |text: String| { + json!({ + "jsonrpc": "2.0", + "id": null, + "method": "session/prompt", + "params": { + "sessionId": session_id, + "prompt": [{"type": "text", "text": text}] + } + }) + }; + let empty_request = serde_json::to_string(&request(String::new())).unwrap(); + assert!(empty_request.len() <= target_size); + let padding_size = target_size - empty_request.len(); + assert!(SENSITIVE_TEXT_CONTENT.len() <= padding_size); + let text = format!( + "{SENSITIVE_TEXT_CONTENT}{}", + "x".repeat(padding_size - SENSITIVE_TEXT_CONTENT.len()) + ); + let request = serde_json::to_string(&request(text)).unwrap(); + assert_eq!(request.len(), target_size); + request + } + + fn base64_prompt_request_with_size(id: i64, session_id: &str, target_size: usize) -> String { + let request = |data: String| { + json!({ + "jsonrpc": "2.0", + "id": id, + "method": "session/prompt", + "params": { + "sessionId": session_id, + "prompt": [{ + "type": "image", + "data": data, + "mimeType": "image/png" + }] + } + }) + }; + let empty_request = serde_json::to_string(&request(String::new())).unwrap(); + assert!(empty_request.len() <= target_size); + let padding_size = target_size - empty_request.len(); + assert!(SENSITIVE_BASE64_CONTENT.len() <= padding_size); + let data = format!( + "{SENSITIVE_BASE64_CONTENT}{}", + "A".repeat(padding_size - SENSITIVE_BASE64_CONTENT.len()) + ); + let request = serde_json::to_string(&request(data)).unwrap(); + assert_eq!(request.len(), target_size); + request + } + + macro_rules! send_client_message { + ($client:ident, $message:expr $(,)?) => { + timeout(Duration::from_secs(1), $client.send($message)) + .await + .expect("WebSocket send should not hang") + .expect("WebSocket message should be sent") + }; + } + + async fn spawn_capturing_server( + limits: WebSocketLimits, + ) -> ( + std::net::SocketAddr, + Arc, + mpsc::UnboundedReceiver, + tokio::task::JoinHandle<()>, + ) { + let (forwarded_tx, forwarded_rx) = mpsc::unbounded_channel(); + let registry = Arc::new(ConnectionRegistry::new(Arc::new(CapturingAgentFactory { + 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.max_frame_size(limits.max_frame_size()) + .max_message_size(limits.max_message_size()), + Some(limits), + ) + } + } + }), + ); + 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(); + }); + (addr, registry, forwarded_rx, server) + } + + fn test_logs() -> Arc>> { + static LOGS: OnceLock>>> = OnceLock::new(); + + LOGS.get_or_init(|| { + let logs = Arc::new(Mutex::new(Vec::new())); + let writer = SharedLogWriter(logs.clone()); + tracing_subscriber::fmt() + .with_env_filter("agent_client_protocol=trace,agent_client_protocol_http=trace") + .without_time() + .with_ansi(false) + .with_writer(move || writer.clone()) + .try_init() + .expect("test logging subscriber should initialize once"); + logs + }) + .clone() + } + + fn logs_for_connection(logs: &Mutex>, connection_id: &str) -> String { + let logs = String::from_utf8(logs.lock().unwrap().clone()).unwrap(); + logs.lines() + .filter(|line| line.contains(connection_id)) + .collect::>() + .join("\n") + } + + fn all_test_logs(logs: &Mutex>) -> String { + String::from_utf8(logs.lock().unwrap().clone()).unwrap() + } + + fn test_raw_is_response_only_shape(raw: &str) -> bool { + serde_json::from_str::(raw).is_ok_and(|value| { + value.as_object().is_some_and(|object| { + !object.contains_key("method") + && (object.contains_key("result") || object.contains_key("error")) + }) + }) + } + + #[test] + fn soft_limit_batch_preflight_skips_notifications_and_response_shapes() { + let batch = json!([ + { + "jsonrpc": "2.0", + "method": "session/update", + "params": {} + }, + { + "jsonrpc": "2.0", + "id": 80, + "result": {"ok": true} + }, + { + "jsonrpc": "2.0", + "id": 81, + "result": {}, + "error": {"code": -1, "message": "conflicting response fields"} + } + ]) + .to_string(); + + assert!(oversized_requests(&batch, batch.len() - 1, 4_096).is_none()); + } + + #[test] + fn soft_limit_preflight_uses_null_id_for_invalid_single_and_empty_batch() { + let invalid = json!({ + "jsonrpc": "2.0", + "id": 82, + "method": 123, + "params": {} + }) + .to_string(); + assert!(matches!( + oversized_requests(&invalid, invalid.len() - 1, 4_096), + Some(OversizedRequests::Single(RequestId::Null)) + )); + + let empty_batch = "[ ]"; + assert!(matches!( + oversized_requests(empty_batch, empty_batch.len() - 1, 4_096), + Some(OversizedRequests::Single(RequestId::Null)) + )); + } + + #[test] + fn soft_limit_batch_preflight_is_bounded_by_response_budget() { + let batch = format!("[{}]", vec!["0"; 4_096].join(",")); + let max_request_bytes = batch.len() - 1; + let max_response_bytes = 256; + + let Some(OversizedRequests::Batch { + response, + request_count: _, + }) = oversized_requests(&batch, max_request_bytes, max_response_bytes) + else { + panic!("oversized invalid batch should be rejected as a batch"); + }; + assert!(matches!( + response, + BoundedSoftLimitResponse::TooLong { max_bytes: 256, .. } + )); + } + #[tokio::test] async fn websocket_buffers_burst_without_polling_slow_subscriber() { let (forwarded_tx, _forwarded_rx) = mpsc::unbounded_channel(); @@ -440,6 +1148,7 @@ mod tests { &mut closed, &connection_id, &connection, + None, ); let finish = connection.shutdown(); futures::join!(message_loop, finish); @@ -507,6 +1216,7 @@ mod tests { &mut closed, &connection_id, &connection, + None, ) .await; @@ -588,13 +1298,15 @@ mod tests { let inbound = RawJsonRpcMessage::notification("test/inbound".to_string(), serde_json::json!({})) .unwrap(); + let inbound = serde_json::to_string(&inbound).unwrap(); let forward = forward_client_text( - serde_json::to_string(&inbound).unwrap(), + &inbound, &mut ws_tx, &mut outbound_rx, &mut closed, &connection_id, &connection, + None, ); futures::pin_mut!(forward); assert!( @@ -626,6 +1338,7 @@ mod tests { #[tokio::test] async fn malformed_ws_frame_returns_parse_error_response_and_continues() { + let logs = test_logs(); let (forwarded_tx, mut forwarded_rx) = mpsc::unbounded_channel(); let registry = Arc::new(ConnectionRegistry::new(Arc::new(CapturingAgentFactory { forwarded: forwarded_tx, @@ -636,7 +1349,7 @@ mod tests { let registry = registry.clone(); move |ws: WebSocketUpgrade| { let registry = registry.clone(); - async move { handle_ws_upgrade(registry, ws) } + async move { handle_ws_upgrade(registry, ws, None) } } }), ); @@ -646,9 +1359,10 @@ mod tests { axum::serve(listener, app).await.unwrap(); }); let (mut client, _) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let malformed = format!("{{not json {SENSITIVE_TEXT_CONTENT}"); client - .send(ClientWsMessage::Text("{not json".into())) + .send(ClientWsMessage::Text(malformed.into())) .await .unwrap(); @@ -663,7 +1377,8 @@ 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_eq!(value["error"]["data"]["line"], "{not json"); + assert!(value["error"]["data"].is_null()); + assert!(!text.contains(SENSITIVE_TEXT_CONTENT)); let parsed = serde_json::from_value::(value).unwrap(); assert!(matches!( @@ -693,6 +1408,61 @@ mod tests { if notification.method.as_ref() == "test/method" )); + let logs = all_test_logs(&logs); + assert!( + !logs.contains(SENSITIVE_TEXT_CONTENT), + "malformed request content leaked into logs:\n{logs}" + ); + + server.abort(); + } + + #[tokio::test] + async fn malformed_response_shape_does_not_receive_error_and_connection_survives() { + let logs = test_logs(); + let limits = WebSocketLimits::new(4_096, 4_096, 2_048); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, _) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let malformed_response = json!({ + "jsonrpc": "2.0", + "id": 61, + "result": {"secret": SENSITIVE_BASE64_CONTENT}, + "error": {"code": -1, "message": "conflicting response fields"} + }); + + send_client_message!( + client, + ClientWsMessage::Text(serde_json::to_string(&malformed_response).unwrap().into()), + ); + assert!( + timeout(Duration::from_secs(1), client.next()) + .await + .is_err(), + "a malformed response shape must not receive a response" + ); + + let notification = + RawJsonRpcMessage::notification("test/after-response".to_string(), json!({})).unwrap(); + send_client_message!( + client, + ClientWsMessage::Text(serde_json::to_string(¬ification).unwrap().into()), + ); + let forwarded = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .expect("connection should remain usable after a malformed response") + .expect("agent channel should remain open"); + assert!(matches!( + forwarded, + RawJsonRpcMessage::Notification(notification) + if notification.method.as_ref() == "test/after-response" + )); + + let logs = all_test_logs(&logs); + assert!( + !logs.contains(SENSITIVE_BASE64_CONTENT), + "malformed response content leaked into logs:\n{logs}" + ); + server.abort(); } @@ -708,7 +1478,7 @@ mod tests { let registry = registry.clone(); move |ws: WebSocketUpgrade| { let registry = registry.clone(); - async move { handle_ws_upgrade(registry, ws) } + async move { handle_ws_upgrade(registry, ws, None) } } }), ); @@ -759,4 +1529,1157 @@ mod tests { server.abort(); } + + #[tokio::test] + async fn oversized_single_request_batch_returns_batch_error_and_connection_survives() { + const MAX_JSON_RPC_REQUEST_SIZE: usize = 1_024; + + let logs = test_logs(); + let limits = WebSocketLimits::new(4_096, 4_096, MAX_JSON_RPC_REQUEST_SIZE); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, response) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let connection_id = response + .headers() + .get(HEADER_CONNECTION_ID) + .expect("upgrade should include a connection ID") + .to_str() + .unwrap() + .to_owned(); + let request = + text_prompt_request_with_size(31, "session-batch-a", MAX_JSON_RPC_REQUEST_SIZE - 1); + let batch = format!("[{request}]"); + assert_eq!(batch.len(), MAX_JSON_RPC_REQUEST_SIZE + 1); + + send_client_message!(client, ClientWsMessage::Text(batch.into())); + + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("server should reject an oversized JSON-RPC batch") + .expect("connection should remain open") + .expect("soft-limit response should be readable"); + let ClientWsMessage::Text(text) = frame else { + panic!("expected text response, got {frame:?}"); + }; + let response: serde_json::Value = serde_json::from_str(&text).unwrap(); + let entries = response.as_array().expect("response should remain a batch"); + assert_eq!(entries.len(), 1); + assert_eq!(entries[0]["id"], 31); + assert_eq!(entries[0]["error"]["code"], -32600); + assert_eq!(entries[0]["error"]["data"]["kind"], "payload_too_large"); + assert!( + timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .is_err(), + "oversized batch must not be forwarded" + ); + + let request = json!({ + "jsonrpc": "2.0", + "id": 32, + "method": "session/prompt", + "params": { + "sessionId": "session-b", + "prompt": [{"type": "text", "text": "small"}] + } + }); + send_client_message!( + client, + ClientWsMessage::Text(serde_json::to_string(&request).unwrap().into()), + ); + let forwarded = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .expect("small request should be forwarded on the same connection") + .expect("agent channel should remain open"); + assert!(matches!( + forwarded, + RawJsonRpcMessage::Request(request) if request.id == RequestId::Number(32) + )); + + let logs = logs_for_connection(&logs, &connection_id); + assert!(logs.contains("Rejecting oversized JSON-RPC request batch")); + assert!(!logs.contains(SENSITIVE_TEXT_CONTENT)); + + server.abort(); + } + + #[tokio::test] + async fn oversized_mixed_batch_returns_errors_for_request_ids_only() { + const MAX_JSON_RPC_REQUEST_SIZE: usize = 1_024; + const TARGET_SIZE: usize = MAX_JSON_RPC_REQUEST_SIZE + 1; + + let logs = test_logs(); + let limits = WebSocketLimits::new(4_096, 4_096, MAX_JSON_RPC_REQUEST_SIZE); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, response) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let connection_id = response + .headers() + .get(HEADER_CONNECTION_ID) + .expect("upgrade should include a connection ID") + .to_str() + .unwrap() + .to_owned(); + let batch = |padding: String| { + json!([ + { + "jsonrpc": "2.0", + "id": 41, + "method": "session/prompt", + "params": {"prompt": padding} + }, + { + "jsonrpc": "2.0", + "id": null, + "method": "session/prompt", + "params": {} + }, + { + "jsonrpc": "2.0", + "method": "session/update", + "params": {"value": "notification"} + }, + { + "jsonrpc": "2.0", + "id": 99, + "result": {"value": "response"} + } + ]) + }; + let empty = serde_json::to_string(&batch(String::new())).unwrap(); + let padding_size = TARGET_SIZE - empty.len(); + let padding = format!( + "{SENSITIVE_BASE64_CONTENT}{}", + "A".repeat(padding_size - SENSITIVE_BASE64_CONTENT.len()) + ); + let batch = serde_json::to_string(&batch(padding)).unwrap(); + assert_eq!(batch.len(), TARGET_SIZE); + + send_client_message!(client, ClientWsMessage::Text(batch.into())); + + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("server should reject an oversized mixed batch") + .expect("connection should remain open") + .expect("soft-limit response should be readable"); + let ClientWsMessage::Text(text) = frame else { + panic!("expected text response, got {frame:?}"); + }; + let response: serde_json::Value = 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"], 41); + assert!(entries[1]["id"].is_null()); + assert!(entries.iter().all(|entry| entry["error"]["code"] == -32600)); + assert!(!text.contains(SENSITIVE_BASE64_CONTENT)); + assert!( + timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .is_err(), + "no entry from a rejected mixed batch may be forwarded" + ); + + let logs = logs_for_connection(&logs, &connection_id); + assert!(logs.contains("Rejecting oversized JSON-RPC request batch")); + assert!( + !logs.contains(SENSITIVE_BASE64_CONTENT), + "oversized batch content leaked into logs:\n{logs}" + ); + + server.abort(); + } + + #[tokio::test] + async fn oversized_malformed_batch_returns_null_id_errors_and_connection_survives() { + const MAX_JSON_RPC_REQUEST_SIZE: usize = 1_024; + const TARGET_SIZE: usize = MAX_JSON_RPC_REQUEST_SIZE + 1; + + let logs = test_logs(); + let limits = WebSocketLimits::new(4_096, 4_096, MAX_JSON_RPC_REQUEST_SIZE); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, response) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let connection_id = response + .headers() + .get(HEADER_CONNECTION_ID) + .expect("upgrade should include a connection ID") + .to_str() + .unwrap() + .to_owned(); + let batch = |padding: String| { + json!([ + { + "jsonrpc": "2.0", + "id": 70, + "params": {"secret": padding} + }, + 42, + { + "jsonrpc": "2.0", + "id": 71, + "result": {"value": "response"}, + "error": {"code": -1, "message": "conflicting response fields"} + }, + { + "jsonrpc": "2.0", + "method": "session/update", + "params": {} + } + ]) + }; + let empty = serde_json::to_string(&batch(String::new())).unwrap(); + let padding_size = TARGET_SIZE - empty.len(); + let padding = format!( + "{SENSITIVE_TEXT_CONTENT}{}", + "x".repeat(padding_size - SENSITIVE_TEXT_CONTENT.len()) + ); + let batch = serde_json::to_string(&batch(padding)).unwrap(); + assert_eq!(batch.len(), TARGET_SIZE); + + send_client_message!(client, ClientWsMessage::Text(batch.into())); + + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("server should reject an oversized malformed batch") + .expect("connection should remain open") + .expect("soft-limit response should be readable"); + let ClientWsMessage::Text(text) = frame else { + panic!("expected text response, got {frame:?}"); + }; + let response: serde_json::Value = serde_json::from_str(&text).unwrap(); + let entries = response.as_array().expect("response should remain a batch"); + assert_eq!(entries.len(), 2); + assert!(entries.iter().all(|entry| entry["id"].is_null())); + assert!(entries.iter().all(|entry| entry["error"]["code"] == -32600)); + assert!(!text.contains(SENSITIVE_TEXT_CONTENT)); + assert!( + timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .is_err(), + "no entry from a rejected malformed batch may be forwarded" + ); + + let request = json!({ + "jsonrpc": "2.0", + "id": 72, + "method": "session/prompt", + "params": { + "sessionId": "session-b", + "prompt": [{"type": "text", "text": "small"}] + } + }); + send_client_message!( + client, + ClientWsMessage::Text(serde_json::to_string(&request).unwrap().into()), + ); + let forwarded = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .expect("another session should continue after the malformed batch rejection") + .expect("agent channel should remain open"); + assert!(matches!( + forwarded, + RawJsonRpcMessage::Request(request) if request.id == RequestId::Number(72) + )); + + let logs = logs_for_connection(&logs, &connection_id); + assert!(logs.contains("Rejecting oversized JSON-RPC request batch")); + assert!( + !logs.contains(SENSITIVE_TEXT_CONTENT), + "oversized malformed batch content leaked into logs:\n{logs}" + ); + + server.abort(); + } + + #[tokio::test] + async fn oversized_batch_error_larger_than_hard_limit_closes_with_message_too_big() { + let requests = (0..20) + .map(|id| json!({"jsonrpc": "2.0", "id": id, "method": "m"})) + .collect::>(); + let batch = serde_json::to_string(&requests).unwrap(); + let limits = WebSocketLimits::new(batch.len() + 64, batch.len() + 64, batch.len() - 1); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, _) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + + send_client_message!(client, ClientWsMessage::Text(batch.into())); + + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("server should bound an oversized generated response") + .expect("server should send a close frame") + .expect("close frame should be readable"); + let ClientWsMessage::Close(Some(frame)) = frame else { + panic!("expected close frame, got {frame:?}"); + }; + assert_eq!(frame.code, CloseCode::Size); + assert!( + timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .is_err(), + "batch with an unrepresentable error must not be forwarded" + ); + + server.abort(); + } + + #[tokio::test] + async fn oversized_dense_invalid_batch_is_bounded_and_not_forwarded() { + let batch = format!("[{}]", vec!["0"; 4_096].join(",")); + let limits = WebSocketLimits::new(batch.len() + 64, batch.len() + 64, batch.len() - 1); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, _) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + + send_client_message!(client, ClientWsMessage::Text(batch.into())); + + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("server should bound the dense invalid batch response") + .expect("server should send a close frame") + .expect("close frame should be readable"); + let ClientWsMessage::Close(Some(frame)) = frame else { + panic!("expected close frame, got {frame:?}"); + }; + assert_eq!(frame.code, CloseCode::Size); + assert!( + timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .is_err(), + "dense invalid batch must not be forwarded" + ); + + server.abort(); + } + + #[tokio::test] + async fn oversized_json_rpc_request_returns_correlated_error_and_connection_survives() { + const MAX_JSON_RPC_REQUEST_SIZE: usize = 1_024; + + let limits = WebSocketLimits::new(4_096, 4_096, MAX_JSON_RPC_REQUEST_SIZE); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, _) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let request = text_prompt_request_with_size(1, "session-a", MAX_JSON_RPC_REQUEST_SIZE + 1); + + send_client_message!(client, ClientWsMessage::Text(request.into())); + + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("server should reject an oversized JSON-RPC request") + .expect("connection should remain open") + .expect("soft-limit response should be readable"); + let ClientWsMessage::Text(text) = frame else { + panic!("expected text response, got {frame:?}"); + }; + let response: serde_json::Value = serde_json::from_str(&text).unwrap(); + assert_eq!(response["id"], 1); + assert_eq!(response["error"]["code"], -32600); + assert_eq!(response["error"]["data"]["kind"], "payload_too_large"); + assert_eq!( + response["error"]["data"]["max_bytes"], + MAX_JSON_RPC_REQUEST_SIZE + ); + assert_eq!( + response["error"]["data"]["actual_bytes"], + MAX_JSON_RPC_REQUEST_SIZE + 1 + ); + assert!( + timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .is_err(), + "oversized request must not be forwarded" + ); + + let request = json!({ + "jsonrpc": "2.0", + "id": 2, + "method": "session/prompt", + "params": { + "sessionId": "session-b", + "prompt": [{"type": "text", "text": "small"}] + } + }); + send_client_message!( + client, + ClientWsMessage::Text(serde_json::to_string(&request).unwrap().into()), + ); + + let forwarded = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .expect("small request should be forwarded on the same connection") + .expect("agent channel should remain open"); + assert!(matches!( + forwarded, + RawJsonRpcMessage::Request(request) if request.id == RequestId::Number(2) + )); + + server.abort(); + } + + #[tokio::test] + async fn soft_limit_error_keeps_existing_websocket_outbound_order() { + const MAX_JSON_RPC_REQUEST_SIZE: usize = 1_024; + + let (forwarded_tx, _forwarded_rx) = mpsc::unbounded_channel(); + let registry = ConnectionRegistry::new(Arc::new(CapturingAgentFactory { + forwarded: forwarded_tx, + })); + let connection_id = ConnectionRegistry::next_connection_id(); + let connection = registry + .create_websocket_connection_with_id(connection_id.clone()) + .await; + let mut outbound_rx = connection.subscribe_all_outbound().unwrap(); + let mut closed = connection.subscribe_closed(); + connection + .push_all_outbound_for_test("queued-first".to_string()) + .unwrap(); + let request = + text_prompt_request_with_size(51, "session-order", MAX_JSON_RPC_REQUEST_SIZE + 1); + let limits = WebSocketLimits::new(4_096, 4_096, MAX_JSON_RPC_REQUEST_SIZE); + let (mut ws_tx, mut ws_rx) = futures::channel::mpsc::unbounded::(); + + assert!( + forward_client_text( + &request, + &mut ws_tx, + &mut outbound_rx, + &mut closed, + &connection_id, + &connection, + Some(limits), + ) + .await + ); + drain_queued_outbound(&mut ws_tx, &mut outbound_rx, &connection_id).await; + + let WsMessage::Text(first) = ws_rx.next().await.unwrap() else { + panic!("expected queued text frame"); + }; + assert_eq!(first.as_str(), "queued-first"); + let WsMessage::Text(second) = ws_rx.next().await.unwrap() else { + panic!("expected soft-limit error frame"); + }; + let response: serde_json::Value = serde_json::from_str(second.as_str()).unwrap(); + assert_eq!(response["id"], 51); + assert_eq!(response["error"]["code"], -32600); + + registry.remove(&connection_id).await; + connection.shutdown().await; + } + + #[tokio::test] + async fn oversized_null_id_request_returns_correlated_error_and_connection_survives() { + const MAX_JSON_RPC_REQUEST_SIZE: usize = 1_024; + + let limits = WebSocketLimits::new(4_096, 4_096, MAX_JSON_RPC_REQUEST_SIZE); + let logs = test_logs(); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, response) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let connection_id = response + .headers() + .get(HEADER_CONNECTION_ID) + .expect("upgrade should include a connection ID") + .to_str() + .unwrap() + .to_owned(); + let request = + null_id_text_prompt_request_with_size("session-null", MAX_JSON_RPC_REQUEST_SIZE + 1); + + send_client_message!(client, ClientWsMessage::Text(request.into())); + + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("server should reject an oversized null-ID request") + .expect("connection should remain open") + .expect("soft-limit response should be readable"); + let ClientWsMessage::Text(text) = frame else { + panic!("expected text response, got {frame:?}"); + }; + let response: serde_json::Value = serde_json::from_str(&text).unwrap(); + assert!(response["id"].is_null()); + assert_eq!(response["error"]["code"], -32600); + assert_eq!(response["error"]["data"]["kind"], "payload_too_large"); + assert_eq!( + response["error"]["data"]["max_bytes"], + MAX_JSON_RPC_REQUEST_SIZE + ); + assert_eq!( + response["error"]["data"]["actual_bytes"], + MAX_JSON_RPC_REQUEST_SIZE + 1 + ); + assert!( + timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .is_err(), + "oversized null-ID request must not be forwarded" + ); + + let request = json!({ + "jsonrpc": "2.0", + "id": 10, + "method": "session/prompt", + "params": { + "sessionId": "session-b", + "prompt": [{"type": "text", "text": "small"}] + } + }); + send_client_message!( + client, + ClientWsMessage::Text(serde_json::to_string(&request).unwrap().into()), + ); + + let forwarded = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .expect("small request should be forwarded on the same connection") + .expect("agent channel should remain open"); + assert!(matches!( + forwarded, + RawJsonRpcMessage::Request(request) if request.id == RequestId::Number(10) + )); + + let logs = logs_for_connection(&logs, &connection_id); + assert!(logs.contains("Rejecting oversized JSON-RPC request")); + assert!(logs.contains("actual_bytes=1025")); + assert!(logs.contains("max_bytes=1024")); + assert!(logs.contains("error_category=\"payload_too_large\"")); + assert!( + !logs.contains(SENSITIVE_TEXT_CONTENT), + "oversized request content leaked into logs:\n{logs}" + ); + + server.abort(); + } + + #[tokio::test] + async fn json_rpc_request_at_soft_limit_is_forwarded() { + const MAX_JSON_RPC_REQUEST_SIZE: usize = 1_024; + + let limits = WebSocketLimits::new(4_096, 4_096, MAX_JSON_RPC_REQUEST_SIZE); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, _) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let request = + text_prompt_request_with_size(3, "session-boundary", MAX_JSON_RPC_REQUEST_SIZE); + + send_client_message!(client, ClientWsMessage::Text(request.into())); + + let forwarded = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .expect("request at the soft limit should be forwarded") + .expect("agent channel should remain open"); + assert!(matches!( + forwarded, + RawJsonRpcMessage::Request(request) if request.id == RequestId::Number(3) + )); + assert!( + timeout(Duration::from_secs(1), client.next()) + .await + .is_err(), + "request at the soft limit must not receive a payload-too-large error" + ); + + server.abort(); + } + + #[tokio::test] + async fn json_rpc_request_below_soft_limit_is_forwarded_without_logging_content() { + const MAX_JSON_RPC_REQUEST_SIZE: usize = 1_024; + + let logs = test_logs(); + let limits = WebSocketLimits::new(4_096, 4_096, MAX_JSON_RPC_REQUEST_SIZE); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, response) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let connection_id = response + .headers() + .get(HEADER_CONNECTION_ID) + .expect("upgrade should include a connection ID") + .to_str() + .unwrap() + .to_owned(); + let request = + text_prompt_request_with_size(5, "session-below", MAX_JSON_RPC_REQUEST_SIZE - 1); + + send_client_message!(client, ClientWsMessage::Text(request.into())); + + let forwarded = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .expect("request below the soft limit should be forwarded") + .expect("agent channel should remain open"); + assert!(matches!( + forwarded, + RawJsonRpcMessage::Request(request) if request.id == RequestId::Number(5) + )); + assert!( + timeout(Duration::from_secs(1), client.next()) + .await + .is_err(), + "request below the soft limit must not receive a payload-too-large error" + ); + + let connection_logs = logs_for_connection(&logs, &connection_id); + assert!(connection_logs.contains("Client → Agent: 1023 bytes")); + let logs = all_test_logs(&logs); + assert!( + !logs.contains(SENSITIVE_TEXT_CONTENT), + "accepted request content leaked into logs:\n{logs}" + ); + + server.abort(); + } + + #[tokio::test] + async fn accepted_string_request_id_is_not_logged() { + let logs = test_logs(); + let limits = WebSocketLimits::new(4_096, 4_096, 2_048); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, response) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let connection_id = response + .headers() + .get(HEADER_CONNECTION_ID) + .expect("upgrade should include a connection ID") + .to_str() + .unwrap() + .to_owned(); + let request = json!({ + "jsonrpc": "2.0", + "id": SENSITIVE_REQUEST_ID_CONTENT, + "method": "session/prompt", + "params": { + "sessionId": "session-sensitive-id", + "prompt": [{"type": "text", "text": "small"}] + } + }); + + send_client_message!( + client, + ClientWsMessage::Text(serde_json::to_string(&request).unwrap().into()), + ); + let forwarded = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .expect("request should reach the agent") + .expect("agent channel should remain open"); + assert!(matches!( + forwarded, + RawJsonRpcMessage::Request(request) + if request.id == RequestId::Str(SENSITIVE_REQUEST_ID_CONTENT.to_string()) + )); + + let logs = logs_for_connection(&logs, &connection_id); + assert!( + !logs.contains(SENSITIVE_REQUEST_ID_CONTENT), + "request ID leaked into logs:\n{logs}" + ); + + server.abort(); + } + + #[tokio::test] + async fn real_protocol_actor_does_not_log_accepted_invalid_or_malformed_content() { + let logs = test_logs(); + let app = AcpHttpServer::new(agent_client_protocol_test::testy::Testy::new) + .with_websocket_limits(WebSocketLimits::new(4_096, 4_096, 2_048)) + .into_router(); + 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 request = json!({ + "jsonrpc": "2.0", + "id": SENSITIVE_REQUEST_ID_CONTENT, + "method": "session/prompt", + "params": { + "sessionId": "missing-sensitive-session", + "prompt": [ + {"type": "text", "text": SENSITIVE_TEXT_CONTENT}, + { + "type": "image", + "data": SENSITIVE_BASE64_CONTENT, + "mimeType": "image/png" + } + ] + } + }); + + send_client_message!( + client, + ClientWsMessage::Text(serde_json::to_string(&request).unwrap().into()), + ); + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("real protocol actor should answer the request") + .expect("connection should remain open") + .expect("response should be readable"); + assert!(matches!(frame, ClientWsMessage::Text(_))); + + let typed_invalid_request = json!({ + "jsonrpc": "2.0", + "id": "typed-invalid-sensitive-request", + "method": "session/prompt", + "params": { + "sessionId": {"invalid": true}, + "prompt": [ + {"type": "text", "text": SENSITIVE_TYPED_INVALID_TEXT_CONTENT}, + { + "type": "image", + "data": SENSITIVE_TYPED_INVALID_BASE64_CONTENT, + "mimeType": "image/png" + } + ] + } + }); + send_client_message!( + client, + ClientWsMessage::Text( + serde_json::to_string(&typed_invalid_request) + .unwrap() + .into(), + ), + ); + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("real protocol actor should reject typed-invalid params") + .expect("connection should remain open") + .expect("invalid-params response should be readable"); + let ClientWsMessage::Text(typed_invalid_response) = frame else { + panic!("expected text frame, got {frame:?}"); + }; + let typed_invalid_json: serde_json::Value = + serde_json::from_str(&typed_invalid_response).unwrap(); + assert_eq!(typed_invalid_json["id"], "typed-invalid-sensitive-request"); + assert_eq!(typed_invalid_json["error"]["code"], -32602); + assert!( + !typed_invalid_response.contains(SENSITIVE_TYPED_INVALID_TEXT_CONTENT), + "typed-invalid prompt leaked into the JSON-RPC error: {typed_invalid_response}" + ); + assert!( + !typed_invalid_response.contains(SENSITIVE_TYPED_INVALID_BASE64_CONTENT), + "typed-invalid base64 leaked into the JSON-RPC error: {typed_invalid_response}" + ); + + send_client_message!( + client, + ClientWsMessage::Text(format!("{{not json {SENSITIVE_MALFORMED_CONTENT}").into()), + ); + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("real protocol actor should answer malformed input") + .expect("connection should remain open") + .expect("parse-error response should be readable"); + assert!(matches!(frame, ClientWsMessage::Text(_))); + + let logs = all_test_logs(&logs); + let sensitive_logs = logs + .lines() + .filter(|line| { + line.contains(SENSITIVE_TEXT_CONTENT) + || line.contains(SENSITIVE_BASE64_CONTENT) + || line.contains(SENSITIVE_MALFORMED_CONTENT) + || line.contains(SENSITIVE_TYPED_INVALID_TEXT_CONTENT) + || line.contains(SENSITIVE_TYPED_INVALID_BASE64_CONTENT) + }) + .collect::>() + .join("\n"); + assert!( + !logs.contains(SENSITIVE_TEXT_CONTENT), + "accepted prompt leaked into logs:\n{sensitive_logs}" + ); + assert!( + !logs.contains(SENSITIVE_BASE64_CONTENT), + "accepted base64 content leaked into logs:\n{sensitive_logs}" + ); + assert!( + !logs.contains(SENSITIVE_MALFORMED_CONTENT), + "malformed content leaked into logs:\n{sensitive_logs}" + ); + assert!( + !logs.contains(SENSITIVE_TYPED_INVALID_TEXT_CONTENT), + "typed-invalid prompt leaked into logs:\n{sensitive_logs}" + ); + assert!( + !logs.contains(SENSITIVE_TYPED_INVALID_BASE64_CONTENT), + "typed-invalid base64 leaked into logs:\n{sensitive_logs}" + ); + let incoming_logs = logs + .lines() + .filter(|line| line.contains("agent_client_protocol::jsonrpc::incoming_actor")) + .collect::>() + .join("\n"); + assert!( + incoming_logs.contains("transport_single"), + "real incoming actor trace was not captured:\n{logs}" + ); + assert!( + !logs.contains(SENSITIVE_REQUEST_ID_CONTENT), + "request ID leaked into core or HTTP logs:\n{logs}" + ); + + server.abort(); + } + + #[tokio::test] + async fn real_protocol_actor_does_not_log_typed_response_content() { + let logs = test_logs(); + let response_processed = Arc::new(tokio::sync::Notify::new()); + let agent_response_processed = response_processed.clone(); + let app = AcpHttpServer::new(move || { + let response_processed = agent_response_processed.clone(); + Agent + .builder() + .name("typed-response-log-test-agent") + .on_receive_dispatch( + async |dispatch: Dispatch, _connection| { + match dispatch { + Dispatch::Response(result, router) => router.route_with_result(result), + Dispatch::Request(..) | Dispatch::Notification(..) => Ok(()), + } + }, + agent_client_protocol::on_receive_dispatch!(), + ) + .with_spawned(move |connection| { + let response_processed = response_processed.clone(); + async move { + connection + .send_request(ReadTextFileRequest::new( + SessionId::new("typed-response-session"), + PathBuf::from("/tmp/typed-response.txt"), + )) + .block_task() + .await?; + response_processed.notify_one(); + Ok(()) + } + }) + }) + .with_websocket_limits(WebSocketLimits::new(4_096, 4_096, 2_048)) + .into_router(); + 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 frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("real protocol actor should send the typed request") + .expect("connection should remain open") + .expect("typed request should be readable"); + let ClientWsMessage::Text(request) = frame else { + panic!("expected text frame, got {frame:?}"); + }; + let request: serde_json::Value = serde_json::from_str(&request).unwrap(); + assert_eq!(request["method"], "fs/read_text_file"); + let callback_response = json!({ + "jsonrpc": "2.0", + "id": request["id"], + "result": serde_json::to_value(ReadTextFileResponse::new( + SENSITIVE_TYPED_RESPONSE_CONTENT + )).unwrap() + }); + send_client_message!( + client, + ClientWsMessage::Text(serde_json::to_string(&callback_response).unwrap().into()), + ); + timeout(Duration::from_secs(1), response_processed.notified()) + .await + .expect("typed response should be parsed and routed"); + + let logs = all_test_logs(&logs); + let response_parse_logs = logs + .lines() + .filter(|line| line.contains("agent_client_protocol::jsonrpc:")) + .collect::>() + .join("\n"); + assert!( + response_parse_logs.contains("parse ok"), + "real typed-response parser trace was not captured:\n{logs}" + ); + assert!( + !response_parse_logs.contains(SENSITIVE_TYPED_RESPONSE_CONTENT), + "typed response content leaked into parser logs:\n{response_parse_logs}" + ); + + server.abort(); + } + + #[tokio::test] + async fn base64_style_request_above_soft_limit_returns_correlated_error_without_logging_content() + { + const MAX_JSON_RPC_REQUEST_SIZE: usize = 1_024; + + let logs = test_logs(); + let limits = WebSocketLimits::new(4_096, 4_096, MAX_JSON_RPC_REQUEST_SIZE); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, response) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let connection_id = response + .headers() + .get(HEADER_CONNECTION_ID) + .expect("upgrade should include a connection ID") + .to_str() + .unwrap() + .to_owned(); + let request = + base64_prompt_request_with_size(4, "session-image", MAX_JSON_RPC_REQUEST_SIZE + 1); + + send_client_message!(client, ClientWsMessage::Text(request.into())); + + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("server should reject an oversized base64-style request") + .expect("connection should remain open") + .expect("soft-limit response should be readable"); + let ClientWsMessage::Text(text) = frame else { + panic!("expected text response, got {frame:?}"); + }; + let response: serde_json::Value = serde_json::from_str(&text).unwrap(); + assert_eq!(response["id"], 4); + assert_eq!(response["error"]["code"], -32600); + assert_eq!(response["error"]["data"]["kind"], "payload_too_large"); + assert_eq!( + response["error"]["data"]["max_bytes"], + MAX_JSON_RPC_REQUEST_SIZE + ); + assert_eq!( + response["error"]["data"]["actual_bytes"], + MAX_JSON_RPC_REQUEST_SIZE + 1 + ); + assert!( + timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .is_err(), + "oversized base64-style request must not be forwarded" + ); + + let logs = logs_for_connection(&logs, &connection_id); + assert!(logs.contains("Rejecting oversized JSON-RPC request")); + assert!(logs.contains("actual_bytes=1025")); + assert!(logs.contains("max_bytes=1024")); + assert!(logs.contains("error_category=\"payload_too_large\"")); + assert!( + !logs.contains(SENSITIVE_BASE64_CONTENT), + "oversized base64 content leaked into logs:\n{logs}" + ); + + server.abort(); + } + + #[tokio::test] + async fn soft_limit_does_not_misrepresent_notifications_or_responses_as_requests() { + const MAX_JSON_RPC_REQUEST_SIZE: usize = 1_024; + const OVERSIZED: usize = MAX_JSON_RPC_REQUEST_SIZE + 1; + + let limits = WebSocketLimits::new(4_096, 4_096, MAX_JSON_RPC_REQUEST_SIZE); + let (addr, _registry, mut forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, _) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + + let notification = |text: String| { + json!({ + "jsonrpc": "2.0", + "method": "session/update", + "params": {"text": text} + }) + }; + let empty = serde_json::to_string(¬ification(String::new())).unwrap(); + let notification = + serde_json::to_string(¬ification("n".repeat(OVERSIZED - empty.len()))).unwrap(); + assert_eq!(notification.len(), OVERSIZED); + send_client_message!(client, ClientWsMessage::Text(notification.into())); + let forwarded = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .expect("oversized notification should reach the agent") + .expect("agent channel should remain open"); + assert!(matches!(forwarded, RawJsonRpcMessage::Notification(_))); + assert!( + timeout(Duration::from_secs(1), client.next()) + .await + .is_err(), + "notification must not receive a JSON-RPC response" + ); + + let response = |data: String| { + json!({ + "jsonrpc": "2.0", + "id": 8, + "result": {"data": data} + }) + }; + let empty = serde_json::to_string(&response(String::new())).unwrap(); + let response = + serde_json::to_string(&response("r".repeat(OVERSIZED - empty.len()))).unwrap(); + assert_eq!(response.len(), OVERSIZED); + send_client_message!(client, ClientWsMessage::Text(response.into())); + let forwarded = timeout(Duration::from_secs(1), forwarded_rx.recv()) + .await + .expect("oversized response should reach the agent") + .expect("agent channel should remain open"); + assert!(matches!(forwarded, RawJsonRpcMessage::Response(_))); + assert!( + timeout(Duration::from_secs(1), client.next()) + .await + .is_err(), + "response must not receive a payload-too-large error" + ); + + server.abort(); + } + + #[tokio::test] + async fn oversized_ws_message_closes_with_message_too_big() { + const MAX_MESSAGE_SIZE: usize = 1_024; + + drop(test_logs()); + let limits = WebSocketLimits::new(2_048, MAX_MESSAGE_SIZE, 512); + let app = AcpHttpServer::new(agent_client_protocol_test::testy::Testy::new) + .with_websocket_limits(limits) + .into_router(); + 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 request = text_prompt_request_with_size(1, "session-a", MAX_MESSAGE_SIZE + 1); + + send_client_message!(client, ClientWsMessage::Text(request.into())); + + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("server should respond to an oversized message") + .expect("server should send a close frame") + .expect("close frame should be readable"); + let ClientWsMessage::Close(Some(frame)) = frame else { + panic!("expected close frame, got {frame:?}"); + }; + assert_eq!(frame.code, CloseCode::Size); + + server.abort(); + } + + #[tokio::test] + async fn oversized_ws_frame_closes_with_message_too_big() { + const MAX_FRAME_SIZE: usize = 1_024; + + let limits = WebSocketLimits::new(MAX_FRAME_SIZE, 2_048, 512); + let (addr, _registry, _forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, _) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let request = text_prompt_request_with_size(6, "session-frame", MAX_FRAME_SIZE + 1); + + send_client_message!(client, ClientWsMessage::Text(request.into())); + + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("server should respond to an oversized frame") + .expect("server should send a close frame") + .expect("close frame should be readable"); + let ClientWsMessage::Close(Some(frame)) = frame else { + panic!("expected close frame, got {frame:?}"); + }; + assert_eq!(frame.code, CloseCode::Size); + + server.abort(); + } + + #[tokio::test] + async fn hard_limit_close_cleans_up_connection() { + const MAX_MESSAGE_SIZE: usize = 1_024; + + let limits = WebSocketLimits::new(2_048, MAX_MESSAGE_SIZE, 512); + let (addr, registry, _forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, response) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let connection_id = response + .headers() + .get(HEADER_CONNECTION_ID) + .expect("upgrade should include a connection ID") + .to_str() + .unwrap() + .to_owned(); + let connection = timeout(Duration::from_secs(1), async { + loop { + if let Some(connection) = registry.get(&connection_id).await { + break connection; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("WebSocket connection should be registered"); + let closed = connection.subscribe_closed(); + let request = text_prompt_request_with_size(7, "session-cleanup", MAX_MESSAGE_SIZE + 1); + + send_client_message!(client, ClientWsMessage::Text(request.into())); + + let frame = timeout(Duration::from_secs(1), client.next()) + .await + .expect("server should respond to an oversized message") + .expect("server should send a close frame") + .expect("close frame should be readable"); + let ClientWsMessage::Close(Some(frame)) = frame else { + panic!("expected close frame, got {frame:?}"); + }; + assert_eq!(frame.code, CloseCode::Size); + + timeout(Duration::from_secs(1), async { + while registry.len().await != 0 || !*closed.borrow() { + tokio::task::yield_now().await; + } + }) + .await + .expect("hard close should remove and shut down the connection"); + + server.abort(); + } + + #[tokio::test] + async fn abrupt_websocket_disconnect_uses_generic_error_path_and_cleans_up() { + let logs = test_logs(); + let limits = WebSocketLimits::new(4_096, 4_096, 1_024); + let (addr, registry, _forwarded_rx, server) = spawn_capturing_server(limits).await; + let (mut client, response) = connect_async(format!("ws://{addr}/acp")).await.unwrap(); + let connection_id = response + .headers() + .get(HEADER_CONNECTION_ID) + .expect("upgrade should include a connection ID") + .to_str() + .unwrap() + .to_owned(); + let connection = timeout(Duration::from_secs(1), async { + loop { + if let Some(connection) = registry.get(&connection_id).await { + break connection; + } + tokio::task::yield_now().await; + } + }) + .await + .expect("WebSocket connection should be registered"); + let closed = connection.subscribe_closed(); + + futures::AsyncWriteExt::close(client.get_mut()) + .await + .unwrap(); + drop(client); + + timeout(Duration::from_secs(1), async { + while registry.len().await != 0 || !*closed.borrow() { + tokio::task::yield_now().await; + } + }) + .await + .expect("generic WebSocket error should clean up the connection"); + + let logs = logs_for_connection(&logs, &connection_id); + assert!( + logs.contains("WebSocket error:"), + "generic error log missing:\n{logs}" + ); + assert!(!logs.contains("message_too_long")); + assert!(!logs.contains("WebSocket close 1009")); + + server.abort(); + } + + #[test] + fn non_capacity_websocket_error_is_not_message_too_long() { + let error = axum::Error::new(tungstenite::Error::Protocol( + tungstenite::error::ProtocolError::ResetWithoutClosingHandshake, + )); + + assert_eq!(message_too_long(&error), None); + } } diff --git a/src/agent-client-protocol/Cargo.toml b/src/agent-client-protocol/Cargo.toml index 491cc4da..cc4122ca 100644 --- a/src/agent-client-protocol/Cargo.toml +++ b/src/agent-client-protocol/Cargo.toml @@ -70,6 +70,7 @@ clap.workspace = true expect-test.workspace = true tokio = { workspace = true, features = ["io-std", "io-util", "macros", "rt", "rt-multi-thread", "sync", "time"] } tokio-util.workspace = true +tracing-subscriber.workspace = true [lints] workspace = true diff --git a/src/agent-client-protocol/src/jsonrpc.rs b/src/agent-client-protocol/src/jsonrpc.rs index a51795a2..735a403a 100644 --- a/src/agent-client-protocol/src/jsonrpc.rs +++ b/src/agent-client-protocol/src/jsonrpc.rs @@ -2479,10 +2479,7 @@ impl RequestCancellationRegistry { ) .is_some() { - tracing::debug!( - ?id, - "peer reused the ID of a request that is still in flight" - ); + tracing::debug!("peer reused the ID of a request that is still in flight"); } generation }; @@ -4444,9 +4441,8 @@ impl Drop for ResponderDropGuard { }, ) { tracing::debug!( - id = ?self.id, method = %self.method, - ?error, + error_code = ?&error.code, "could not complete abandoned JSON-RPC batch response slot" ); } @@ -4589,7 +4585,7 @@ impl Responder { mut self, response: Result, ) -> Result<(), crate::Error> { - tracing::debug!(id = ?self.id, "respond called"); + tracing::debug!(method = %self.method, "respond called"); self.drop_guard.disarm(); (self.send_fn)(response) } @@ -4606,7 +4602,7 @@ impl Responder { /// Respond to the JSON-RPC request with an error. pub fn respond_with_error(self, error: crate::Error) -> Result<(), crate::Error> { - tracing::debug!(id = ?self.id, ?error, "respond_with_error called"); + tracing::debug!(method = %self.method, error_code = ?&error.code, "respond_with_error called"); self.respond_with_result(Err(error)) } @@ -4791,7 +4787,7 @@ impl ResponseRouter { /// Route an error response to the waiting task. pub fn route_with_error(self, error: crate::Error) -> Result<(), crate::Error> { - tracing::debug!(id = ?self.id, ?error, "error routed to awaiter"); + tracing::debug!(id = ?self.id, error_code = ?&error.code, "error routed to awaiter"); self.route_with_result(Err(error)) } } @@ -5025,24 +5021,21 @@ impl Dispatch { /// * `Ok(Ok(typed))` if this dispatch matches the requested type for its variant /// * `Ok(Err(self))` if it does not match the requested type for its variant /// * `Err` if its method matches the requested type but parsing fails - #[tracing::instrument(skip(self), fields(Request = ?std::any::type_name::(), Notif = ?std::any::type_name::()), level = "trace", ret)] + #[tracing::instrument(skip(self), fields(Request = ?std::any::type_name::(), Notif = ?std::any::type_name::()), level = "trace")] pub(crate) fn into_typed_dispatch( self, ) -> Result, Dispatch>, crate::Error> { - tracing::debug!( - message = ?self, - "into_typed_dispatch" - ); + tracing::debug!(method = %self.method(), "into_typed_dispatch"); match self { Dispatch::Request(message, responder) => { if Req::matches_method(&message.method) { match Req::parse_message(&message.method, &message.params) { Ok(req) => { - tracing::trace!(?req, "parsed ok"); + tracing::trace!("request parsed"); Ok(Ok(Dispatch::Request(req, responder.cast()))) } Err(err) => { - tracing::trace!(?err, "parse error"); + tracing::trace!(error_code = ?&err.code, "parse error"); Err(err) } } @@ -5056,11 +5049,11 @@ impl Dispatch { if Notif::matches_method(&message.method) { match Notif::parse_message(&message.method, &message.params) { Ok(notif) => { - tracing::trace!(?notif, "parse ok"); + tracing::trace!("notification parsed"); Ok(Ok(Dispatch::Notification(notif))) } Err(err) => { - tracing::trace!(?err, "parse error"); + tracing::trace!(error_code = ?&err.code, "parse error"); Err(err) } } @@ -5078,11 +5071,14 @@ impl Dispatch { Ok(value) => { match ::from_value(method, value) { Ok(parsed) => { - tracing::trace!(?parsed, "parse ok"); + tracing::trace!( + response_type = std::any::type_name::(), + "parse ok" + ); Ok(parsed) } Err(err) => { - tracing::trace!(?err, "parse error"); + tracing::trace!(error_code = ?&err.code, "parse error"); return Err(err); } } diff --git a/src/agent-client-protocol/src/jsonrpc/handlers.rs b/src/agent-client-protocol/src/jsonrpc/handlers.rs index efc904b7..79980571 100644 --- a/src/agent-client-protocol/src/jsonrpc/handlers.rs +++ b/src/agent-client-protocol/src/jsonrpc/handlers.rs @@ -109,14 +109,14 @@ where Dispatch::Request(message, responder) => { tracing::debug!( request_type = std::any::type_name::(), - message = ?message, + method = %message.method, "RequestHandler::handle_request" ); if Req::matches_method(&message.method) { match Req::parse_message(&message.method, &message.params) { Ok(req) => { tracing::trace!( - ?req, + request_type = std::any::type_name::(), "RequestHandler::handle_request: parse completed" ); let typed_responder = responder.cast(); @@ -147,7 +147,7 @@ where } Err(err) => { tracing::trace!( - ?err, + error_code = ?&err.code, "RequestHandler::handle_request: parse errored" ); Err(err) @@ -239,14 +239,14 @@ where Dispatch::Notification(message) => { tracing::debug!( request_type = std::any::type_name::(), - message = ?message, + method = %message.method, "NotificationHandler::handle_dispatch" ); if Notif::matches_method(&message.method) { match Notif::parse_message(&message.method, &message.params) { Ok(notif) => { tracing::trace!( - ?notif, + notification_type = std::any::type_name::(), "NotificationHandler::handle_notification: parse completed" ); let result = (self.to_future_hack)( @@ -272,7 +272,7 @@ where } Err(err) => { tracing::trace!( - ?err, + error_code = ?&err.code, "NotificationHandler::handle_notification: parse errored" ); Err(err) diff --git a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs index b326f33b..cc9e2112 100644 --- a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs @@ -100,7 +100,16 @@ pub(super) async fn incoming_protocol_actor( }; message }; - tracing::trace!(message = ?message_result, actor = "incoming_protocol_actor"); + let event = match &message_result { + IncomingProtocolMsg::Transport(TransportFrame::Single(_)) => "transport_single", + IncomingProtocolMsg::Transport(TransportFrame::Batch(_)) => "transport_batch", + IncomingProtocolMsg::Transport(TransportFrame::Malformed { .. }) => { + "transport_malformed" + } + IncomingProtocolMsg::TransportClosed => "transport_closed", + IncomingProtocolMsg::DynamicHandler(_) => "dynamic_handler", + }; + tracing::trace!(event, actor = "incoming_protocol_actor"); match message_result { IncomingProtocolMsg::TransportClosed => { connection.begin_incoming_close(); @@ -132,7 +141,7 @@ pub(super) async fn incoming_protocol_actor( for (message, destination) in entries { match message { Ok(RawJsonRpcMessage::Request(request)) => { - tracing::trace!(method = %request.method, id = ?request.id, "Handling request"); + tracing::trace!(method = %request.method, "Handling request"); let request_method = request.method.to_string(); let request_id = request.id.clone(); let destination = destination @@ -211,7 +220,7 @@ pub(super) async fn incoming_protocol_actor( Response::Error { id, error } => (id, Err(error)), }; - tracing::trace!(?id, "Handling response"); + tracing::trace!("Handling response"); if let Some(pending_reply) = pending_replies.remove(&id) { let result = protocol_compat .incoming_response(&pending_reply.method, result); @@ -267,14 +276,13 @@ pub(super) async fn incoming_protocol_actor( } } else { tracing::warn!( - ?id, "incoming_actor: received response for unknown id, no subscriber found" ); } } Err(error) => { tracing::warn!( - ?error, + error_code = ?&error.code, "Invalid transport input, sending error response" ); let destination = destination.expect( @@ -332,7 +340,7 @@ async fn handle_dynamic_handler_message( new_pending_messages.push(m); } Err(err) => { - tracing::warn!(?err, handler = ?handler.dyn_describe_chain(), "Dynamic handler errored on pending message"); + tracing::warn!(error_code = ?&err.code, handler = ?handler.dyn_describe_chain(), "Dynamic handler errored on pending message"); handle_handler_error(connection, reply_target, method, err)?; } } @@ -482,7 +490,15 @@ pub(super) fn dispatch_from_response( } #[tracing::instrument( - skip(connection, dispatch, dynamic_handlers, handler, pending_messages), + skip( + counterpart, + connection, + dispatch, + dynamic_handlers, + handler, + pending_messages, + request_cancellations + ), fields(method = dispatch.method()), level = "trace", )] @@ -495,11 +511,15 @@ async fn dispatch_dispatch( pending_messages: &mut Vec, request_cancellations: &super::RequestCancellationRegistry, ) -> Result<(), crate::Error> { - tracing::trace!(?dispatch, "dispatch_dispatch"); + let dispatch_kind = match &dispatch { + Dispatch::Notification(_) => "notification", + Dispatch::Request(_, _) => "request", + Dispatch::Response(_, _) => "response", + }; + tracing::trace!(dispatch_kind, "dispatch_dispatch"); let mut retry_any = false; - let id = dispatch.id().cloned(); let method = dispatch.method().to_string(); let error_target = dispatch.handler_error_target(); let _handler_attempt = error_target @@ -512,12 +532,7 @@ async fn dispatch_dispatch( } Ok(false) => {} Err(err) => { - tracing::warn!( - ?method, - ?id, - ?err, - "Request cancellation notification errored" - ); + tracing::warn!(?method, error_code = ?&err.code, "Request cancellation notification errored"); return handle_handler_error(connection, error_target, method, err); } } @@ -529,18 +544,18 @@ async fn dispatch_dispatch( .await { Ok(Handled::Yes) => { - tracing::trace!(?method, ?id, handler = ?handler.describe_chain(), "Handler accepted message"); + tracing::trace!(?method, handler = ?handler.describe_chain(), "Handler accepted message"); return Ok(()); } Ok(Handled::No { message: m, retry }) => { - tracing::trace!(?method, ?id, handler = ?handler.describe_chain(), "Handler declined message"); + tracing::trace!(?method, handler = ?handler.describe_chain(), "Handler declined message"); dispatch = m; retry_any |= retry; } Err(err) => { - tracing::warn!(?method, ?id, ?err, handler = ?handler.describe_chain(), "Handler errored"); + tracing::warn!(?method, error_code = ?&err.code, handler = ?handler.describe_chain(), "Handler errored"); return handle_handler_error(connection, error_target, method, err); } } @@ -553,18 +568,18 @@ async fn dispatch_dispatch( .await { Ok(Handled::Yes) => { - tracing::trace!(?method, ?id, handler = ?dynamic_handler.dyn_describe_chain(), "Dynamic handler accepted message"); + tracing::trace!(?method, handler = ?dynamic_handler.dyn_describe_chain(), "Dynamic handler accepted message"); return Ok(()); } Ok(Handled::No { message: m, retry }) => { - tracing::trace!(?method, ?id, handler = ?dynamic_handler.dyn_describe_chain(), "Dynamic handler declined message"); + tracing::trace!(?method, handler = ?dynamic_handler.dyn_describe_chain(), "Dynamic handler declined message"); retry_any |= retry; dispatch = m; } Err(err) => { - tracing::warn!(?method, ?id, ?err, handler = ?dynamic_handler.dyn_describe_chain(), "Dynamic handler errored"); + tracing::warn!(?method, error_code = ?&err.code, handler = ?dynamic_handler.dyn_describe_chain(), "Dynamic handler errored"); return handle_handler_error(connection, error_target, method, err); } } @@ -588,8 +603,7 @@ async fn dispatch_dispatch( Err(err) => { tracing::warn!( ?method, - ?id, - ?err, + error_code = ?&err.code, handler = "default", "Default handler errored" ); @@ -656,7 +670,7 @@ fn handle_handler_error( None => { tracing::warn!( %method, - ?error, + error_code = ?&error.code, "Ignoring message-processing error because there is no request to answer" ); Ok(()) diff --git a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs index b5114939..a208fe0d 100644 --- a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs @@ -8,11 +8,27 @@ use crate::schema::v1::RequestId; pub type OutgoingMessageTx = mpsc::UnboundedSender; +fn outgoing_message_kind(message: &OutgoingMessage) -> &'static str { + match message { + OutgoingMessage::CloseAfterDraining { .. } => "close_after_draining", + OutgoingMessage::BatchDispatchComplete { .. } => "batch_dispatch_complete", + OutgoingMessage::BatchHandlerAttemptComplete { .. } => "batch_handler_attempt_complete", + OutgoingMessage::AbandonedBatchResponse { .. } => "abandoned_batch_response", + OutgoingMessage::Request { .. } => "request", + OutgoingMessage::Notification { .. } => "notification", + OutgoingMessage::Response { .. } => "response", + OutgoingMessage::UncorrelatedErrorResponse { .. } => "uncorrelated_error_response", + } +} + pub(crate) fn send_raw_message( tx: &OutgoingMessageTx, message: OutgoingMessage, ) -> Result<(), crate::Error> { - tracing::debug!(?message, ?tx, "send_raw_message"); + tracing::debug!( + message_kind = outgoing_message_kind(&message), + "send_raw_message" + ); tx.unbounded_send(message) .map_err(crate::util::internal_error) } @@ -33,7 +49,10 @@ pub(super) async fn outgoing_protocol_actor( let mut drain_waiters = Vec::new(); while let Some(message) = outgoing_rx.next().await { - tracing::debug!(?message, "outgoing_protocol_actor"); + tracing::debug!( + message_kind = outgoing_message_kind(&message), + "outgoing_protocol_actor" + ); // Create the message to be sent over the transport let (json_rpc_message, destination) = match message { @@ -66,7 +85,6 @@ pub(super) async fn outgoing_protocol_actor( destination, } => { tracing::warn!( - ?id, %method, "Completing abandoned JSON-RPC batch request with Internal Error" ); @@ -101,9 +119,8 @@ pub(super) async fn outgoing_protocol_actor( && let Err(error) = readiness.await { tracing::warn!( - ?id, %method, - ?error, + error_code = ?&error.code, "Outgoing request readiness failed" ); if let Some(pending_reply) = pending_replies.remove(&id) { @@ -123,7 +140,7 @@ pub(super) async fn outgoing_protocol_actor( { Ok(request) => request, Err(error) => { - tracing::warn!(?id, %method, ?error, "Failed to prepare outgoing request"); + tracing::warn!(%method, error_code = ?&error.code, "Failed to prepare outgoing request"); if let Some(pending_reply) = pending_replies.remove(&id) { pending_reply.fail(error); } @@ -149,7 +166,7 @@ pub(super) async fn outgoing_protocol_actor( Ok(messages) => messages, Err(error) => { tracing::warn!( - ?error, + error_code = ?&error.code, "Dropping outgoing notification after preparation failed" ); continue; @@ -161,7 +178,7 @@ pub(super) async fn outgoing_protocol_actor( Ok(message) => message, Err(error) => { tracing::warn!( - ?error, + error_code = ?&error.code, "Dropping outgoing notification after serialization failed" ); continue; @@ -180,11 +197,11 @@ pub(super) async fn outgoing_protocol_actor( destination, } => match protocol_compat.outgoing_response(&method, response) { Ok(value) => { - tracing::debug!(?id, "Sending success response"); + tracing::debug!(%method, "Sending success response"); (RawJsonRpcMessage::response(id, Ok(value)), destination) } Err(error) => { - tracing::warn!(?id, %method, ?error, "Sending error response"); + tracing::warn!(%method, error_code = ?&error.code, "Sending error response"); (RawJsonRpcMessage::response(id, Err(error)), destination) } }, diff --git a/src/agent-client-protocol/src/role.rs b/src/agent-client-protocol/src/role.rs index 5c537403..91c1dcf0 100644 --- a/src/agent-client-protocol/src/role.rs +++ b/src/agent-client-protocol/src/role.rs @@ -149,11 +149,16 @@ where Counterpart: Role + HasPeer, Peer: Role, { + let dispatch_kind = match &dispatch { + Dispatch::Notification(_) => "notification", + Dispatch::Request(_, _) => "request", + Dispatch::Response(_, _) => "response", + }; tracing::trace!( method = %dispatch.method(), ?counterpart, ?peer, - ?dispatch, + dispatch_kind, "handle_incoming_dispatch: enter" ); diff --git a/src/agent-client-protocol/src/session.rs b/src/agent-client-protocol/src/session.rs index 47200c07..26267702 100644 --- a/src/agent-client-protocol/src/session.rs +++ b/src/agent-client-protocol/src/session.rs @@ -629,7 +629,11 @@ where let mut output = String::new(); loop { let update = self.read_update().await?; - tracing::trace!(?update, "read_to_string update"); + let update_kind = match &update { + SessionMessage::SessionMessage(_) => "session_message", + SessionMessage::StopReason(_) => "stop_reason", + }; + tracing::trace!(update_kind, "read_to_string update"); match update { SessionMessage::SessionMessage(dispatch) => MatchDispatch::new(dispatch) .if_notification(async |notif: SessionNotification| match notif.update { @@ -767,7 +771,7 @@ where ) -> Result, crate::Error> { // If this is a message for our session, grab it. tracing::trace!( - ?message, + method = %message.method(), handler_session_id = ?self.session_id, "ActiveSessionHandler::handle_dispatch" ); diff --git a/src/agent-client-protocol/src/util.rs b/src/agent-client-protocol/src/util.rs index ed704bcb..302b163f 100644 --- a/src/agent-client-protocol/src/util.rs +++ b/src/agent-client-protocol/src/util.rs @@ -9,21 +9,22 @@ mod typed; pub use typed::{MatchDispatch, MatchDispatchFrom, TypeNotification}; /// Cast from `N` to `M` by serializing/deserialization to/from JSON. +/// +/// Conversion errors include only the failed phase. They deliberately omit +/// the original JSON and serializer error text because response values may +/// contain file contents, terminal output, or other sensitive data. pub fn json_cast(params: N) -> Result where N: serde::Serialize, M: serde::de::DeserializeOwned, { - let json = serde_json::to_value(params).map_err(|e| { + let json = serde_json::to_value(params).map_err(|_| { crate::Error::parse_error().data(serde_json::json!({ - "error": e.to_string(), "phase": "serialization" })) })?; - let m = serde_json::from_value(json.clone()).map_err(|e| { + let m = serde_json::from_value(json).map_err(|_| { crate::Error::parse_error().data(serde_json::json!({ - "error": e.to_string(), - "json": json, "phase": "deserialization" })) })?; @@ -36,21 +37,22 @@ where /// [`Error::invalid_params`](`crate::Error::invalid_params`) (`-32602`) /// instead of a parse error, which is the correct JSON-RPC error code for /// malformed method parameters. +/// +/// Conversion errors include only the failed phase. They deliberately omit +/// the original JSON and serializer error text because method parameters may +/// contain prompts, encoded files, or other sensitive content. pub fn json_cast_params(params: N) -> Result where N: serde::Serialize, M: serde::de::DeserializeOwned, { - let json = serde_json::to_value(params).map_err(|e| { + let json = serde_json::to_value(params).map_err(|_| { crate::Error::internal_error().data(serde_json::json!({ - "error": e.to_string(), "phase": "serialization" })) })?; - let m = serde_json::from_value(json.clone()).map_err(|e| { + let m = serde_json::from_value(json).map_err(|_| { crate::Error::invalid_params().data(serde_json::json!({ - "error": e.to_string(), - "json": json, "phase": "deserialization" })) })?; diff --git a/src/agent-client-protocol/tests/jsonrpc_error_handling.rs b/src/agent-client-protocol/tests/jsonrpc_error_handling.rs index b1376a36..acb7e9ea 100644 --- a/src/agent-client-protocol/tests/jsonrpc_error_handling.rs +++ b/src/agent-client-protocol/tests/jsonrpc_error_handling.rs @@ -16,6 +16,49 @@ use futures::{AsyncRead, AsyncWrite, StreamExt as _}; use serde::{Deserialize, Serialize}; use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; +const SENSITIVE_SERIALIZER_ERROR: &str = "TOP_SECRET_SERIALIZER_ERROR_SENTINEL"; + +#[derive(Clone, Copy)] +struct FailingSerializer; + +impl Serialize for FailingSerializer { + fn serialize(&self, _serializer: S) -> Result + where + S: serde::Serializer, + { + Err(serde::ser::Error::custom(SENSITIVE_SERIALIZER_ERROR)) + } +} + +#[test] +fn json_cast_serialization_errors_do_not_echo_serializer_text() { + let parse_error = + agent_client_protocol::util::json_cast::<_, serde_json::Value>(FailingSerializer) + .expect_err("the custom serializer should fail"); + assert_eq!( + parse_error.code, + agent_client_protocol::ErrorCode::ParseError + ); + assert_eq!( + parse_error.data, + Some(serde_json::json!({"phase": "serialization"})) + ); + assert!(!format!("{parse_error:?}").contains(SENSITIVE_SERIALIZER_ERROR)); + + let params_error = + agent_client_protocol::util::json_cast_params::<_, serde_json::Value>(FailingSerializer) + .expect_err("the custom parameter serializer should fail"); + assert_eq!( + params_error.code, + agent_client_protocol::ErrorCode::InternalError + ); + assert_eq!( + params_error.data, + Some(serde_json::json!({"phase": "serialization"})) + ); + assert!(!format!("{params_error:?}").contains(SENSITIVE_SERIALIZER_ERROR)); +} + /// Test helper to block and wait for a JSON-RPC response. async fn recv( response: SentRequest, @@ -783,10 +826,6 @@ async fn test_bad_request_params_return_invalid_params_and_connection_stays_aliv "code": -32602, "message": "Invalid params", "data": { - "error": "missing field `message`", - "json": { - "content": "hello" - }, "phase": "deserialization" } } diff --git a/src/agent-client-protocol/tests/protocol_v2.rs b/src/agent-client-protocol/tests/protocol_v2.rs index 3dc93dfb..054a6a86 100644 --- a/src/agent-client-protocol/tests/protocol_v2.rs +++ b/src/agent-client-protocol/tests/protocol_v2.rs @@ -2342,9 +2342,10 @@ async fn raw_proxy_session_new_validates_success_and_preserves_valid_response() .block_task() .await .expect_err("missing sessionId must be rejected"); - assert!( - error.to_string().contains("sessionId"), - "unexpected malformed session response error: {error:?}" + assert_eq!(error.code, agent_client_protocol::ErrorCode::ParseError); + assert_eq!( + error.data, + Some(serde_json::json!({"phase": "deserialization"})) ); let valid = UntypedMessage::new( diff --git a/src/agent-client-protocol/tests/session_ordering.rs b/src/agent-client-protocol/tests/session_ordering.rs index cf67f988..32b0f762 100644 --- a/src/agent-client-protocol/tests/session_ordering.rs +++ b/src/agent-client-protocol/tests/session_ordering.rs @@ -1,4 +1,8 @@ -use std::time::Duration; +use std::{ + io, + sync::{Arc, Mutex}, + time::Duration, +}; use agent_client_protocol::{ ActiveSession, Agent, Channel, Client, Conductor, ConnectionTo, RawJsonRpcMessage, Responder, @@ -20,6 +24,20 @@ use agent_client_protocol::{ }; const TIMEOUT: Duration = Duration::from_secs(10); +const SENSITIVE_SESSION_UPDATE_CONTENT: &str = "TOP_SECRET_SESSION_UPDATE_SENTINEL"; + +#[derive(Clone)] +struct SharedLogWriter(Arc>>); + +impl io::Write for SharedLogWriter { + fn write(&mut self, buffer: &[u8]) -> io::Result { + io::Write::write(&mut *self.0.lock().unwrap(), buffer) + } + + fn flush(&mut self) -> io::Result<()> { + io::Write::flush(&mut *self.0.lock().unwrap()) + } +} // Compile-time regressions for the callback future bounds on the two non-blocking // session helpers. The callbacks themselves are `'static`, but their future @@ -99,6 +117,17 @@ mod callback_future_lifetimes { #[tokio::test(flavor = "current_thread")] async fn on_session_start_callback_can_consume_later_session_messages() { + let logs = Arc::new(Mutex::new(Vec::new())); + let writer = SharedLogWriter(logs.clone()); + let subscriber = tracing_subscriber::fmt() + .with_env_filter("agent_client_protocol::session=trace") + .without_time() + .with_ansi(false) + .with_writer(move || writer.clone()) + .finish(); + let dispatch = tracing::Dispatch::new(subscriber); + let _guard = tracing::dispatcher::set_default(&dispatch); + let session_id = SessionId::new("ordered-session"); let new_session_id = session_id.clone(); let prompt_session_id = session_id.clone(); @@ -121,7 +150,7 @@ async fn on_session_start_callback_can_consume_later_session_messages() { connection.send_notification(SessionNotification::new( request.session_id, SessionUpdate::AgentMessageChunk(ContentChunk::new(ContentBlock::Text( - TextContent::new("ordered response"), + TextContent::new(SENSITIVE_SESSION_UPDATE_CONTENT), ))), ))?; responder.respond(PromptResponse::new(StopReason::EndTurn)) @@ -146,7 +175,7 @@ async fn on_session_start_callback_can_consume_later_session_messages() { let text = result_rx .await .map_err(|_| agent_client_protocol::Error::internal_error())?; - assert_eq!(text, "ordered response"); + assert_eq!(text, SENSITIVE_SESSION_UPDATE_CONTENT); Ok(()) }); @@ -154,6 +183,16 @@ async fn on_session_start_callback_can_consume_later_session_messages() { .await .expect("session callback deadlocked the incoming dispatch loop") .expect("session connection failed"); + + let logs = String::from_utf8(logs.lock().unwrap().clone()).unwrap(); + assert!( + logs.contains("read_to_string update"), + "session update trace was not captured:\n{logs}" + ); + assert!( + !logs.contains(SENSITIVE_SESSION_UPDATE_CONTENT), + "session update content leaked into logs:\n{logs}" + ); } #[tokio::test(flavor = "current_thread")]