From 3e132c33d1e32d34279168ad23145433e7025b7e Mon Sep 17 00:00:00 2001 From: daniel Date: Tue, 25 Aug 2026 14:28:37 +0100 Subject: [PATCH 1/3] feat(acp): add unstable v2 session injection --- Cargo.lock | 3 +- Cargo.toml | 2 +- src/agent-client-protocol/CHANGELOG.md | 3 + src/agent-client-protocol/Cargo.toml | 4 + .../src/schema/v2_impls.rs | 30 +++++ src/agent-client-protocol/src/session/v2.rs | 38 ++++++ .../tests/schema_session_inject.rs | 84 +++++++++++++ src/agent-client-protocol/tests/session_v2.rs | 112 ++++++++++++++++++ 8 files changed, 273 insertions(+), 3 deletions(-) create mode 100644 src/agent-client-protocol/tests/schema_session_inject.rs diff --git a/Cargo.lock b/Cargo.lock index 74c2bcb..897ca11 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -140,8 +140,7 @@ dependencies = [ [[package]] name = "agent-client-protocol-schema" version = "1.7.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca98360c7bb8cc97d7acd49e2a8a851c3f7bee6b2f0535036d8ab86b5fcd223d" +source = "git+https://github.com/danielkov/agent-client-protocol?rev=af121986c3a7e6a1fd5176485d7c809b5654c088#af121986c3a7e6a1fd5176485d7c809b5654c088" dependencies = [ "anyhow", "derive_more", diff --git a/Cargo.toml b/Cargo.toml index c5b8a21..24f04bf 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,7 +38,7 @@ agent-client-protocol-trace-viewer = { path = "src/agent-client-protocol-trace-v yopo = { package = "agent-client-protocol-yopo", path = "src/yopo" } # Protocol -agent-client-protocol-schema = { version = "=1.7.0", features = ["tracing"] } +agent-client-protocol-schema = { git = "https://github.com/danielkov/agent-client-protocol", rev = "af121986c3a7e6a1fd5176485d7c809b5654c088", features = ["tracing"] } # Core async runtime tokio = { version = "1.52", default-features = false } diff --git a/src/agent-client-protocol/CHANGELOG.md b/src/agent-client-protocol/CHANGELOG.md index 4003afe..a49de91 100644 --- a/src/agent-client-protocol/CHANGELOG.md +++ b/src/agent-client-protocol/CHANGELOG.md @@ -4,6 +4,9 @@ ### Added +- *(unstable-v2)* Expose pending session injection through the + `unstable_session_inject` feature, including typed JSON-RPC dispatch and + `V2Session` helpers to inject, replace, and revoke typed content. - *(unstable-v2)* Add runnable draft-v2 agent and one-shot client examples. The agent implements the complete baseline session lifecycle; the client handles permissions, projects chunk and snapshot updates by message ID, and waits for diff --git a/src/agent-client-protocol/Cargo.toml b/src/agent-client-protocol/Cargo.toml index 1eb857a..ffbf490 100644 --- a/src/agent-client-protocol/Cargo.toml +++ b/src/agent-client-protocol/Cargo.toml @@ -43,6 +43,10 @@ unstable_mcp_over_acp = ["agent-client-protocol-schema/unstable_mcp_over_acp"] unstable_plan_operations = ["agent-client-protocol-schema/unstable_plan_operations"] unstable_session_compaction = ["agent-client-protocol-schema/unstable_session_compaction"] unstable_session_fork = ["agent-client-protocol-schema/unstable_session_fork"] +unstable_session_inject = [ + "unstable_protocol_v2", + "agent-client-protocol-schema/unstable_session_inject", +] unstable_tool_call_name = ["agent-client-protocol-schema/unstable_tool_call_name"] unstable_protocol_v2 = ["agent-client-protocol-schema/unstable_protocol_v2"] diff --git a/src/agent-client-protocol/src/schema/v2_impls.rs b/src/agent-client-protocol/src/schema/v2_impls.rs index 8cae486..bc6fb38 100644 --- a/src/agent-client-protocol/src/schema/v2_impls.rs +++ b/src/agent-client-protocol/src/schema/v2_impls.rs @@ -263,6 +263,24 @@ impl_v2_jsonrpc_request!( "session/set_config_option" ); impl_v2_jsonrpc_request!(v2::PromptRequest, v2::PromptResponse, "session/prompt"); +#[cfg(feature = "unstable_session_inject")] +impl_v2_jsonrpc_request!( + v2::InjectSessionRequest, + v2::InjectSessionResponse, + "session/inject" +); +#[cfg(feature = "unstable_session_inject")] +impl_v2_jsonrpc_request!( + v2::RevokeInjectSessionRequest, + v2::RevokeInjectSessionResponse, + "session/revoke_inject" +); +#[cfg(feature = "unstable_session_inject")] +impl_v2_jsonrpc_request!( + v2::ReplaceInjectSessionRequest, + v2::ReplaceInjectSessionResponse, + "session/replace_inject" +); #[cfg(feature = "unstable_mcp_over_acp")] impl_v2_jsonrpc_request!(v2::MessageMcpRequest, v2::MessageMcpResponse, "mcp/message"); @@ -316,6 +334,12 @@ impl_v2_jsonrpc_request_enum!(v2::ClientRequest { CloseSessionRequest => "session/close", SetSessionConfigOptionRequest => "session/set_config_option", PromptRequest => "session/prompt", + #[cfg(feature = "unstable_session_inject")] + InjectSessionRequest => "session/inject", + #[cfg(feature = "unstable_session_inject")] + RevokeInjectSessionRequest => "session/revoke_inject", + #[cfg(feature = "unstable_session_inject")] + ReplaceInjectSessionRequest => "session/replace_inject", #[cfg(feature = "unstable_mcp_over_acp")] MessageMcpRequest => "mcp/message", [ext] ExtMethodRequest, @@ -340,6 +364,12 @@ impl_v2_jsonrpc_response_enum!(v2::AgentResponse { CloseSessionResponse => "session/close", SetSessionConfigOptionResponse => "session/set_config_option", PromptResponse => "session/prompt", + #[cfg(feature = "unstable_session_inject")] + InjectSessionResponse => "session/inject", + #[cfg(feature = "unstable_session_inject")] + RevokeInjectSessionResponse => "session/revoke_inject", + #[cfg(feature = "unstable_session_inject")] + ReplaceInjectSessionResponse => "session/replace_inject", #[cfg(feature = "unstable_mcp_over_acp")] MessageMcpResponse => "mcp/message", [ext] ExtMethodResponse, diff --git a/src/agent-client-protocol/src/session/v2.rs b/src/agent-client-protocol/src/session/v2.rs index 4d78c2f..eba50c2 100644 --- a/src/agent-client-protocol/src/session/v2.rs +++ b/src/agent-client-protocol/src/session/v2.rs @@ -809,6 +809,44 @@ where ) } + /// Inject content for pending delivery to this session. + #[cfg(feature = "unstable_session_inject")] + pub fn inject( + &self, + mode: v2::SessionInjectMode, + content: Vec, + ) -> SentRequest { + self.connection.send_request_to( + Agent, + v2::InjectSessionRequest::new(self.session_id.clone(), mode, content), + ) + } + + /// Revoke a pending injected message. + #[cfg(feature = "unstable_session_inject")] + pub fn revoke_inject( + &self, + message_id: impl Into, + ) -> SentRequest { + self.connection.send_request_to( + Agent, + v2::RevokeInjectSessionRequest::new(self.session_id.clone(), message_id), + ) + } + + /// Replace the content of a pending injected message. + #[cfg(feature = "unstable_session_inject")] + pub fn replace_inject( + &self, + message_id: impl Into, + content: Vec, + ) -> SentRequest { + self.connection.send_request_to( + Agent, + v2::ReplaceInjectSessionRequest::new(self.session_id.clone(), message_id, content), + ) + } + /// Ask the agent to cancel the session's current foreground work. /// /// This is independent from cancelling a prompt's [`SentRequest`]. diff --git a/src/agent-client-protocol/tests/schema_session_inject.rs b/src/agent-client-protocol/tests/schema_session_inject.rs new file mode 100644 index 0000000..2f2cac4 --- /dev/null +++ b/src/agent-client-protocol/tests/schema_session_inject.rs @@ -0,0 +1,84 @@ +#![cfg(feature = "unstable_session_inject")] + +use agent_client_protocol::{JsonRpcMessage, JsonRpcResponse, schema::v2}; +use serde_json::json; + +#[test] +fn v2_session_inject_requests_serialize_and_dispatch() { + let content = vec![v2::ContentBlock::Text(v2::TextContent::new("steer now"))]; + let inject = + v2::InjectSessionRequest::new("session-1", v2::SessionInjectMode::Steer, content.clone()); + let untyped = inject.to_untyped_message().unwrap(); + assert_eq!(untyped.method, "session/inject"); + assert_eq!( + untyped.params, + json!({ + "sessionId": "session-1", + "mode": "steer", + "content": [{ "type": "text", "text": "steer now" }] + }) + ); + assert!(matches!( + v2::ClientRequest::parse_message(untyped.method(), untyped.params()).unwrap(), + v2::ClientRequest::InjectSessionRequest(_) + )); + assert!(matches!( + v2::AgentResponse::from_value( + "session/inject", + json!({ "messageId": "message-1" }), + ) + .unwrap(), + v2::AgentResponse::InjectSessionResponse(response) + if response.message_id == v2::MessageId::new("message-1") + )); + + let revoke = v2::RevokeInjectSessionRequest::new("session-1", v2::MessageId::new("message-1")); + assert_eq!( + revoke.to_untyped_message().unwrap().params, + json!({ "sessionId": "session-1", "messageId": "message-1" }) + ); + assert!(matches!( + v2::ClientRequest::parse_message( + "session/revoke_inject", + &json!({ "sessionId": "session-1", "messageId": "message-1" }), + ) + .unwrap(), + v2::ClientRequest::RevokeInjectSessionRequest(_) + )); + assert!(matches!( + v2::AgentResponse::from_value("session/revoke_inject", json!({})).unwrap(), + v2::AgentResponse::RevokeInjectSessionResponse(_) + )); + + let replace = + v2::ReplaceInjectSessionRequest::new("session-1", v2::MessageId::new("message-1"), content); + assert_eq!( + replace.to_untyped_message().unwrap().params, + json!({ + "sessionId": "session-1", + "messageId": "message-1", + "content": [{ "type": "text", "text": "steer now" }] + }) + ); + assert!(matches!( + v2::ClientRequest::parse_message( + "session/replace_inject", + &json!({ + "sessionId": "session-1", + "messageId": "message-1", + "content": [{ "type": "text", "text": "steer now" }] + }), + ) + .unwrap(), + v2::ClientRequest::ReplaceInjectSessionRequest(_) + )); + assert!(matches!( + v2::AgentResponse::from_value( + "session/replace_inject", + json!({ "messageId": "message-1" }), + ) + .unwrap(), + v2::AgentResponse::ReplaceInjectSessionResponse(response) + if response.message_id == v2::MessageId::new("message-1") + )); +} diff --git a/src/agent-client-protocol/tests/session_v2.rs b/src/agent-client-protocol/tests/session_v2.rs index 87943da..e828f12 100644 --- a/src/agent-client-protocol/tests/session_v2.rs +++ b/src/agent-client-protocol/tests/session_v2.rs @@ -1004,6 +1004,118 @@ async fn v2_session_commands_cover_configuration_and_close() { .expect("v2 session command test failed"); } +#[cfg(feature = "unstable_session_inject")] +#[tokio::test(flavor = "current_thread")] +async fn v2_session_inject_helpers_preserve_typed_content_and_message_ids() { + let session_id = v2::SessionId::new("inject-session"); + let agent_session_id = session_id.clone(); + + let agent = Agent + .v2() + .on_receive_request( + async |request: v2::InitializeRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(initialize_response(request.protocol_version)) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async move |_request: v2::NewSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + responder.respond(v2::NewSessionResponse::new(agent_session_id.clone())) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |request: v2::InjectSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + assert_eq!(request.session_id, v2::SessionId::new("inject-session")); + assert_eq!(request.mode, v2::SessionInjectMode::Steer); + assert_eq!( + request.content, + vec![v2::ContentBlock::Text(v2::TextContent::new("steer"))] + ); + responder.respond(v2::InjectSessionResponse::new(v2::MessageId::new( + "message-1", + ))) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |request: v2::RevokeInjectSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + assert_eq!(request.session_id, v2::SessionId::new("inject-session")); + assert_eq!(request.message_id, v2::MessageId::new("message-1")); + responder.respond(v2::RevokeInjectSessionResponse::new()) + }, + agent_client_protocol::on_receive_request!(), + ) + .on_receive_request( + async |request: v2::ReplaceInjectSessionRequest, + responder: Responder, + _connection: V2ConnectionTo| { + assert_eq!(request.session_id, v2::SessionId::new("inject-session")); + assert_eq!(request.message_id, v2::MessageId::new("message-1")); + assert_eq!( + request.content, + vec![v2::ContentBlock::Text(v2::TextContent::new("replacement"))] + ); + responder.respond(v2::ReplaceInjectSessionResponse::new(request.message_id)) + }, + agent_client_protocol::on_receive_request!(), + ); + + let client = Client.v2().connect_with(agent, async move |connection| { + connection + .send_request(v2::InitializeRequest::new( + ProtocolVersion::V2, + implementation(), + )) + .block_task() + .await?; + + let opened = connection + .build_session(cwd()?) + .start_session() + .block_task() + .await?; + let (session, _) = opened.into_parts(); + + let injected = session + .inject( + v2::SessionInjectMode::Steer, + vec![v2::ContentBlock::Text(v2::TextContent::new("steer"))], + ) + .block_task() + .await?; + assert_eq!(injected.message_id, v2::MessageId::new("message-1")); + + let replaced = session + .replace_inject( + injected.message_id, + vec![v2::ContentBlock::Text(v2::TextContent::new("replacement"))], + ) + .block_task() + .await?; + assert_eq!(replaced.message_id, v2::MessageId::new("message-1")); + + session + .revoke_inject(replaced.message_id) + .block_task() + .await?; + Ok(()) + }); + + tokio::time::timeout(TIMEOUT, client) + .await + .expect("v2 session inject helper test timed out") + .expect("v2 session inject helper test failed"); +} + #[tokio::test(flavor = "current_thread")] async fn dropping_v2_session_does_not_unregister_update_handling() { let session_id = v2::SessionId::new("dropped-handle"); From c1ba83b2bffd47a6ce49aa87ea0dca7c3f83ec73 Mon Sep 17 00:00:00 2001 From: daniel Date: Tue, 25 Aug 2026 14:30:13 +0100 Subject: [PATCH 2/3] build(acp): pin reviewed session injection schema --- Cargo.lock | 2 +- Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 897ca11..1df89af 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -140,7 +140,7 @@ dependencies = [ [[package]] name = "agent-client-protocol-schema" version = "1.7.0" -source = "git+https://github.com/danielkov/agent-client-protocol?rev=af121986c3a7e6a1fd5176485d7c809b5654c088#af121986c3a7e6a1fd5176485d7c809b5654c088" +source = "git+https://github.com/danielkov/agent-client-protocol?rev=6e7e044f9464c4fd652d90699a09e9edc8b3bbad#6e7e044f9464c4fd652d90699a09e9edc8b3bbad" dependencies = [ "anyhow", "derive_more", diff --git a/Cargo.toml b/Cargo.toml index 24f04bf..21417a8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,7 +38,7 @@ agent-client-protocol-trace-viewer = { path = "src/agent-client-protocol-trace-v yopo = { package = "agent-client-protocol-yopo", path = "src/yopo" } # Protocol -agent-client-protocol-schema = { git = "https://github.com/danielkov/agent-client-protocol", rev = "af121986c3a7e6a1fd5176485d7c809b5654c088", features = ["tracing"] } +agent-client-protocol-schema = { git = "https://github.com/danielkov/agent-client-protocol", rev = "6e7e044f9464c4fd652d90699a09e9edc8b3bbad", features = ["tracing"] } # Core async runtime tokio = { version = "1.52", default-features = false } From 2f039993d1d6ed8da35b38c31f54a7cbb7338c70 Mon Sep 17 00:00:00 2001 From: daniel Date: Tue, 25 Aug 2026 18:35:01 +0100 Subject: [PATCH 3/3] feat(acp): add tracked response receipts --- src/agent-client-protocol/src/jsonrpc.rs | 281 +++++++++++++++--- .../src/jsonrpc/incoming_actor.rs | 1 + .../src/jsonrpc/outgoing_actor.rs | 130 ++++++-- src/agent-client-protocol/src/lib.rs | 4 +- .../tests/jsonrpc_batch.rs | 269 ++++++++++++++++- 5 files changed, 622 insertions(+), 63 deletions(-) diff --git a/src/agent-client-protocol/src/jsonrpc.rs b/src/agent-client-protocol/src/jsonrpc.rs index 9cff5a6..5068cea 100644 --- a/src/agent-client-protocol/src/jsonrpc.rs +++ b/src/agent-client-protocol/src/jsonrpc.rs @@ -2692,6 +2692,120 @@ pub fn is_cancel_request_notification(notification: &N) } } +/// Resolves when a tracked JSON-RPC response frame is accepted by the outgoing +/// transport-frame queue. +/// +/// This receipt does not indicate that any bytes were written to the transport. +/// For a response in a batch, it resolves only after every response in the batch +/// is ready and the aggregate batch frame is accepted by the queue. It resolves +/// with an error if enqueueing fails or the outgoing actor tears down first. +/// +/// # Batch handler deadlocks +/// +/// A batch cannot be enqueued until all of its handlers return. Do not await a +/// receipt from inside a batch handler. Return from the handler first, or spawn +/// receipt-dependent side effects so the handler can return immediately. +#[must_use = "a response receipt must be awaited to observe enqueue completion"] +pub struct ResponseReceipt { + receiver: oneshot::Receiver>, +} + +impl std::fmt::Debug for ResponseReceipt { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ResponseReceipt") + .finish_non_exhaustive() + } +} + +impl std::future::Future for ResponseReceipt { + type Output = Result<(), crate::Error>; + + fn poll( + mut self: std::pin::Pin<&mut Self>, + context: &mut std::task::Context<'_>, + ) -> std::task::Poll { + match std::pin::Pin::new(&mut self.receiver).poll(context) { + std::task::Poll::Ready(Ok(result)) => std::task::Poll::Ready(result), + std::task::Poll::Ready(Err(_)) => { + std::task::Poll::Ready(Err(response_receipt_teardown_error())) + } + std::task::Poll::Pending => std::task::Poll::Pending, + } + } +} + +struct ResponseReceiptSender { + state: Arc, +} + +impl std::fmt::Debug for ResponseReceiptSender { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ResponseReceiptSender") + .finish_non_exhaustive() + } +} + +struct ResponseReceiptState { + sender: Mutex>>>, +} + +impl ResponseReceiptSender { + fn channel() -> (Self, ResponseReceipt) { + let (sender, receiver) = oneshot::channel(); + ( + Self { + state: Arc::new(ResponseReceiptState { + sender: Mutex::new(Some(sender)), + }), + }, + ResponseReceipt { receiver }, + ) + } + + fn resolve(self, result: Result<(), crate::Error>) { + self.state.resolve(result); + } +} + +impl ResponseReceiptState { + fn resolve(&self, result: Result<(), crate::Error>) { + if let Some(sender) = self + .sender + .lock() + .expect("response receipt mutex poisoned") + .take() + { + drop(sender.send(result)); + } + } +} + +impl Drop for ResponseReceiptState { + fn drop(&mut self) { + if let Some(sender) = self + .sender + .get_mut() + .expect("response receipt mutex poisoned") + .take() + { + drop(sender.send(Err(response_receipt_teardown_error()))); + } + } +} + +fn response_receipt_teardown_error() -> crate::Error { + crate::util::internal_error( + "outgoing JSON-RPC actor stopped before the response frame was enqueued", + ) +} + +struct CompletedResponseFrame { + frame: TransportFrame, + receipts: Vec, +} + /// Messages send to be serialized over the transport. #[derive(Clone)] enum ResponseDestination { @@ -2719,6 +2833,7 @@ impl ResponseDestination { responses: (0..slot_count).map(|_| None).collect(), abandoned: (0..slot_count).map(|_| None).collect(), active_handler_attempts: (0..slot_count).map(|_| 0).collect(), + receipts: (0..slot_count).map(|_| None).collect(), dispatch_complete: false, emitted: false, })); @@ -2737,14 +2852,18 @@ impl ResponseDestination { ) } - fn complete(self, response: RawJsonRpcMessage) -> Option { + fn complete( + self, + response: RawJsonRpcMessage, + receipt: Option, + ) -> Option { match self { - Self::Individual(slot) => slot.complete(response), - Self::Batch(slot) => slot.complete(response).map(batch_response_frame), + Self::Individual(slot) => slot.complete(response, receipt), + Self::Batch(slot) => slot.complete(response, receipt).map(batch_response_frame), } } - fn abandon(self, fallback: RawJsonRpcMessage) -> Option { + fn abandon(self, fallback: RawJsonRpcMessage) -> Option { match self { Self::Individual(_) => None, Self::Batch(slot) => slot.abandon(fallback).map(batch_response_frame), @@ -2769,7 +2888,7 @@ impl ResponseDestination { }) } - fn finish_handler_attempt(self) -> Option { + fn finish_handler_attempt(self) -> Option { match self { Self::Individual(_) => None, Self::Batch(slot) => slot.finish_handler_attempt().map(batch_response_frame), @@ -2783,21 +2902,33 @@ struct IndividualResponseSlot { } impl IndividualResponseSlot { - fn complete(self, response: RawJsonRpcMessage) -> Option { + fn complete( + self, + response: RawJsonRpcMessage, + receipt: Option, + ) -> Option { if self.completed.swap(true, Ordering::AcqRel) { tracing::warn!("Ignoring duplicate completion of JSON-RPC request"); return None; } - Some(TransportFrame::Single(response)) + Some(CompletedResponseFrame { + frame: TransportFrame::Single(response), + receipts: receipt.into_iter().collect(), + }) } } -fn batch_response_frame(responses: Vec) -> TransportFrame { - TransportFrame::Batch( - TransportBatch::from_messages(responses) - .expect("a completed JSON-RPC response batch is non-empty"), - ) +fn batch_response_frame( + (responses, receipts): (Vec, Vec), +) -> CompletedResponseFrame { + CompletedResponseFrame { + frame: TransportFrame::Batch( + TransportBatch::from_messages(responses) + .expect("a completed JSON-RPC response batch is non-empty"), + ), + receipts, + } } #[derive(Clone)] @@ -2814,7 +2945,7 @@ impl std::fmt::Debug for BatchDispatchCompletion { } impl BatchDispatchCompletion { - fn complete(self) -> Option { + fn complete(self) -> Option { let mut state = self .state .lock() @@ -2841,23 +2972,25 @@ fn promote_abandoned_response(state: &mut BatchResponseState, index: usize) { } } -fn take_completed_batch(state: &mut BatchResponseState) -> Option> { +fn take_completed_batch( + state: &mut BatchResponseState, +) -> Option<(Vec, Vec)> { if !state.dispatch_complete || state.remaining != 0 || state.emitted { return None; } state.emitted = true; - Some( - state - .responses - .iter_mut() - .map(|response| { - response - .take() - .expect("completed JSON-RPC batch has every response slot") - }) - .collect(), - ) + let responses = state + .responses + .iter_mut() + .map(|response| { + response + .take() + .expect("completed JSON-RPC batch has every response slot") + }) + .collect(); + let receipts = state.receipts.iter_mut().filter_map(Option::take).collect(); + Some((responses, receipts)) } #[derive(Clone)] @@ -2884,7 +3017,9 @@ impl BatchResponseSlot { state.active_handler_attempts[self.index] += 1; } - fn finish_handler_attempt(self) -> Option> { + fn finish_handler_attempt( + self, + ) -> Option<(Vec, Vec)> { let mut state = self .state .lock() @@ -2898,7 +3033,11 @@ impl BatchResponseSlot { take_completed_batch(&mut state) } - fn complete(self, response: RawJsonRpcMessage) -> Option> { + fn complete( + self, + response: RawJsonRpcMessage, + receipt: Option, + ) -> Option<(Vec, Vec)> { let mut state = self .state .lock() @@ -2924,11 +3063,15 @@ impl BatchResponseSlot { state.abandoned[self.index] = None; state.responses[self.index] = Some(response); + state.receipts[self.index] = receipt; state.remaining -= 1; take_completed_batch(&mut state) } - fn abandon(self, fallback: RawJsonRpcMessage) -> Option> { + fn abandon( + self, + fallback: RawJsonRpcMessage, + ) -> Option<(Vec, Vec)> { let mut state = self .state .lock() @@ -2959,6 +3102,7 @@ struct BatchResponseState { responses: Vec>, abandoned: Vec>, active_handler_attempts: Vec, + receipts: Vec>, dispatch_complete: bool, emitted: bool, } @@ -3138,6 +3282,8 @@ enum OutgoingMessage { response: Result, destination: ResponseDestination, + + receipt: Option, }, /// Send an Error Response that cannot be correlated to a request ID. @@ -4452,7 +4598,13 @@ pub struct Responder { /// /// For incoming requests: serializes to JSON and sends over the wire. /// For incoming responses: sends to the waiting oneshot channel. - send_fn: Box) -> Result<(), crate::Error> + Send>, + send_fn: Box< + dyn FnOnce( + Result, + Option, + ) -> Result<(), crate::Error> + + Send, + >, /// Completes an abandoned batch slot unless an explicit response disarms it. drop_guard: ResponderDropGuard, @@ -4533,17 +4685,20 @@ impl Responder { id, cancellation, destination, - send_fn: Box::new(move |response: Result| { - send_raw_message( - &message_tx, - OutgoingMessage::Response { - id: id_clone, - method: method_clone, - response, - destination: send_destination, - }, - ) - }), + send_fn: Box::new( + move |response: Result, receipt| { + send_raw_message( + &message_tx, + OutgoingMessage::Response { + id: id_clone, + method: method_clone, + response, + destination: send_destination, + receipt, + }, + ) + }, + ), drop_guard, } } @@ -4619,9 +4774,9 @@ impl Responder { id: self.id, cancellation: self.cancellation, destination: self.destination, - send_fn: Box::new(move |input: Result| { + send_fn: Box::new(move |input: Result, receipt| { let t_value = wrap_fn(&method, input); - (self.send_fn)(t_value) + (self.send_fn)(t_value, receipt) }), drop_guard: self.drop_guard, } @@ -4634,7 +4789,47 @@ impl Responder { ) -> Result<(), crate::Error> { tracing::debug!(id = ?self.id, "respond called"); self.drop_guard.disarm(); - (self.send_fn)(response) + (self.send_fn)(response, None) + } + + /// Respond to the JSON-RPC request with either a value (`Ok`) or an error (`Err`) + /// and return a receipt for enqueue completion. + /// + /// The receipt resolves successfully when the response's single + /// [`TransportFrame`], or its aggregate batch frame, is accepted by the + /// outgoing transport-frame queue. It does not wait for bytes to be written. + /// It resolves with an error if enqueueing fails or the outgoing actor tears + /// down first. + /// + /// # Errors + /// + /// Returns an error immediately if the response cannot enter the outgoing + /// protocol queue. After this method returns a receipt, enqueue or actor + /// teardown failures are reported by awaiting that receipt. + /// + /// # Batch handler deadlocks + /// + /// A batch frame cannot be enqueued until all batch handlers return. Do not + /// await the receipt inside a batch handler. Return first, or spawn any side + /// effect that awaits the receipt so the handler can return immediately. + pub fn respond_with_result_tracked( + mut self, + response: Result, + ) -> Result { + tracing::debug!(id = ?self.id, "tracked respond called"); + let (sender, receipt) = ResponseReceiptSender::channel(); + self.drop_guard.disarm(); + (self.send_fn)(response, Some(sender))?; + Ok(receipt) + } + + /// Respond to the JSON-RPC request with a value and return a receipt for + /// enqueue completion. + /// + /// See [`respond_with_result_tracked`](Self::respond_with_result_tracked) for + /// receipt semantics and the batch-handler deadlock warning. + pub fn respond_tracked(self, response: T) -> Result { + self.respond_with_result_tracked(Ok(response)) } /// Respond to the JSON-RPC request with a value. diff --git a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs index b65fcca..a4ab7a3 100644 --- a/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/incoming_actor.rs @@ -647,6 +647,7 @@ fn handle_handler_error( method: reply_target.method, response: Err(error), destination: reply_target.destination, + receipt: None, }, ), Some(HandlerErrorTarget::Response(reply_target)) => { diff --git a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs index effe164..0d115c3 100644 --- a/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs +++ b/src/agent-client-protocol/src/jsonrpc/outgoing_actor.rs @@ -1,9 +1,14 @@ // Types re-exported from crate root +use std::sync::{Arc, Weak}; + use futures::StreamExt as _; use futures::channel::mpsc; use crate::jsonrpc::protocol_compat::ProtocolCompat; -use crate::jsonrpc::{OutgoingMessage, PendingReplies, RawJsonRpcMessage, TransportFrame}; +use crate::jsonrpc::{ + CompletedResponseFrame, OutgoingMessage, PendingReplies, RawJsonRpcMessage, + ResponseReceiptSender, ResponseReceiptState, TransportFrame, response_receipt_teardown_error, +}; use crate::schema::v1::RequestId; pub type OutgoingMessageTx = mpsc::UnboundedSender; @@ -17,6 +22,47 @@ pub(crate) fn send_raw_message( .map_err(crate::util::internal_error) } +#[derive(Default)] +struct ResponseReceiptRegistry { + pending: Vec>, +} + +impl ResponseReceiptRegistry { + fn register(&mut self, sender: &ResponseReceiptSender) { + self.pending.retain(|state| state.strong_count() != 0); + self.pending.push(Arc::downgrade(&sender.state)); + } +} + +impl Drop for ResponseReceiptRegistry { + fn drop(&mut self) { + for state in self.pending.drain(..).filter_map(|state| state.upgrade()) { + state.resolve(Err(response_receipt_teardown_error())); + } + } +} + +fn enqueue_completed_response( + transport_tx: &mpsc::UnboundedSender, + completed: CompletedResponseFrame, +) -> Result<(), crate::Error> { + match transport_tx.unbounded_send(completed.frame) { + Ok(()) => { + for receipt in completed.receipts { + receipt.resolve(Ok(())); + } + Ok(()) + } + Err(error) => { + let error = crate::Error::into_internal_error(error); + for receipt in completed.receipts { + receipt.resolve(Err(error.clone())); + } + Err(error) + } + } +} + /// Outgoing protocol actor: Converts application-level OutgoingMessage to protocol-level RawJsonRpcMessage. /// /// This actor handles JSON-RPC protocol semantics: @@ -31,12 +77,13 @@ pub(super) async fn outgoing_protocol_actor( protocol_compat: ProtocolCompat, ) -> Result<(), crate::Error> { let mut drain_waiters = Vec::new(); + let mut receipt_registry = ResponseReceiptRegistry::default(); while let Some(message) = outgoing_rx.next().await { tracing::debug!(?message, "outgoing_protocol_actor"); // Create the message to be sent over the transport - let (json_rpc_message, destination) = match message { + let (json_rpc_message, destination, receipt) = match message { OutgoingMessage::CloseAfterDraining { done } => { // Reject later sends while preserving every message that was // already accepted into this receiver's buffer. @@ -46,17 +93,13 @@ pub(super) async fn outgoing_protocol_actor( } OutgoingMessage::BatchDispatchComplete { completion } => { if let Some(frame) = completion.complete() { - transport_tx - .unbounded_send(frame) - .map_err(crate::Error::into_internal_error)?; + enqueue_completed_response(&transport_tx, frame)?; } continue; } OutgoingMessage::BatchHandlerAttemptComplete { destination } => { if let Some(frame) = destination.finish_handler_attempt() { - transport_tx - .unbounded_send(frame) - .map_err(crate::Error::into_internal_error)?; + enqueue_completed_response(&transport_tx, frame)?; } continue; } @@ -79,9 +122,7 @@ pub(super) async fn outgoing_protocol_actor( ); let fallback = RawJsonRpcMessage::response(id, fallback); if let Some(frame) = destination.abandon(fallback) { - transport_tx - .unbounded_send(frame) - .map_err(crate::Error::into_internal_error)?; + enqueue_completed_response(&transport_tx, frame)?; } continue; } @@ -180,14 +221,23 @@ pub(super) async fn outgoing_protocol_actor( method, response, destination, + receipt, } => match protocol_compat.outgoing_response_to(&id, &method, response) { Ok(value) => { tracing::debug!(?id, "Sending success response"); - (RawJsonRpcMessage::response(id, Ok(value)), destination) + ( + RawJsonRpcMessage::response(id, Ok(value)), + destination, + receipt, + ) } Err(error) => { tracing::warn!(?id, %method, ?error, "Sending error response"); - (RawJsonRpcMessage::response(id, Err(error)), destination) + ( + RawJsonRpcMessage::response(id, Err(error)), + destination, + receipt, + ) } }, OutgoingMessage::UncorrelatedErrorResponse { error, destination } => { @@ -196,14 +246,16 @@ pub(super) async fn outgoing_protocol_actor( ( RawJsonRpcMessage::response(RequestId::Null, Err(error)), destination, + None, ) } }; - if let Some(frame) = destination.complete(json_rpc_message) { - transport_tx - .unbounded_send(frame) - .map_err(crate::Error::into_internal_error)?; + if let Some(receipt) = receipt.as_ref() { + receipt_registry.register(receipt); + } + if let Some(frame) = destination.complete(json_rpc_message, receipt) { + enqueue_completed_response(&transport_tx, frame)?; } } @@ -216,3 +268,47 @@ pub(super) async fn outgoing_protocol_actor( } Ok(()) } + +#[cfg(test)] +mod tests { + use futures::executor::block_on; + use futures::future::join; + + use super::*; + + #[test] + fn actor_teardown_fails_every_registered_response_receipt() { + let (first_sender, first_receipt) = ResponseReceiptSender::channel(); + let (second_sender, second_receipt) = ResponseReceiptSender::channel(); + let mut registry = ResponseReceiptRegistry::default(); + registry.register(&first_sender); + registry.register(&second_sender); + + drop(registry); + + let (first_result, second_result) = block_on(join(first_receipt, second_receipt)); + assert!(first_result.is_err()); + assert!(second_result.is_err()); + drop((first_sender, second_sender)); + } + + #[test] + fn response_transport_queue_failure_fails_every_receipt() { + let (transport_tx, transport_rx) = mpsc::unbounded(); + drop(transport_rx); + let (first_sender, first_receipt) = ResponseReceiptSender::channel(); + let (second_sender, second_receipt) = ResponseReceiptSender::channel(); + let completed = CompletedResponseFrame { + frame: TransportFrame::Single(RawJsonRpcMessage::response( + RequestId::Null, + Ok(serde_json::Value::Null), + )), + receipts: vec![first_sender, second_sender], + }; + + assert!(enqueue_completed_response(&transport_tx, completed).is_err()); + let (first_result, second_result) = block_on(join(first_receipt, second_receipt)); + assert!(first_result.is_err()); + assert!(second_result.is_err()); + } +} diff --git a/src/agent-client-protocol/src/lib.rs b/src/agent-client-protocol/src/lib.rs index 4f71dbd..b94cddf 100644 --- a/src/agent-client-protocol/src/lib.rs +++ b/src/agent-client-protocol/src/lib.rs @@ -125,8 +125,8 @@ pub use jsonrpc::{ HandleConnectionClose, HandleDispatchFrom, Handled, INCOMING_TRANSPORT_CLOSED_REASON, IntoHandled, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, Lines, NullClose, NullHandler, RawConnectionContext, RawJsonRpcMessage, RawJsonRpcParams, Responder, - ResponseRouter, SentRequest, TransportBatch, TransportBatchEntry, TransportFrame, - UntypedMessage, is_incoming_transport_closed, + ResponseReceipt, ResponseRouter, SentRequest, TransportBatch, TransportBatchEntry, + TransportFrame, UntypedMessage, is_incoming_transport_closed, run::{ChainRun, NullRun, RunWithConnectionTo}, }; pub use jsonrpc::{RequestCancellation, is_cancel_request_notification}; diff --git a/src/agent-client-protocol/tests/jsonrpc_batch.rs b/src/agent-client-protocol/tests/jsonrpc_batch.rs index 97df192..ecb812b 100644 --- a/src/agent-client-protocol/tests/jsonrpc_batch.rs +++ b/src/agent-client-protocol/tests/jsonrpc_batch.rs @@ -17,7 +17,8 @@ use std::{ use agent_client_protocol::{ Agent, ByteStreams, Channel, ConnectTo, ConnectionTo, Dispatch, Error, HandleDispatchFrom, Handled, JsonRpcMessage, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, - RawJsonRpcMessage, Responder, TransportBatch, TransportBatchEntry, TransportFrame, + RawJsonRpcMessage, Responder, ResponseReceipt, TransportBatch, TransportBatchEntry, + TransportFrame, role::{Role, UntypedRole}, schema::ProtocolVersion, schema::v1, @@ -237,6 +238,79 @@ async fn next_deferred_response( .expect("deferred responder channel closed unexpectedly") } +async fn next_response_receipt( + rx: &mut mpsc::UnboundedReceiver<(String, ResponseReceipt)>, +) -> (String, ResponseReceipt) { + tokio::time::timeout(TIMEOUT, rx.next()) + .await + .expect("timed out waiting for response receipt") + .expect("response receipt channel closed unexpectedly") +} + +fn start_tracked_server( + responder_tx: mpsc::UnboundedSender<(String, Responder)>, + receipt_tx: mpsc::UnboundedSender<(String, ResponseReceipt)>, +) -> ( + DuplexStream, + BufReader, + JoinHandle>, +) { + let (peer_writer, sdk_reader) = tokio::io::duplex(8192); + let (sdk_writer, peer_reader) = tokio::io::duplex(8192); + let transport = ByteStreams::new(sdk_writer.compat_write(), sdk_reader.compat()); + + let server = UntypedRole.builder().on_receive_request( + async move |request: TestRequest, + responder: Responder, + connection: ConnectionTo| { + let message = request.message; + if message == "drop responder" { + drop(responder); + return Ok(()); + } + if message.starts_with("deferred") { + return responder_tx + .unbounded_send((message, responder)) + .map_err(agent_client_protocol::Error::into_internal_error); + } + if message == "tracked then notification" { + let receipt = responder.respond_tracked(TestResponse { + result: "tracked response".into(), + })?; + let notification_connection = connection.clone(); + connection.spawn(async move { + receipt.await?; + notification_connection.send_notification(TestNotification { + message: "after tracked response".into(), + }) + })?; + return Ok(()); + } + if message == "untracked" { + return responder.respond(TestResponse { + result: "untracked response".into(), + }); + } + + let response = if message == "tracked error" { + Err(agent_client_protocol::Error::internal_error().data("tracked error")) + } else { + Ok(TestResponse { + result: format!("echo: {message}"), + }) + }; + let receipt = responder.respond_with_result_tracked(response)?; + receipt_tx + .unbounded_send((message, receipt)) + .map_err(agent_client_protocol::Error::into_internal_error) + }, + agent_client_protocol::on_receive_request!(), + ); + + let server_task = tokio::task::spawn_local(server.connect_to(transport)); + (peer_writer, BufReader::new(peer_reader), server_task) +} + fn start_deferred_server( responder_tx: mpsc::UnboundedSender<(String, Responder)>, ) -> ( @@ -281,6 +355,199 @@ async fn finish_server( .expect("server connection failed"); } +#[tokio::test(flavor = "current_thread")] +async fn tracked_individual_response_orders_notification_and_leaves_untracked_unchanged() { + tokio::task::LocalSet::new() + .run_until(async { + let (responder_tx, _responder_rx) = mpsc::unbounded(); + let (receipt_tx, _receipt_rx) = mpsc::unbounded(); + let (mut peer_writer, mut peer_reader, server_task) = + start_tracked_server(responder_tx, receipt_tx); + + write_json_line( + &mut peer_writer, + &json!({ + "jsonrpc": "2.0", + "id": 1, + "method": "test/echo", + "params": { "message": "tracked then notification" } + }), + ) + .await; + + let response = read_json_line(&mut peer_reader).await; + assert_eq!(response["id"], json!(1)); + assert_eq!(response["result"], json!({ "result": "tracked response" })); + let notification = read_json_line(&mut peer_reader).await; + assert_eq!(notification["method"], json!("test/notify")); + assert_eq!( + notification["params"], + json!({ "message": "after tracked response" }) + ); + + write_json_line( + &mut peer_writer, + &json!({ + "jsonrpc": "2.0", + "id": 2, + "method": "test/echo", + "params": { "message": "untracked" } + }), + ) + .await; + let untracked = read_json_line(&mut peer_reader).await; + assert_eq!(untracked["id"], json!(2)); + assert_eq!( + untracked["result"], + json!({ "result": "untracked response" }) + ); + + finish_server(peer_writer, server_task).await; + }) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn tracked_error_response_receipt_resolves_when_frame_is_enqueued() { + tokio::task::LocalSet::new() + .run_until(async { + let (responder_tx, _responder_rx) = mpsc::unbounded(); + let (receipt_tx, mut receipt_rx) = mpsc::unbounded(); + let (mut peer_writer, mut peer_reader, server_task) = + start_tracked_server(responder_tx, receipt_tx); + + write_json_line( + &mut peer_writer, + &json!({ + "jsonrpc": "2.0", + "id": 3, + "method": "test/echo", + "params": { "message": "tracked error" } + }), + ) + .await; + + let (message, receipt) = next_response_receipt(&mut receipt_rx).await; + assert_eq!(message, "tracked error"); + receipt + .await + .expect("error response frame should be enqueued"); + let response = read_json_line(&mut peer_reader).await; + assert_eq!(response["id"], json!(3)); + assert_eq!(response["error"]["code"], json!(-32603)); + assert_eq!(response["error"]["data"], json!("tracked error")); + + finish_server(peer_writer, server_task).await; + }) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn tracked_batch_receipts_resolve_together_after_aggregate_is_ready() { + tokio::task::LocalSet::new() + .run_until(async { + let (responder_tx, mut responder_rx) = mpsc::unbounded(); + let (receipt_tx, mut receipt_rx) = mpsc::unbounded(); + let (mut peer_writer, mut peer_reader, server_task) = + start_tracked_server(responder_tx, receipt_tx); + + write_json_line( + &mut peer_writer, + &json!([ + { + "jsonrpc": "2.0", + "id": 4, + "method": "test/echo", + "params": { "message": "immediate tracked" } + }, + { + "jsonrpc": "2.0", + "id": 5, + "method": "test/echo", + "params": { "message": "deferred tracked" } + } + ]), + ) + .await; + + let (message, mut first_receipt) = next_response_receipt(&mut receipt_rx).await; + assert_eq!(message, "immediate tracked"); + let (message, responder) = next_deferred_response(&mut responder_rx).await; + assert_eq!(message, "deferred tracked"); + assert!( + tokio::time::timeout(Duration::from_millis(25), &mut first_receipt) + .await + .is_err(), + "a batch receipt resolved before its sibling response was ready" + ); + + let second_receipt = responder + .respond_tracked(TestResponse { + result: "echo: deferred tracked".into(), + }) + .expect("deferred tracked response should be accepted"); + let (first_result, second_result) = join(first_receipt, second_receipt).await; + first_result.expect("first batch receipt should resolve"); + second_result.expect("second batch receipt should resolve"); + + let response = read_json_line(&mut peer_reader).await; + let responses = response + .as_array() + .expect("tracked batch should emit one aggregate array"); + assert_eq!(responses.len(), 2); + + finish_server(peer_writer, server_task).await; + }) + .await; +} + +#[tokio::test(flavor = "current_thread")] +async fn tracked_batch_receipt_resolves_when_sibling_responder_is_abandoned() { + tokio::task::LocalSet::new() + .run_until(async { + let (responder_tx, _responder_rx) = mpsc::unbounded(); + let (receipt_tx, mut receipt_rx) = mpsc::unbounded(); + let (mut peer_writer, mut peer_reader, server_task) = + start_tracked_server(responder_tx, receipt_tx); + + write_json_line( + &mut peer_writer, + &json!([ + { + "jsonrpc": "2.0", + "id": 6, + "method": "test/echo", + "params": { "message": "tracked sibling" } + }, + { + "jsonrpc": "2.0", + "id": 7, + "method": "test/echo", + "params": { "message": "drop responder" } + } + ]), + ) + .await; + + let (message, receipt) = next_response_receipt(&mut receipt_rx).await; + assert_eq!(message, "tracked sibling"); + receipt + .await + .expect("tracked sibling should resolve with abandoned fallback batch"); + let response = read_json_line(&mut peer_reader).await; + let responses = response + .as_array() + .expect("abandoned sibling should still emit one aggregate array"); + assert_eq!(responses.len(), 2); + assert!(responses.iter().any(|response| { + response["id"] == json!(7) && response["error"]["code"] == json!(-32603) + })); + + finish_server(peer_writer, server_task).await; + }) + .await; +} + #[tokio::test(flavor = "current_thread")] async fn mixed_batch_returns_one_array_with_each_response_bearing_entry() { tokio::task::LocalSet::new()