diff --git a/src/agent-client-protocol/CHANGELOG.md b/src/agent-client-protocol/CHANGELOG.md index d9f72ffe..ae6c4e39 100644 --- a/src/agent-client-protocol/CHANGELOG.md +++ b/src/agent-client-protocol/CHANGELOG.md @@ -2,6 +2,11 @@ ## [Unreleased] +### Added + +- Allow mapped `SentRequest` values to be consumed without implementing `JsonRpcResponse`, and + accept one-shot mappers. + ### Changed - **Breaking:** Make `Channel` the batch-aware `TransportFrame` boundary and remove the hidden diff --git a/src/agent-client-protocol/src/jsonrpc.rs b/src/agent-client-protocol/src/jsonrpc.rs index 033d22fa..2f958cee 100644 --- a/src/agent-client-protocol/src/jsonrpc.rs +++ b/src/agent-client-protocol/src/jsonrpc.rs @@ -4342,7 +4342,7 @@ pub struct SentRequest { method: String, task_tx: TaskTx, response_rx: oneshot::Receiver, - to_result: Box Result + Send>, + to_result: Box Result + Send>, cancellation: SentRequestCancellation, /// Cancellation markers of other (incoming) requests whose cancellation /// should be forwarded to this request. See @@ -4591,7 +4591,7 @@ impl SentRequest { } } -impl SentRequest { +impl SentRequest { /// The id of the outgoing request. #[must_use] pub fn id(&self) -> &RequestId { @@ -4604,10 +4604,14 @@ impl SentRequest { &self.method } - /// Create a new response that maps the result of the response to a new type. + /// Map a successful JSON-RPC response into an application type. + /// + /// The mapped type does not need to implement [`JsonRpcResponse`]. The + /// mapper runs at most once and may consume captured state. JSON-RPC error + /// responses bypass the mapper. pub fn map( self, - map_fn: impl Fn(T) -> Result + 'static + Send, + map_fn: impl FnOnce(T) -> Result + 'static + Send, ) -> SentRequest { SentRequest { id: self.id, @@ -4685,7 +4689,7 @@ impl SentRequest { #[track_caller] pub fn forward_response_to(self, responder: Responder) -> Result<(), crate::Error> where - T: Send, + T: JsonRpcResponse, { let this = self.forward_cancellation_from(responder.cancellation()); @@ -4716,7 +4720,6 @@ impl SentRequest { ) -> Result<(), crate::Error> where F: Future> + 'static + Send, - T: Send, { let task_tx = self.task_tx.clone(); let method = self.method; @@ -4826,10 +4829,7 @@ impl SentRequest { /// - Linear control flow is more natural than callbacks /// /// For handler callbacks, use [`on_receiving_result`](Self::on_receiving_result) instead. - pub async fn block_task(self) -> Result - where - T: Send, - { + pub async fn block_task(self) -> Result { let response = await_response_forwarding_cancellation( self.response_rx, &self.cancellation, @@ -4928,7 +4928,7 @@ impl SentRequest { ) -> Result<(), crate::Error> where F: Future> + 'static + Send, - T: Send, + T: JsonRpcResponse, { self.on_receiving_result(async move |result| match result { Ok(value) => task(value, responder).await, @@ -5008,16 +5008,14 @@ impl SentRequest { ) -> Result<(), crate::Error> where F: Future> + 'static + Send, - T: Send, { - self.consume_with(async move |response| { - match response { - // Run the user's callback on the peer's result. - Ok(result) => task(result).await, - // A response that was never delivered fails the consuming - // task instead of invoking the callback. - Err(err) => Err(err), - } + self.consume_with(move |response| match response { + // Invoke the callback before constructing its future so the + // response value does not need to be `Send` across an await. + Ok(result) => Either::Left(task(result)), + // A response that was never delivered fails the consuming + // task instead of invoking the callback. + Err(err) => Either::Right(future::ready(Err(err))), }) } } diff --git a/src/agent-client-protocol/tests/jsonrpc_advanced.rs b/src/agent-client-protocol/tests/jsonrpc_advanced.rs index 78de0d01..528b8349 100644 --- a/src/agent-client-protocol/tests/jsonrpc_advanced.rs +++ b/src/agent-client-protocol/tests/jsonrpc_advanced.rs @@ -14,7 +14,7 @@ use serde::{Deserialize, Serialize}; use tokio_util::compat::{TokioAsyncReadCompatExt, TokioAsyncWriteCompatExt}; /// Test helper to block and wait for a JSON-RPC response. -async fn recv( +async fn recv( response: SentRequest, ) -> Result { let (tx, rx) = tokio::sync::oneshot::channel(); @@ -215,6 +215,68 @@ async fn test_bidirectional_communication() { .await; } +#[tokio::test(flavor = "current_thread")] +async fn test_map_response_to_application_types() { + use std::rc::Rc; + use tokio::task::LocalSet; + + LocalSet::new() + .run_until(async { + let (server_reader, server_writer, client_reader, client_writer) = setup_test_streams(); + + let server_transport = + agent_client_protocol::ByteStreams::new(server_writer, server_reader); + let server = UntypedRole.builder().on_receive_request( + async |request: PingRequest, + responder: Responder, + _connection: ConnectionTo| { + responder.respond(PongResponse { + value: request.value + 1, + }) + }, + agent_client_protocol::on_receive_request!(), + ); + + tokio::task::spawn_local(async move { + server.connect_to(server_transport).await.ok(); + }); + + let client_transport = + agent_client_protocol::ByteStreams::new(client_writer, client_reader); + let mapper_state = String::from("consumed by the mapper"); + + let result = UntypedRole + .builder() + .connect_with( + client_transport, + async |cx| -> Result<(), agent_client_protocol::Error> { + let response = recv(cx.send_request(PingRequest { value: 10 }).map( + move |response| { + let mapper_state = mapper_state; + assert_eq!(mapper_state, "consumed by the mapper"); + Ok(response.value) + }, + )) + .await?; + + assert_eq!(response, 11_u32); + let non_send_response = cx + .send_request(PingRequest { value: 20 }) + .map(|response| Ok(Rc::new(response.value))) + .block_task() + .await?; + + assert_eq!(*non_send_response, 21_u32); + Ok(()) + }, + ) + .await; + + assert!(result.is_ok(), "Test failed: {result:?}"); + }) + .await; +} + // ============================================================================ // Test 2: Request IDs are properly tracked // ============================================================================